diff --git a/.clangd b/.clangd index 99f2765a557c..c8d6fdda360a 100644 --- a/.clangd +++ b/.clangd @@ -29,7 +29,7 @@ CompileFlags: # Tweak the clangd parse settings for all files CompileFlags: Compiler: clang++ - CompilationDatabase: . + CompilationDatabase: cpp/build Add: # report all errors - "-ferror-limit=0" diff --git a/.coderabbit.yaml b/.coderabbit.yaml new file mode 100644 index 000000000000..bb78fe3508c5 --- /dev/null +++ b/.coderabbit.yaml @@ -0,0 +1,40 @@ +# 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. + +# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json +# https://docs.coderabbit.ai/getting-started/configure-coderabbit/ +# In PR, comment "@coderabbitai configuration" to get the full config including defaults +language: "en-US" +reviews: + profile: chill + auto_title_placeholder: '@coderabbitai title' + auto_title_instructions: 'Should follow the format: "[fix/feat/doc/infra/...] \". Keep it concise.' + commit_status: false + collapse_walkthrough: true + assess_linked_issues: true + related_issues: true + related_prs: true + suggested_labels: true + auto_apply_labels: true + suggested_reviewers: true + auto_assign_reviewers: true + poem: false + auto_review: + drafts: true + base_branches: ["main", "release/.+"] +knowledge_base: + code_guidelines: + enabled: true + filePatterns: ["**/CODING_GUIDELINES.md"] diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index b407b6624098..883d39817aa3 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,17 +1,24 @@ +@coderabbitai summary -# PR title - -Please write the PR title by following template: + + ## Description + ## Test Coverage @@ -31,29 +38,40 @@ See details below for each supported subcommand.
-`run [--disable-fail-fast --skip-test --stage-list "A10-1, xxx" --gpu-type "A30, H100_PCIe" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-[Post-Merge]-1, xxx"]` +`run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]` Launch build/test pipelines. All previously running jobs will be killed. +`--reuse-test (optional)pipeline-id ` *(OPTIONAL)* : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline. + +`--disable-reuse-test ` *(OPTIONAL)* : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes. + `--disable-fail-fast ` *(OPTIONAL)* : Disable fail fast on build/tests/infra failures. `--skip-test ` *(OPTIONAL)* : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does **NOT** update GitHub check status. -`--stage-list "A10-1, xxx"` *(OPTIONAL)* : Only run the specified test stages. Examples: "A10-1, xxx". Note: Does **NOT** update GitHub check status. +`--stage-list "A10-PyTorch-1, xxx"` *(OPTIONAL)* : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does **NOT** update GitHub check status. `--gpu-type "A30, H100_PCIe"` *(OPTIONAL)* : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does **NOT** update GitHub check status. +`--test-backend "pytorch, cpp"` *(OPTIONAL)* : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does **NOT** update GitHub pipeline status. + `--only-multi-gpu-test ` *(OPTIONAL)* : Only run the multi-GPU tests. Note: Does **NOT** update GitHub check status. `--disable-multi-gpu-test ` *(OPTIONAL)* : Disable the multi-GPU tests. Note: Does **NOT** update GitHub check status. -`--add-multi-gpu-test ` *(OPTIONAL)* : Force run the multi-GPU tests. Will also run L0 pre-merge pipeline. +`--add-multi-gpu-test ` *(OPTIONAL)* : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline. `--post-merge ` *(OPTIONAL)* : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline. -`--extra-stage "H100_PCIe-[Post-Merge]-1, xxx"` *(OPTIONAL)* : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-[Post-Merge]-1, xxx". +`--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"` *(OPTIONAL)* : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx". + +`--detailed-log ` *(OPTIONAL)* : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job. + +`--debug ` *(OPTIONAL)* : **Experimental feature**. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the `stage-list` parameter to access the appropriate container environment. Note: Does **NOT** update GitHub check status. -For guidance on mapping tests to stage names, see `docs/source/reference/ci-overview.md`. +For guidance on mapping tests to stage names, see `docs/source/reference/ci-overview.md` +and the `scripts/test_to_stage_mapping.py` helper. ### kill diff --git a/.github/workflows/blossom-ci.yml b/.github/workflows/blossom-ci.yml index 7690a85e22d2..b2b253b2f6c0 100644 --- a/.github/workflows/blossom-ci.yml +++ b/.github/workflows/blossom-ci.yml @@ -40,7 +40,7 @@ jobs: startsWith(github.event.comment.body, '/bot skip --comment') || startsWith(github.event.comment.body, '/bot reuse-pipeline') || startsWith(github.event.comment.body, '/bot kill')) && contains( - fromJson('["byshiue","chuangz0","funatiq","hypdeb","jdemouth-nvidia","joyang-nv","lowsfer","Tabrizian","yweng0828","Shixiaowei02","MartinMarciniszyn","schetlur-nv","dcampora","pcastonguay","Naveassaf","lfr-0531","nekorobov","PerkzZheng","kaiyux","nv-guomingz","LinPoly","thorjohnsen","jiahanc","latency1024","tburt-nv","zeroepoch","chzblych","niukuo","ZhanruiSunCh","EmmaQiaoCh","yiqingy0","achartier","suyoggupta","amukkara","mk-nvidia","QiJune","lucaslie","davidmlw","hlu1","nvzhou","syuoni","NVGaryJi","symphonylyh","hello-11","zongfeijing","Jackch-NV","jinyangyuan-nvidia","LarryXFly","crazydemo","jaedeok-nvidia","wm2012011492","rosenrodt","zhuoyao1012","xinhe-nv","Yuening-wa","Shunkangz","zhengd-nv","yibinl-nvidia","StanleySun639","KingsleyLiu-NV","kxdc","yingcanw","BestJuly","ChristinaZ","bobboli","xueweilnvidia","kunlunl","cherichy","lucifer1004","Autumn1998","litaotju","peaceh-nv","liji-nv","SimengLiu-nv","yuxianq","yechank-nvidia","vallis-neria","DylanChen-NV","Tracin","zhhuang-nv","ISEEKYAN","xupinjie","tongyuantongyu","laikhtewari","zhuolingwang","dominicshanshan","jershi425","shifangx","StudyingShao","Superjomn","dongjiyingdjy","guangyunh-nv","wili-65535","tiffany940107","DanBlanaru","mikeiovine","djns99","ruodil","xiaoweiw-nv","xuwchen","bashimao","yizhang-nv","hyukn","nvpohanh","yuki-666","juney-nvidia","barry-delaney","Kefeng-Duan","MinaHuai","yilin-void","jhaotingc","jmydurant","katec846","CarstyYou","Njuapp","Jie-Fang","nvbrantz","inocsin","ruoqianguo","chenfeiz0326","ming-wei","eopXD","longlee0622","dongfengy","georgeliu95","evezhier","rakib-hasan","shangz-ai","JyChang012","wangsiping1997","yuanjings-nvda","tomeras91","roikoren755","amirkl94","shaharmor98","danielafrimi","amitz-nv","hijkzzz","rzilberstein-nvidia","dc3671","hchings","yuhengxnv","dongxuy04","qiaoxj07","omera-nv","DomBrown","brb-nv","FrankD412","yuhsuan-t","Fridah-nv","a-mccarthy","HuiGao-NV","alexmsettle","meenchen","sugunav14","cjluo-nv","kyleliang-nv","chang-l","WeiHaocheng","qixiang-99","BatshevaBlack","ebarilanM","xmchen1987","lingjiew","heyuhhh","netanel-haber","jiefangz-nv","wyw1267","yunruis","sklevtsov-nvidia","jgangani","pamelap-nvidia","ixlmar","GalSha","Dido0o0","rabiel","nvzhihanj","milesial","fzmu727","zackyoray","RoeyAzran1992","viraatc","v-shobhit","yuanjingx87","uchihatmtkinu","nvrohanv","vegaluisjose","qsang-nv","ChunhuanLin","timlee0212","venkywonka","zbpatel","tijyojwad","shyeh25","zihaok","nv-yilinf","ttyio","farazkh80","yuantailing","JennyLiu-nv","moraxu","IzzyPutterman","nvchenghaoz","nvxuanyuc","poweiw","stnie","zhanga5","nzmora-nvidia","greg-kwasniewski1","linda-stadter","Tom-Zheng","vanshilshah97","ixlmar","MatthiasKohl","Wanli-Jiang", "arekay", "davidclark-nv", "2ez4bz", "tcherckez-nvidia", "MrGeva", "galagam", "limin2021", "dhansen-nvidia","talorabr","kanghui0204","wu6u3tw","hvagadia","xavier-nvidia","raayandhar","dbari","nvjullin","elvischenv","zhenhuaw-me","weireweire","yifeizhang-c","jiaganc","ziyixiong-nv","FelixXidddd","JunyiXu-nv","bo-nv","zerollzeng","RayenTian","ameynaik-hub"]'), + fromJson('["byshiue","chuangz0","funatiq","hypdeb","jdemouth-nvidia","joyang-nv","lowsfer","Tabrizian","yweng0828","Shixiaowei02","MartinMarciniszyn","schetlur-nv","dcampora","pcastonguay","Naveassaf","lfr-0531","nekorobov","PerkzZheng","kaiyux","nv-guomingz","LinPoly","thorjohnsen","jiahanc","latency1024","tburt-nv","zeroepoch","chzblych","niukuo","ZhanruiSunCh","EmmaQiaoCh","yiqingy0","achartier","suyoggupta","amukkara","mk-nvidia","QiJune","lucaslie","davidmlw","hlu1","nvzhou","syuoni","NVGaryJi","symphonylyh","hello-11","zongfeijing","Jackch-NV","jinyangyuan-nvidia","LarryXFly","crazydemo","jaedeok-nvidia","wm2012011492","rosenrodt","zhuoyao1012","xinhe-nv","Yuening-wa","Shunkangz","zhengd-nv","yibinl-nvidia","StanleySun639","KingsleyLiu-NV","kxdc","yingcanw","BestJuly","ChristinaZ","bobboli","xueweilnvidia","kunlunl","cherichy","lucifer1004","Autumn1998","litaotju","peaceh-nv","liji-nv","SimengLiu-nv","yuxianq","yechank-nvidia","vallis-neria","DylanChen-NV","Tracin","zhhuang-nv","ISEEKYAN","xupinjie","tongyuantongyu","laikhtewari","zhuolingwang","dominicshanshan","jershi425","shifangx","StudyingShao","Superjomn","dongjiyingdjy","guangyunh-nv","wili-65535","tiffany940107","DanBlanaru","mikeiovine","djns99","ruodil","xiaoweiw-nv","xuwchen","bashimao","yizhang-nv","hyukn","nvpohanh","yuki-666","juney-nvidia","barry-delaney","Kefeng-Duan","MinaHuai","yilin-void","jhaotingc","jmydurant","katec846","CarstyYou","Njuapp","Jie-Fang","nvbrantz","inocsin","ruoqianguo","chenfeiz0326","ming-wei","eopXD","longlee0622","dongfengy","georgeliu95","evezhier","rakib-hasan","shangz-ai","JyChang012","wangsiping1997","yuanjings-nvda","tomeras91","roikoren755","amirkl94","shaharmor98","danielafrimi","amitz-nv","hijkzzz","rzilberstein-nvidia","dc3671","hchings","yuhengxnv","dongxuy04","qiaoxj07","omera-nv","DomBrown","brb-nv","FrankD412","yuhsuan-t","Fridah-nv","a-mccarthy","HuiGao-NV","alexmsettle","meenchen","sugunav14","cjluo-nv","kyleliang-nv","chang-l","WeiHaocheng","qixiang-99","BatshevaBlack","ebarilanM","xmchen1987","lingjiew","heyuhhh","netanel-haber","jiefangz-nv","wyw1267","yunruis","sklevtsov-nvidia","jgangani","pamelap-nvidia","ixlmar","GalSha","Dido0o0","rabiel","nvzhihanj","milesial","fzmu727","zackyoray","RoeyAzran1992","viraatc","v-shobhit","yuanjingx87","uchihatmtkinu","nvrohanv","vegaluisjose","qsang-nv","ChunhuanLin","timlee0212","venkywonka","zbpatel","tijyojwad","shyeh25","zihaok","nv-yilinf","ttyio","farazkh80","yuantailing","JennyLiu-nv","moraxu","IzzyPutterman","nvchenghaoz","nvxuanyuc","poweiw","stnie","zhanga5","nzmora-nvidia","greg-kwasniewski1","linda-stadter","Tom-Zheng","vanshilshah97","ixlmar","MatthiasKohl","Wanli-Jiang", "arekay", "davidclark-nv", "2ez4bz", "tcherckez-nvidia", "MrGeva", "galagam", "limin2021", "dhansen-nvidia","talorabr","kanghui0204","wu6u3tw","hvagadia","xavier-nvidia","raayandhar","dbari","nvjullin","elvischenv","zhenhuaw-me","weireweire","yifeizhang-c","jiaganc","ziyixiong-nv","FelixXidddd","JunyiXu-nv","bo-nv","zerollzeng","RayenTian","ameynaik-hub","raymochen","shuyixiong","johncalesp","leslie-fang25","reasonsolo","zhou-yuxin","vadiklyutiy","yali-arch","NVShreyas","h-guo18","pengbowang-nv"]'), github.actor) steps: - name: Check if comment is issued by authorized person diff --git a/.github/workflows/bot-command.yml b/.github/workflows/bot-command.yml index 573e7f499ab6..6689ab619d38 100644 --- a/.github/workflows/bot-command.yml +++ b/.github/workflows/bot-command.yml @@ -46,17 +46,22 @@ jobs: "Run `/bot [-h|--help]` to print this help message.\n\n" + "See details below for each supported subcommand.\n\n" + "
\n\n" + - "`run [--disable-fail-fast --skip-test --stage-list \"A10-1, xxx\" --gpu-type \"A30, H100_PCIe\" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage \"H100_PCIe-[Post-Merge]-1, xxx\"]`\n\n" + + "`run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list \"A10-PyTorch-1, xxx\" --gpu-type \"A30, H100_PCIe\" --test-backend \"pytorch, cpp\" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage \"H100_PCIe-TensorRT-Post-Merge-1, xxx\" --detailed-log --debug(experimental)]`\n\n" + "Launch build/test pipelines. All previously running jobs will be killed.\n\n" + + "`--reuse-test (optional)pipeline-id ` *(OPTIONAL)* : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.\n\n" + + "`--disable-reuse-test ` *(OPTIONAL)* : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.\n\n" + "`--disable-fail-fast ` *(OPTIONAL)* : Disable fail fast on build/tests/infra failures.\n\n" + "`--skip-test ` *(OPTIONAL)* : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does **NOT** update GitHub check status.\n\n" + - "`--stage-list \"A10-1, xxx\"` *(OPTIONAL)* : Only run the specified test stages. Examples: \"A10-1, xxx\". Note: Does **NOT** update GitHub check status.\n\n" + + "`--stage-list \"A10-PyTorch-1, xxx\"` *(OPTIONAL)* : Only run the specified test stages. Examples: \"A10-PyTorch-1, xxx\". Note: Does **NOT** update GitHub check status.\n\n" + "`--gpu-type \"A30, H100_PCIe\"` *(OPTIONAL)* : Only run the test stages on the specified GPU types. Examples: \"A30, H100_PCIe\". Note: Does **NOT** update GitHub check status.\n\n" + + "`--test-backend \"pytorch, cpp\"` *(OPTIONAL)* : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: \"pytorch, cpp\" (does not run test stages with tensorrt or triton backend). Note: Does **NOT** update GitHub pipeline status.\n\n" + "`--only-multi-gpu-test ` *(OPTIONAL)* : Only run the multi-GPU tests. Note: Does **NOT** update GitHub check status.\n\n" + "`--disable-multi-gpu-test ` *(OPTIONAL)* : Disable the multi-GPU tests. Note: Does **NOT** update GitHub check status.\n\n" + - "`--add-multi-gpu-test ` *(OPTIONAL)* : Force run the multi-GPU tests. Will also run L0 pre-merge pipeline.\n\n" + + "`--add-multi-gpu-test ` *(OPTIONAL)* : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.\n\n" + "`--post-merge ` *(OPTIONAL)* : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.\n\n" + - "`--extra-stage \"H100_PCIe-[Post-Merge]-1, xxx\"` *(OPTIONAL)* : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage \"H100_PCIe-[Post-Merge]-1, xxx\".\n\n" + + "`--extra-stage \"H100_PCIe-TensorRT-Post-Merge-1, xxx\"` *(OPTIONAL)* : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage \"H100_PCIe-TensorRT-Post-Merge-1, xxx\".\n\n" + + "`--detailed-log ` *(OPTIONAL)* : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.\n\n" + + "`--debug ` *(OPTIONAL)* : **Experimental feature**. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the `stage-list` parameter to access the appropriate container environment. Note: Does **NOT** update GitHub check status.\n\n" + "### kill\n\n" + "`kill `\n\n" + "Kill all running builds associated with pull request.\n\n" + diff --git a/README.md b/README.md index 99b00e26195a..15449460963d 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ TensorRT-LLM [![python](https://img.shields.io/badge/python-3.10-green)](https://www.python.org/downloads/release/python-31012/) [![cuda](https://img.shields.io/badge/cuda-12.9.0-green)](https://developer.nvidia.com/cuda-downloads) [![trt](https://img.shields.io/badge/TRT-10.11.0-green)](https://developer.nvidia.com/tensorrt) -[![version](https://img.shields.io/badge/release-1.0.0rc3-green)](./tensorrt_llm/version.py) +[![version](https://img.shields.io/badge/release-1.0.0rc5-green)](./tensorrt_llm/version.py) [![license](https://img.shields.io/badge/license-Apache%202-blue)](./LICENSE) [Architecture](./docs/source/torch/arch_overview.md)   |   [Performance](./docs/source/performance/perf-overview.md)   |   [Examples](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html)   |   [Documentation](./docs/source/)   |   [Roadmap](https://github.com/NVIDIA/TensorRT-LLM/issues?q=is%3Aissue%20state%3Aopen%20label%3Aroadmap) @@ -223,6 +223,23 @@ To get started with TensorRT-LLM, visit our documentation: - [Benchmarking Performance](https://nvidia.github.io/TensorRT-LLM/performance/performance-tuning-guide/benchmarking-default-performance.html#benchmarking-with-trtllm-bench) - [Release Notes](https://nvidia.github.io/TensorRT-LLM/release-notes.html) +## Deprecation Policy + +Deprecation is used to inform developers that some APIs and tools are no longer recommended for use. Beginning with version 1.0, TensorRT-LLM has the following deprecation policy: + +1. Communication of Deprecation + - Deprecation notices are documented in the Release Notes. + - Deprecated APIs, methods, classes, or parameters include a statement in the source code indicating when they were deprecated. + - If used, deprecated methods, classes, or parameters issue runtime deprecation warnings. +2. Migration Period + - TensorRT-LLM provides a 3-month migration period after deprecation. + - During this period, deprecated APIs, tools, or parameters continue to work but trigger warnings. +3. Scope of Deprecation + - Full API/Method/Class Deprecation: The entire API/method/class is marked for removal. + - Partial Deprecation: If only specific parameters of an API/method are deprecated (e.g., param1 in LLM.generate(param1, param2)), the method itself remains functional, but the deprecated parameters will be removed in a future release. +4. Removal After Migration Period + - After the 3-month migration period ends, deprecated APIs, tools, or parameters are removed in a manner consistent with semantic versioning (major version changes may include breaking removals). + ## Useful Links - [Quantized models on Hugging Face](https://huggingface.co/collections/nvidia/model-optimizer-66aa84f7966b3150262481a4): A growing collection of quantized (e.g., FP8, FP4) and optimized LLMs, including [DeepSeek FP4](https://huggingface.co/nvidia/DeepSeek-R1-FP4), ready for fast inference with TensorRT-LLM. - [NVIDIA Dynamo](https://github.com/ai-dynamo/dynamo): A datacenter scale distributed inference serving framework that works seamlessly with TensorRT-LLM. diff --git a/benchmarks/cpp/__init__.py b/benchmarks/cpp/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/benchmarks/cpp/disaggServerBenchmark.cpp b/benchmarks/cpp/disaggServerBenchmark.cpp index d0b5fb8c8642..ab009802757a 100644 --- a/benchmarks/cpp/disaggServerBenchmark.cpp +++ b/benchmarks/cpp/disaggServerBenchmark.cpp @@ -636,6 +636,8 @@ class DisaggExecutorServer : texec::DecodingMode::Auto(), benchmarkParams.executorLookaheadConfig, benchmarkParams.medusaChoices)); executorConfig.setExtendedRuntimePerfKnobConfig(extendedRuntimePerfKnobConfig); + executorConfig.setCacheTransceiverConfig( + texec::CacheTransceiverConfig(texec::CacheTransceiverConfig::BackendType::DEFAULT)); constexpr int maxIterationsForRequestStats = 1000; if (mEnableCollectKvCacheTransferTime) { diff --git a/benchmarks/cpp/utils/__init__.py b/benchmarks/cpp/utils/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index a76b3e21558f..6732db6eaa7f 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -30,6 +30,7 @@ project(tensorrt_llm LANGUAGES CXX) 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_MICRO_BENCHMARKS "Build C++ micro benchmarks" OFF) option(NVTX_DISABLE "Disable all NVTX features" ON) option(WARNING_IS_ERROR "Treat all warnings as errors" OFF) @@ -198,7 +199,7 @@ set(TRT_LIB TensorRT::NvInfer) get_filename_component(TRT_LLM_ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR} PATH) set(3RDPARTY_DIR ${TRT_LLM_ROOT_DIR}/3rdparty) -if(BINDING_TYPE STREQUAL "pybind") +if(BINDING_TYPE STREQUAL "pybind" OR BUILD_DEEP_EP) add_subdirectory(${3RDPARTY_DIR}/pybind11 ${CMAKE_CURRENT_BINARY_DIR}/pybind11) endif() @@ -217,7 +218,7 @@ include_directories( ${3RDPARTY_DIR}/cutlass/tools/util/include ${3RDPARTY_DIR}/NVTX/include ${3RDPARTY_DIR}/json/include) -if(BINDING_TYPE STREQUAL "pybind") +if(BINDING_TYPE STREQUAL "pybind" OR BUILD_DEEP_EP) include_directories(${3RDPARTY_DIR}/pybind11/include) endif() if(BINDING_TYPE STREQUAL "nanobind") diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 6f9c2f82dd60..c39fee6f940e 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -70,28 +70,20 @@ class BaseCacheTransceiver class CacheTransceiver : public BaseCacheTransceiver { public: - enum class CommType : std::uint8_t - { - UNKNOWN = 0, - MPI = 1, - UCX = 2, - NIXL = 3 - }; - - CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, CommType commType, + CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, executor::kv_cache::CacheState::ModelConfig const& cacheStateModelCfg, runtime::WorldConfig const& worldConfig, nvinfer1::DataType dataType, executor::kv_cache::CacheState::AttentionType attentionType = executor::kv_cache::CacheState::AttentionType::kDEFAULT, std::optional cacheTransceiverConfig = std::nullopt); - CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, CommType commType, - std::vector numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, - runtime::WorldConfig const& worldConfig, nvinfer1::DataType dataType, + CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, std::vector numKvHeadsPerLayer, + SizeType32 sizePerHead, SizeType32 tokensPerBlock, runtime::WorldConfig const& worldConfig, + nvinfer1::DataType dataType, executor::kv_cache::CacheState::AttentionType attentionType = executor::kv_cache::CacheState::AttentionType::kDEFAULT, std::optional cacheTransceiverConfig = std::nullopt) - : CacheTransceiver(cacheManager, commType, + : CacheTransceiver(cacheManager, executor::kv_cache::CacheState::ModelConfig{numKvHeadsPerLayer, sizePerHead, tokensPerBlock}, worldConfig, dataType, attentionType, cacheTransceiverConfig) { @@ -118,7 +110,6 @@ class CacheTransceiver : public BaseCacheTransceiver void setContextState(LlmRequest* llmRequest); - CommType mCommType; std::unique_ptr mDataResponder; std::unique_ptr mDataRequester; std::vector>> mResponderFutures; diff --git a/cpp/include/tensorrt_llm/batch_manager/decoderBuffers.h b/cpp/include/tensorrt_llm/batch_manager/decoderBuffers.h index 831a4179ecb8..2af03c0af710 100644 --- a/cpp/include/tensorrt_llm/batch_manager/decoderBuffers.h +++ b/cpp/include/tensorrt_llm/batch_manager/decoderBuffers.h @@ -16,6 +16,7 @@ #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" @@ -38,8 +39,8 @@ class DecoderInputBuffers using SizeType32 = runtime::SizeType32; using TensorPtr = runtime::ITensor::SharedPtr; - explicit DecoderInputBuffers(SizeType32 maxNumSequences, SizeType32 maxBatchSize, SizeType32 maxDecoderSteps, - runtime::BufferManager const& manager); + explicit DecoderInputBuffers( + SizeType32 maxBatchSize, SizeType32 maxDecoderSteps, runtime::BufferManager const& manager); void setupMedusaLogits(SizeType32 maxNumSequences, runtime::ModelConfig const& modelConfig); @@ -56,11 +57,13 @@ class DecoderInputBuffers //! Buffers for decoder forward + //! Requests for considered in decoder forward + RequestVector decoderRequests; + //! Batch slots for all decoder steps, [maxDecoderSteps][maxBatchSize] std::vector forwardBatchSlots; - //! Logits for all batch slots, [maxNumSequences] - //! The vector is sparse, only slots in forwardBatchSlots are used. + //! Logits of decoder requests std::vector logits; //! Logits for speculative decoding (Medusa) diff --git a/cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h b/cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h index 26d20cc9fa39..9a577b61ad51 100644 --- a/cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h +++ b/cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h @@ -29,6 +29,7 @@ class GrammarCompiler; namespace tensorrt_llm::batch_manager { +class DecoderInputBuffers; class GuidedDecoder { @@ -40,8 +41,7 @@ class GuidedDecoder GuidedDecoder(executor::GuidedDecodingConfig const& guidedDecodingConfig, SizeType32 maxNumSequences, SizeType32 vocabSizePadded, nvinfer1::DataType logitsDtype, runtime::BufferManager const& runtimeBufferManager); void build(ScheduledRequests const& scheduledRequests); - void execute(ScheduledRequests const& scheduledRequests, runtime::BufferManager const& runtimeBufferManager, - std::vector const& decoderBuffersLogits); + void execute(DecoderInputBuffers const& decoderInputBuffers, runtime::BufferManager const& runtimeBufferManager); private: executor::GuidedDecodingConfig::GuidedDecodingBackend mGuidedDecodingBackend; diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index b863c3276740..a0234cbbe49b 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -31,6 +31,7 @@ #include "tensorrt_llm/runtime/worldConfig.h" #include +#include #include #include #include @@ -68,6 +69,9 @@ using VecUniqueTokens = tensorrt_llm::runtime::VecUniqueTokens; using LoraTaskIdType = tensorrt_llm::runtime::LoraTaskIdType; using BlocksPerWindow = std::map>; +// Type alias for multimodal hash key (hash array + start offset) +using MmKey = std::pair, SizeType32>; + template using OptionalRef = tensorrt_llm::common::OptionalRef; @@ -107,6 +111,10 @@ struct BlockKey std::optional loraTaskId = std::nullopt; VecUniqueTokens uniqueTokens; + // Extra keys for multimodal data (similar to VLLM's approach) + // Each extra key is a pair of (mm_hash, start_offset_in_block) + std::vector extraKeys; + BlockKey() = default; explicit BlockKey(VecTokens const& tokens, std::optional loraTaskId = std::nullopt) @@ -119,23 +127,25 @@ struct BlockKey } } - BlockKey(bool usesExtraIds, std::optional loraTaskId, VecUniqueTokens uniqueTokens) - : usesExtraIds(usesExtraIds) + explicit BlockKey(bool usesExtraIds, std::optional loraTaskId, VecUniqueTokens uniqueTokens, + std::vector extraKeys = {}) + : usesExtraIds{usesExtraIds} , loraTaskId{loraTaskId} , uniqueTokens{std::move(uniqueTokens)} + , extraKeys{std::move(extraKeys)} { } bool operator==(BlockKey const& other) const noexcept { - return ( - usesExtraIds == other.usesExtraIds && loraTaskId == other.loraTaskId && uniqueTokens == other.uniqueTokens); + return (usesExtraIds == other.usesExtraIds && loraTaskId == other.loraTaskId + && uniqueTokens == other.uniqueTokens && extraKeys == other.extraKeys); } int partialMatch(BlockKey const& other) const noexcept { SizeType32 numMatched{0}; - if (loraTaskId == other.loraTaskId) + if (loraTaskId == other.loraTaskId && extraKeys == other.extraKeys) { auto [matchEnd, otherMatchEnd] = std::mismatch( uniqueTokens.begin(), uniqueTokens.end(), other.uniqueTokens.begin(), other.uniqueTokens.end()); @@ -180,6 +190,8 @@ struct KvCacheStats SizeType32 missedBlocks; // Measuring the KV Cache reuse rate. cacheHitRate = reusedBlocks / (reusedBlocks + missedBlocks). float cacheHitRate; + // Number of free blocks for every configured attention-window size. + std::map numFreeBlocksPerWindowSize; }; // Basic building block of a paged KV cache - a single @@ -854,6 +866,9 @@ class BlockManager std::shared_ptr eventManager = nullptr, bool enableHashKey = false, bool enablePartialReuse = true, bool copyOnPartialReuse = true); + BlockManager(BlockManager const&) = delete; + BlockManager& operator=(BlockManager const&) = delete; + //! \brief Calculate the proportional share each window size receives of the total memory pool //! \details Example: (uniqueWindowSizeToLayers={1024: [1], 4096: [0, 4, 5], 8192: [2, 3]}) //! Would Return: {1024: 0.0345, 4096: 0.4138, 8192: 0.5517} [sums to 1.0]. @@ -1454,6 +1469,11 @@ class KVCacheManager : public BaseKVCacheManager return mBlockManager.getNumMissedBlocks(); } + [[nodiscard]] std::map getNumFreeBlocksPerWindowSize() const + { + return mBlockManager.getNumFreeBlocksPerWindowSize(); + } + [[nodiscard]] KvCacheStats getKvCacheStats() const override { KvCacheStats kvCacheStats; @@ -1468,6 +1488,7 @@ class KVCacheManager : public BaseKVCacheManager kvCacheStats.cacheHitRate = kvCacheStats.reusedBlocks == 0 ? 0 : static_cast(kvCacheStats.reusedBlocks) / static_cast(kvCacheStats.reusedBlocks + kvCacheStats.missedBlocks); + kvCacheStats.numFreeBlocksPerWindowSize = getNumFreeBlocksPerWindowSize(); return kvCacheStats; } diff --git a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h index cb8d6edb91fc..cb79f89a8ae3 100644 --- a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h +++ b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h @@ -826,6 +826,7 @@ class GenericLlmRequest mState = mEncoderTokens.has_value() || mEncoderInputFeatures ? LlmRequestState::kENCODER_INIT : LlmRequestState::kCONTEXT_INIT; mContextCurrentPosition = 0; + mPrepopulatedPromptLen = 0; mContextChunkSize = mPromptLen; mSeqSlot.reset(); } @@ -1564,7 +1565,9 @@ class GenericLlmRequest /// Returns whether the position is at the beginning of the context. [[nodiscard]] bool isFirstContextChunk() const noexcept { - return mContextCurrentPosition == 0; + // The number of cached token is encountered in mContextCurrentPosition, + // so the start position of the context is mPrepopulatedPromptLen. + return mContextCurrentPosition == mPrepopulatedPromptLen; } /// Move the cursor forward one chunk. When not chunked, move forward to the end of the context. diff --git a/cpp/include/tensorrt_llm/batch_manager/logitsPostProcessor.h b/cpp/include/tensorrt_llm/batch_manager/logitsPostProcessor.h index 9610b96763b4..048a84ecca34 100644 --- a/cpp/include/tensorrt_llm/batch_manager/logitsPostProcessor.h +++ b/cpp/include/tensorrt_llm/batch_manager/logitsPostProcessor.h @@ -24,28 +24,29 @@ namespace tensorrt_llm::runtime { -class TllmRuntime; +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&, - runtime::BufferManager::CudaStreamPtr const&, + std::vector> const&, CudaStreamPtr const&, std::vector> const&)>; constexpr static auto name{"LogitsPostProcessor"}; LogitsPostProcessor() = default; - bool operator()(RequestVector const& contextRequests, RequestVector const& generationRequests, - bool replicateLogitsPostProcessor, std::vector& seqSlotLogits, - runtime::WorldConfig const& worldConfig, runtime::TllmRuntime& runtime, + bool operator()(DecoderInputBuffers& inputBuffers, bool replicateLogitsPostProcessor, + runtime::WorldConfig const& worldConfig, CudaStreamPtr const& stream, std::optional logitsPostProcessorBatched = std::nullopt) const; }; diff --git a/cpp/include/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.h b/cpp/include/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.h index 1757a9f076ee..cea23a4e7ec9 100644 --- a/cpp/include/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.h +++ b/cpp/include/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.h @@ -46,8 +46,7 @@ class MakeDecodingBatchInputOutput : Algorithm MakeDecodingBatchInputOutput() = default; - std::unique_ptr operator()(RequestVector const& contextRequests, - RequestVector const& generationRequests, DecoderInputBuffers const& inputBuffers, + std::unique_ptr operator()(DecoderInputBuffers& inputBuffers, runtime::decoder::DecoderState& decoderState, runtime::ModelConfig const& modelConfig, SizeType32 maxNumSequences, OptionalRef fusedRuntimeBuffers) const; diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index 1cd651cd07ca..bba3c31a0148 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -1430,18 +1430,29 @@ class LogitsPostProcessorConfig class CacheTransceiverConfig { public: - explicit CacheTransceiverConfig(std::optional maxNumTokens = std::nullopt); + enum class BackendType : std::uint8_t + { + DEFAULT = 0, + MPI = 1, + UCX = 2, + NIXL = 3 + }; + explicit CacheTransceiverConfig( + std::optional backendType = std::nullopt, std::optional maxNumTokens = std::nullopt); bool operator==(CacheTransceiverConfig const& other) const; + void setBackendType(std::optional backendType); + void setMaxTokensInBuffer(std::optional maxTokensInBuffer); - [[nodiscard]] std::optional getMaxNumTokens() const; - void setMaxNumTokens(size_t maxNumTokens); + [[nodiscard]] std::optional getMaxTokensInBuffer() const; + [[nodiscard]] std::optional getBackendType() const; private: + std::optional mBackendType; /// @brief The maximum number of tokens that the CacheTransceiver's pre-allocated buffer can hold. If the number of /// kvCache tokens to be transferred for a single request is greater than this value, the performance of the cache /// transfer may be degraded. - std::optional mMaxNumTokens; + std::optional mMaxTokensInBuffer; }; /// @brief Configuration class for the model executor diff --git a/cpp/kernels/fmha_v2/fmha_test.py b/cpp/kernels/fmha_v2/fmha_test.py index 3523ee1d1002..f9f28978e661 100644 --- a/cpp/kernels/fmha_v2/fmha_test.py +++ b/cpp/kernels/fmha_v2/fmha_test.py @@ -150,14 +150,17 @@ def test_trtllm_sage_attention_fmha(d, s): @pytest.mark.parametrize('dtype', ["-bf16", "-e4m3", "-e4m3 -bf16-output"], ids=["bf16", "e4m3", "e4m3-bf16"]) @pytest.mark.parametrize('s', [1024, 4096], ids=["seqlen-1024", "seqlen-4096"]) -def test_trtllm_context_mla_attention_fmha(dtype, s): +@pytest.mark.parametrize( + 'input_layout', ["", "-paged-kv", "-contiguous-q-kv", "-separate-q-k-v"], + ids=["packed-qkv", "paged-kv", "q-contiguous-kv", "separate-q-k-v"]) +def test_trtllm_context_mla_attention_fmha(dtype, s, input_layout): # use higher error tolerance for bf16 and s = 4096. epsilon = '' if dtype == "-bf16" and s == 4096: epsilon += ' -epsilon 0.03' sm_version = getSMVersion() - if sm_version != 89: + if dtype in ["-e4m3", "-e4m3 -bf16-output"] and sm_version != 89: pytest.skip("FP8 MLAs only supported on sm89 currently.") # Context phase kernels. @@ -167,6 +170,14 @@ def test_trtllm_context_mla_attention_fmha(dtype, s): shell=True, check=True) + if sm_version == 90: + # Now only hopper-style supports separate-q-k-v + subprocess.run( + f"bin/fmha.exe -v 0 -runs 1 -min-s 1024 -s {s} -b 8 -h 8 -d 192 -dv 128 {dtype} \ + -causal-mask {epsilon} {input_layout}", + shell=True, + check=True) + @pytest.mark.parametrize('dtype', ["-bf16", "-e4m3", "-e4m3 -bf16-output"], ids=["bf16", "e4m3", "e4m3-bf16"]) diff --git a/cpp/kernels/fmha_v2/setup.py b/cpp/kernels/fmha_v2/setup.py index 8d3549f56fdc..e7a39864551d 100644 --- a/cpp/kernels/fmha_v2/setup.py +++ b/cpp/kernels/fmha_v2/setup.py @@ -101,6 +101,7 @@ class InputLayout(IntEnum): PACKED_QKV = 0 CONTIGUOUS_Q_KV = 1 Q_PAGED_KV = 2 + SEPARATE_Q_K_V = 3 spec_fields = ( @@ -1431,6 +1432,7 @@ def get_makefile_code(specs_names): {loop_step}, {kv_loop_step}, {head_size}, + {head_size_v}, {q_tile_buffers}, {kv_tile_buffers}, NUM_COMPUTE_GROUPS, @@ -1453,6 +1455,7 @@ def get_makefile_code(specs_names): {loop_step}, {kv_loop_step}, {head_size}, + {head_size_v}, {q_tile_buffers}, {kv_tile_buffers}, NUM_COMPUTE_GROUPS, @@ -1472,6 +1475,7 @@ def get_makefile_code(specs_names): {loop_step}, {kv_loop_step}, {head_size}, + {head_size_v}, {q_tile_buffers}, {kv_tile_buffers}, NUM_COMPUTE_GROUPS, @@ -1491,6 +1495,7 @@ def get_makefile_code(specs_names): {loop_step}, {kv_loop_step}, {head_size}, + {head_size_v}, {q_tile_buffers}, {kv_tile_buffers}, NUM_COMPUTE_GROUPS, @@ -1814,6 +1819,8 @@ def encode_name(kernel_spec): qkv_layout_tag = '_qkv' elif kernel_spec.input_layout == InputLayout.Q_PAGED_KV: qkv_layout_tag = '_q_paged_kv' + elif kernel_spec.input_layout == InputLayout.SEPARATE_Q_K_V: + qkv_layout_tag = '_q_k_v' else: qkv_layout_tag = '_q_kv' # for SM90 kernels, let's also differentiate ldgsts and tma kernels @@ -2881,6 +2888,7 @@ def get_kernel_traits_code(specs_names): {loop_step}, {kv_loop_step}, {head_size}, + {head_size_v}, {q_tile_buffers}, {kv_tile_buffers}, NUM_COMPUTE_GROUPS, @@ -3092,13 +3100,13 @@ def get_cubin_header(kernel_traits, specs_names): 'tma_', '').replace('ldgsts_', '').replace('causal_', '').replace( 'alibi_', '').replace('softmax_', '').replace( - 'sliding_or_chunked_', - '').replace('custom_mask_', '').replace( - 'qkv_', '').replace('q_kv_', '').replace( - 'q_paged_kv_', '').replace('ws_', '').replace( - 'softcapping_', - '').replace('sage_', - '').replace('output_', '')) + 'sliding_or_chunked_', '').replace( + 'custom_mask_', '').replace('qkv_', '').replace( + 'q_kv_', '').replace('q_paged_kv_', '').replace( + 'q_k_v_', '').replace('ws_', '').replace( + 'softcapping_', + '').replace('sage_', + '').replace('output_', '')) flash_attention = 'flash_attention' in kname warp_specialization = 'tma_ws' in kname toks = tname.split('_') @@ -3183,6 +3191,8 @@ def get_cubin_header(kernel_traits, specs_names): attention_input_layout = InputLayout.CONTIGUOUS_Q_KV elif '_q_paged_kv' in kname: attention_input_layout = InputLayout.Q_PAGED_KV + elif '_q_k_v' in kname: + attention_input_layout = InputLayout.SEPARATE_Q_K_V attention_input_layout_value = attention_input_layout.value @@ -3418,43 +3428,7 @@ def get_lname_from_kname(kname: str) -> str: # The source code of paged context fmha kernels are not in this repo, but we have cubins for them. # Other kernels are for passing CI cases. def modify_cubin_header(cubin_header): - # for paged context fmha cases - target = "#ifndef EXCLUDE_SM_90" - - first_addition = """extern unsigned char cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_tma_ws_sm90_cu_cubin[]; -extern unsigned char cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_tma_ws_sm90_cu_cubin[];""" - - second_addition = """extern uint32_t cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_tma_ws_sm90_cu_cubin_len; -extern uint32_t cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_tma_ws_sm90_cu_cubin_len;""" - - third_addition = """{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_causal_tma_ws_sm90_kernel", 213248, 384, 64, 1, 0, false, true, true, true, false, false, false, false, nullptr}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_causal_tma_ws_sm90_kernel", 213248, 384, 64, 1, 2, false, true, true, true, false, false, false, false, nullptr},""" - result = cubin_header - offset = 0 - pos = -1 - - def add_kernel_line(result, target, addition, pos, offset): - if pos == -1: - pos = result.find(target) - else: - pos = result.find(target, pos + len(target) + offset) - if pos != -1: - end_pos = result.find('\n', pos) - if end_pos == -1: - end_pos = len(result) - result = result[:end_pos + 1] + addition + result[end_pos:] - offset += len(addition) - return result, offset, pos - - result, offset, pos = add_kernel_line(result, target, first_addition, pos, - offset) - - result, offset, pos = add_kernel_line(result, target, second_addition, pos, - offset) - - result, offset, pos = add_kernel_line(result, target, third_addition, pos, - offset) # for CI cases def add_kernel_line(result, target, addition): @@ -3672,7 +3646,8 @@ def enumerate_hgmma_flash_warpspec_kernels(specs, sm=90, dtype='fp16'): # use specialized kernels for cases without alibi scales. # there is a numeric issues when applying the exp2f scale optimization and alibi scale at the same time. combinations = product([False, True], [False, True], \ - [InputLayout.PACKED_QKV, InputLayout.CONTIGUOUS_Q_KV, InputLayout.Q_PAGED_KV], [False, True]) + [InputLayout.PACKED_QKV, InputLayout.CONTIGUOUS_Q_KV, + InputLayout.Q_PAGED_KV, InputLayout.SEPARATE_Q_K_V], [False, True]) for (alibi, return_softmax, input_layout, enable_attn_logit_softcapping) in combinations: # alibi and enable_attn_logit_softcapping shouldn't be used together. @@ -3776,6 +3751,49 @@ def enumerate_hgmma_flash_warpspec_kernels(specs, sm=90, dtype='fp16'): return_softmax_stats=return_softmax, scheduling_mode=scheduling_mode, input_layout=input_layout)) + ''' + smem size = (q_step * d * q_buffers * NUM_COMPUTE_GROUPS + + (kv_step * d + kv_step * dv) * kv_buffers) * ele_size + Originally, head size is padded to next_power_of_2 and next_power_of_2. + For fp16/bf16 context MLA (d=192/dv=128), d is padded to 256, and dv remains 128, + if kv_step=64, then smem_size = 160 KB, it is OK but wastes much smem. + if kv_step=128, then smem_size = 256 KB, it is too big for Hopper (228KB smem per SM). + But in fact, 'next multiply of 128 bytes' is needed only, due to TMA 128B swizzle mode. + Then for fp16/bf16 context MLA, d remains 192 (192 * 2 = 128 * 3), and dv remains 128, + if kv_step = 128, then smem_size = 208 KB, smem is fully utilized. + ''' + specs.append( + kernel_spec( + sm=sm, + sm_mma=90, + dtype=dtype, + seq_len=0, # support any sequence length + head_size=192, + head_size_v=128, + warps_m=4, #4x1 warpgroups + warps_n=1, + version=2, + interleaved=False, + ldgsts_q= + False, # for Hopper kernels, ldgsts = False signals TMA usage. + ldgsts_k=False, + ldgsts_v=False, + share_smem_k_v=False, + loop_step=64, + q_tile_buffers=1, # only used by warp specialized kernels + has_noloop=0, + noloop_step=64, + kv_loop_step=128, + kv_tile_buffers=2, # only used by warp specialized kernels + unroll_threshold=1, + has_scale_max=False, + flash_attention=True, + warp_specialization=True, + alibi=alibi, + enable_attn_logit_softcapping=enable_attn_logit_softcapping, + return_softmax_stats=return_softmax, + scheduling_mode=scheduling_mode, + input_layout=input_layout)) # Note this will be used in TRT-LLM. @@ -6323,6 +6341,7 @@ def enumerate_kernels(): and kspec.version == 2 and kspec.cross_mha == False and kspec.flash_attention == True + and kspec.input_layout != InputLayout.SEPARATE_Q_K_V or (kspec.sm == 90 and kspec.dtype in ['fp16', 'bf16', 'fp16_fp32'] and kspec.head_size <= 256 @@ -6341,6 +6360,18 @@ def enumerate_kernels(): and kspec.flash_attention == True and kspec.warp_specialization == False and kspec.tiled == True) + # Deepseek MLA (hopper-style context 192/128) + or (kspec.sm == 90 + and kspec.dtype == 'bf16' + and kspec.head_size == 192 + and kspec.head_size_v == 128 + and kspec.sage_block_sizes is None + and kspec.version == 2 + and kspec.cross_mha == False + and kspec.flash_attention == True + and kspec.warp_specialization == True + and kspec.alibi == False + and kspec.enable_attn_logit_softcapping == False) # SageAttention (warp_spec, head_size in (80, 128), packed QKV, padding mask) or (kspec.sm == 90 and kspec.head_size in [80, 128] diff --git a/cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv.h b/cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv.h index 642071841f4a..73d640cd9cbc 100644 --- a/cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv.h +++ b/cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv.h @@ -111,7 +111,8 @@ struct Gmem_tile_qkv inline __device__ Gmem_tile_qkv( Params const& params, int qkv_offset, Block_info const& binfo, int tidx, int cta_row_offset = 0) - : params_qkv_stride_in_bytes_(params.qkv_stride_in_bytes) + // in PACKED_QKV, q_stride = k_stride = v_stride + : params_qkv_stride_in_bytes_(params.q_stride_in_bytes) , qkv_ptr_(reinterpret_cast(params.qkv_ptr)) { @@ -132,7 +133,7 @@ struct Gmem_tile_qkv preds_[0] = fmha::pack_predicates(preds); // The row offset in the batched GEMM. For each seq element, we store QKV in that order. - int64_t row_offset = (int64_t) (row + cta_row_offset) * params.qkv_stride_in_bytes; + int64_t row_offset = (int64_t) (row + cta_row_offset) * params_qkv_stride_in_bytes_; // Add the block index. int idx; if (HEADS_INTERLEAVED) diff --git a/cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h b/cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h index d380201610a5..7e05ef3caf30 100644 --- a/cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h +++ b/cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h @@ -172,7 +172,7 @@ struct Gmem_tile_qkv template inline __device__ Gmem_tile_qkv(bert::Fused_multihead_attention_params_v2 const& params, int qkv_offset, Block_info const& binfo, int tidx, int cta_row_offset = 0, int cta_col_offset_in_bytes = 0) - : Gmem_tile_qkv(params.qkv_ptr, params.qkv_stride_in_bytes, params.d, params.dv, params.h, qkv_offset, binfo, + : Gmem_tile_qkv(params.qkv_ptr, params.q_stride_in_bytes, params.d, params.dv, params.h, qkv_offset, binfo, tidx, params.h_kv, cta_row_offset, cta_col_offset_in_bytes) { } @@ -181,7 +181,7 @@ struct Gmem_tile_qkv template inline __device__ Gmem_tile_qkv(Params const& params, int qkv_offset, Block_info const& binfo, int tidx, int cta_row_offset = 0, int cta_col_offset_in_bytes = 0) - : Gmem_tile_qkv(params.qkv_ptr, params.qkv_stride_in_bytes, params.d, params.dv, params.h, qkv_offset, binfo, + : Gmem_tile_qkv(params.qkv_ptr, params.q_stride_in_bytes, params.d, params.dv, params.h, qkv_offset, binfo, tidx, cta_row_offset, cta_col_offset_in_bytes) { } @@ -741,7 +741,7 @@ struct Gmem_tile_contiguous_kv inline __device__ Gmem_tile_contiguous_kv(bert::Fused_multihead_attention_params_v2 const& params, int qkv_offset, // q = 0, k = 1, v = 2. Block_info const& binfo, int tidx, int cta_row_offset = 0, int cta_col_offset_in_bytes = 0) - : Gmem_tile_contiguous_kv(params.kv_ptr, params.kv_stride_in_bytes, params.h_kv, params.h_q_per_kv, qkv_offset, + : Gmem_tile_contiguous_kv(params.kv_ptr, params.k_stride_in_bytes, params.h_kv, params.h_q_per_kv, qkv_offset, binfo, tidx, cta_row_offset, cta_col_offset_in_bytes) { } @@ -1070,35 +1070,11 @@ struct Gmem_tile_paged_kv // Do not load/store if the thread is in the padded area col_in_bytes_ = cta_col_offset_in_bytes + col * BYTES_PER_LDG; - // In DeepSeek, V is a prefix of K, and they share the same memory space. - // Therefore, when generating the cubin, only `kv_stride_in_bytes` field is needed. - // However, for ease of testing, the FMHA has been designed to support independent K and V, - // which requires an additional `v_stride_in_bytes` field. -#ifdef GENERATE_CUBIN - // The head offset. - head_stride_in_bytes_ = (int64_t) (binfo.bidh / params.h_q_per_kv) * params.kv_stride_in_bytes; - token_stride_in_bytes_ = BYTES_PER_ELEMENT * params.d; -#else - int64_t kv_stride_in_bytes; - if (qkv_offset == 1) - { - kv_stride_in_bytes = params.kv_stride_in_bytes; - } - else if (params.v_stride_in_bytes != 0) - { - kv_stride_in_bytes = params.v_stride_in_bytes; - } - else - { - kv_stride_in_bytes = params.kv_stride_in_bytes * params.dv / params.d; - } + int64_t kv_stride_in_bytes = qkv_offset == 1 ? params.k_stride_in_bytes : params.v_stride_in_bytes; // The head offset. head_stride_in_bytes_ = (int64_t) (binfo.bidh / params.h_q_per_kv) * kv_stride_in_bytes; - // In DeepSeek MLA, params.kv_stride_in_bytes == params.v_stride_in_bytes, - // token_stride_in_bytes_ of both K and V = d * sizeof(dtype), - // so the stride of V != VALID_BYTES_PER_ROW + // When V is padded (like MLA), we cannot use VALID_BYTES_PER_ROW token_stride_in_bytes_ = kv_stride_in_bytes >> paged_kv_log2_block_size_; -#endif // Take the CTA offset to modify the sequence length. // Actually we don't need that for flash attention. @@ -1552,7 +1528,7 @@ struct Gmem_tile_qkv_interleaved inline __device__ Gmem_tile_qkv_interleaved( Params const& params, int qkv_select, Block_info const& block_info, int tidx, int cta_row_offset = 0) : actual_seqlen_(block_info.actual_seqlen - cta_row_offset) - , total_(params.qkv_stride_in_bytes) + , total_(params.q_stride_in_bytes) , kv_ptr_(reinterpret_cast(params.qkv_ptr)) { diff --git a/cpp/kernels/fmha_v2/src/fmha/hopper/gmem_tile_o_packed.h b/cpp/kernels/fmha_v2/src/fmha/hopper/gmem_tile_o_packed.h index cda927b54d8a..75946bac612b 100644 --- a/cpp/kernels/fmha_v2/src/fmha/hopper/gmem_tile_o_packed.h +++ b/cpp/kernels/fmha_v2/src/fmha/hopper/gmem_tile_o_packed.h @@ -846,8 +846,8 @@ struct Gmem_tile_o_gmma_32bit_8bit #pragma unroll for (int di = 0; di < N_GROUPS; ++di) { - int32_t const coords[4] = {di * N_PER_GROUP, bidh_, 0, row_tma_}; - fmha::utmastg<4, fmha::cudaTmaDescType::TILED>( + const int32_t coords[3] = {di * N_PER_GROUP, bidh_, row_tma_}; + fmha::utmastg<3, fmha::cudaTmaDescType::TILED>( desc_o_, smem_base_ + di * 16 * N_BYTES_PER_GROUP, coords); } tmastg_arrive(); diff --git a/cpp/kernels/fmha_v2/src/fmha/hopper/gmem_tile_qkv_packed.h b/cpp/kernels/fmha_v2/src/fmha/hopper/gmem_tile_qkv_packed.h index 26ca608064f5..37589621d4e6 100644 --- a/cpp/kernels/fmha_v2/src/fmha/hopper/gmem_tile_qkv_packed.h +++ b/cpp/kernels/fmha_v2/src/fmha/hopper/gmem_tile_qkv_packed.h @@ -107,7 +107,8 @@ struct Gmem_tile_tma_qkv template inline __device__ Gmem_tile_tma_qkv(Params const& params, cudaTmaDesc const* p_desc, int qkv_offset, Block_info const& block_info, int tidx, int cta_row_offset = 0) - : params_qkv_stride_in_bytes_(params.qkv_stride_in_bytes) + // in PACKED_QKV, q_stride = k_stride = v_stride + : params_qkv_stride_in_bytes_(params.q_stride_in_bytes) , actual_seqlen_(block_info.actual_seqlen) , qkv_ptr_(reinterpret_cast(params.qkv_ptr)) , p_desc_(p_desc) diff --git a/cpp/kernels/fmha_v2/src/fmha/hopper/utils_hgmma.h b/cpp/kernels/fmha_v2/src/fmha/hopper/utils_hgmma.h index c03f6a9d4d01..9948d7c09516 100644 --- a/cpp/kernels/fmha_v2/src/fmha/hopper/utils_hgmma.h +++ b/cpp/kernels/fmha_v2/src/fmha/hopper/utils_hgmma.h @@ -577,6 +577,41 @@ struct Hgmma_rfa_fp16<128, TB> } }; +//////////////////////////////////////////////////////////////////////////////////////////////////// +// 64x192x16 +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct Hgmma_rfa_fp16<192, TB> +{ + static inline __device__ void mma(const uint32_t (&a)[4], uint64_t desc_b, uint32_t (&acc)[48]) + { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 && defined(__CUDA_ARCH_FEAT_SM90_ALL) + int const trans_b = TB ? 1 : 0; + asm volatile( + "wgmma.mma_async.sync.aligned.m64n192k16.f16.f16.f16 " + "{" + " %0, %1, %2, %3, %4, %5, %6, %7,\n" + " %8, %9, %10, %11, %12, %13, %14, %15,\n" + " %16, %17, %18, %19, %20, %21, %22, %23,\n" + " %24, %25, %26, %27, %28, %29, %30, %31,\n" + " %32, %33, %34, %35, %36, %37, %38, %39,\n" + " %40, %41, %42, %43, %44, %45, %46, %47 \n" + "},\n" + "{ %48, %49, %50, %51 }, %52, 1, 1, 1, %53;\n" + + : "+r"(acc[0]), "+r"(acc[1]), "+r"(acc[2]), "+r"(acc[3]), "+r"(acc[4]), "+r"(acc[5]), "+r"(acc[6]), + "+r"(acc[7]), "+r"(acc[8]), "+r"(acc[9]), "+r"(acc[10]), "+r"(acc[11]), "+r"(acc[12]), "+r"(acc[13]), + "+r"(acc[14]), "+r"(acc[15]), "+r"(acc[16]), "+r"(acc[17]), "+r"(acc[18]), "+r"(acc[19]), "+r"(acc[20]), + "+r"(acc[21]), "+r"(acc[22]), "+r"(acc[23]), "+r"(acc[24]), "+r"(acc[25]), "+r"(acc[26]), "+r"(acc[27]), + "+r"(acc[28]), "+r"(acc[29]), "+r"(acc[30]), "+r"(acc[31]), "+r"(acc[32]), "+r"(acc[33]), "+r"(acc[34]), + "+r"(acc[35]), "+r"(acc[36]), "+r"(acc[37]), "+r"(acc[38]), "+r"(acc[39]), "+r"(acc[40]), "+r"(acc[41]), + "+r"(acc[42]), "+r"(acc[43]), "+r"(acc[44]), "+r"(acc[45]), "+r"(acc[46]), "+r"(acc[47]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "l"(desc_b), "n"(trans_b)); +#endif + } +}; + //////////////////////////////////////////////////////////////////////////////////////////////////// // 64x256x16 //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -758,6 +793,54 @@ struct Hgmma_rfa_fp32<128, TB> } }; +//////////////////////////////////////////////////////////////////////////////////////////////////// +// 64x192x16 +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct Hgmma_rfa_fp32<192, TB> +{ + static inline __device__ void mma(const uint32_t (&a)[4], uint64_t desc_b, uint32_t (&acc)[96]) + { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 && defined(__CUDA_ARCH_FEAT_SM90_ALL) + int const trans_b = TB ? 1 : 0; + asm volatile( + "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16\n" + "{\n" + " %0, %1, %2, %3, %4, %5, %6, %7,\n" + " %8, %9, %10, %11, %12, %13, %14, %15,\n" + " %16, %17, %18, %19, %20, %21, %22, %23,\n" + " %24, %25, %26, %27, %28, %29, %30, %31,\n" + " %32, %33, %34, %35, %36, %37, %38, %39,\n" + " %40, %41, %42, %43, %44, %45, %46, %47,\n" + " %48, %49, %50, %51, %52, %53, %54, %55,\n" + " %56, %57, %58, %59, %60, %61, %62, %63,\n" + " %64, %65, %66, %67, %68, %69, %70, %71,\n" + " %72, %73, %74, %75, %76, %77, %78, %79,\n" + " %80, %81, %82, %83, %84, %85, %86, %87,\n" + " %88, %89, %90, %91, %92, %93, %94, %95 \n" + "},\n" + "{ %96, %97, %98, %99 }, %100, 1, 1, 1, %101;\n" + + : "+r"(acc[0]), "+r"(acc[1]), "+r"(acc[2]), "+r"(acc[3]), "+r"(acc[4]), "+r"(acc[5]), "+r"(acc[6]), + "+r"(acc[7]), "+r"(acc[8]), "+r"(acc[9]), "+r"(acc[10]), "+r"(acc[11]), "+r"(acc[12]), "+r"(acc[13]), + "+r"(acc[14]), "+r"(acc[15]), "+r"(acc[16]), "+r"(acc[17]), "+r"(acc[18]), "+r"(acc[19]), "+r"(acc[20]), + "+r"(acc[21]), "+r"(acc[22]), "+r"(acc[23]), "+r"(acc[24]), "+r"(acc[25]), "+r"(acc[26]), "+r"(acc[27]), + "+r"(acc[28]), "+r"(acc[29]), "+r"(acc[30]), "+r"(acc[31]), "+r"(acc[32]), "+r"(acc[33]), "+r"(acc[34]), + "+r"(acc[35]), "+r"(acc[36]), "+r"(acc[37]), "+r"(acc[38]), "+r"(acc[39]), "+r"(acc[40]), "+r"(acc[41]), + "+r"(acc[42]), "+r"(acc[43]), "+r"(acc[44]), "+r"(acc[45]), "+r"(acc[46]), "+r"(acc[47]), "+r"(acc[48]), + "+r"(acc[49]), "+r"(acc[50]), "+r"(acc[51]), "+r"(acc[52]), "+r"(acc[53]), "+r"(acc[54]), "+r"(acc[55]), + "+r"(acc[56]), "+r"(acc[57]), "+r"(acc[58]), "+r"(acc[59]), "+r"(acc[60]), "+r"(acc[61]), "+r"(acc[62]), + "+r"(acc[63]), "+r"(acc[64]), "+r"(acc[65]), "+r"(acc[66]), "+r"(acc[67]), "+r"(acc[68]), "+r"(acc[69]), + "+r"(acc[70]), "+r"(acc[71]), "+r"(acc[72]), "+r"(acc[73]), "+r"(acc[74]), "+r"(acc[75]), "+r"(acc[76]), + "+r"(acc[77]), "+r"(acc[78]), "+r"(acc[79]), "+r"(acc[80]), "+r"(acc[81]), "+r"(acc[82]), "+r"(acc[83]), + "+r"(acc[84]), "+r"(acc[85]), "+r"(acc[86]), "+r"(acc[87]), "+r"(acc[88]), "+r"(acc[89]), "+r"(acc[90]), + "+r"(acc[91]), "+r"(acc[92]), "+r"(acc[93]), "+r"(acc[94]), "+r"(acc[95]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "l"(desc_b), "n"(trans_b)); +#endif + } +}; + //////////////////////////////////////////////////////////////////////////////////////////////////// // 64x256x16 //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/cpp/kernels/fmha_v2/src/fmha/hopper/utils_hgmma_bf16.h b/cpp/kernels/fmha_v2/src/fmha/hopper/utils_hgmma_bf16.h index c7a5da4e6120..627d5c316bda 100644 --- a/cpp/kernels/fmha_v2/src/fmha/hopper/utils_hgmma_bf16.h +++ b/cpp/kernels/fmha_v2/src/fmha/hopper/utils_hgmma_bf16.h @@ -369,6 +369,54 @@ struct Hgmma_rfa_bf16<128, TB> } }; +//////////////////////////////////////////////////////////////////////////////////////////////////// +// 64x192x16 +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct Hgmma_rfa_bf16<192, TB> +{ + static inline __device__ void mma(uint32_t const (&a)[4], uint64_t desc_b, uint32_t (&acc)[96]) + { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 && defined(__CUDA_ARCH_FEAT_SM90_ALL) + int const trans_b = TB ? 1 : 0; + asm volatile( + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16\n" + "{\n" + " %0, %1, %2, %3, %4, %5, %6, %7,\n" + " %8, %9, %10, %11, %12, %13, %14, %15,\n" + " %16, %17, %18, %19, %20, %21, %22, %23,\n" + " %24, %25, %26, %27, %28, %29, %30, %31,\n" + " %32, %33, %34, %35, %36, %37, %38, %39,\n" + " %40, %41, %42, %43, %44, %45, %46, %47,\n" + " %48, %49, %50, %51, %52, %53, %54, %55,\n" + " %56, %57, %58, %59, %60, %61, %62, %63,\n" + " %64, %65, %66, %67, %68, %69, %70, %71,\n" + " %72, %73, %74, %75, %76, %77, %78, %79,\n" + " %80, %81, %82, %83, %84, %85, %86, %87,\n" + " %88, %89, %90, %91, %92, %93, %94, %95 \n" + "},\n" + "{ %96, %97, %98, %99 }, %100, 1, 1, 1, %101;\n" + + : "+r"(acc[0]), "+r"(acc[1]), "+r"(acc[2]), "+r"(acc[3]), "+r"(acc[4]), "+r"(acc[5]), "+r"(acc[6]), + "+r"(acc[7]), "+r"(acc[8]), "+r"(acc[9]), "+r"(acc[10]), "+r"(acc[11]), "+r"(acc[12]), "+r"(acc[13]), + "+r"(acc[14]), "+r"(acc[15]), "+r"(acc[16]), "+r"(acc[17]), "+r"(acc[18]), "+r"(acc[19]), "+r"(acc[20]), + "+r"(acc[21]), "+r"(acc[22]), "+r"(acc[23]), "+r"(acc[24]), "+r"(acc[25]), "+r"(acc[26]), "+r"(acc[27]), + "+r"(acc[28]), "+r"(acc[29]), "+r"(acc[30]), "+r"(acc[31]), "+r"(acc[32]), "+r"(acc[33]), "+r"(acc[34]), + "+r"(acc[35]), "+r"(acc[36]), "+r"(acc[37]), "+r"(acc[38]), "+r"(acc[39]), "+r"(acc[40]), "+r"(acc[41]), + "+r"(acc[42]), "+r"(acc[43]), "+r"(acc[44]), "+r"(acc[45]), "+r"(acc[46]), "+r"(acc[47]), "+r"(acc[48]), + "+r"(acc[49]), "+r"(acc[50]), "+r"(acc[51]), "+r"(acc[52]), "+r"(acc[53]), "+r"(acc[54]), "+r"(acc[55]), + "+r"(acc[56]), "+r"(acc[57]), "+r"(acc[58]), "+r"(acc[59]), "+r"(acc[60]), "+r"(acc[61]), "+r"(acc[62]), + "+r"(acc[63]), "+r"(acc[64]), "+r"(acc[65]), "+r"(acc[66]), "+r"(acc[67]), "+r"(acc[68]), "+r"(acc[69]), + "+r"(acc[70]), "+r"(acc[71]), "+r"(acc[72]), "+r"(acc[73]), "+r"(acc[74]), "+r"(acc[75]), "+r"(acc[76]), + "+r"(acc[77]), "+r"(acc[78]), "+r"(acc[79]), "+r"(acc[80]), "+r"(acc[81]), "+r"(acc[82]), "+r"(acc[83]), + "+r"(acc[84]), "+r"(acc[85]), "+r"(acc[86]), "+r"(acc[87]), "+r"(acc[88]), "+r"(acc[89]), "+r"(acc[90]), + "+r"(acc[91]), "+r"(acc[92]), "+r"(acc[93]), "+r"(acc[94]), "+r"(acc[95]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "l"(desc_b), "n"(trans_b)); +#endif + } +}; + //////////////////////////////////////////////////////////////////////////////////////////////////// // 64x256x16 //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/cpp/kernels/fmha_v2/src/fmha/hopper/utils_tma.h b/cpp/kernels/fmha_v2/src/fmha/hopper/utils_tma.h index a13b6282929f..841ab3887739 100644 --- a/cpp/kernels/fmha_v2/src/fmha/hopper/utils_tma.h +++ b/cpp/kernels/fmha_v2/src/fmha/hopper/utils_tma.h @@ -104,6 +104,19 @@ inline __device__ void utmastg(cudaTmaDesc const* p_desc, // TMA desc uint32_t smem_ptr, // src smem address int32_t const (&coord)[DIM]); // coord +// 3D, TILED +template <> +inline __device__ void utmastg<3, fmha::cudaTmaDescType::TILED>( + cudaTmaDesc const* p_desc, uint32_t smem_ptr, const int32_t (&coord)[3]) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("cp.async.bulk.tensor.3d.global.shared::cta.bulk_group [%0, {%1, %2, %3}], [%4];\n" ::"l"( + reinterpret_cast(p_desc)), + "r"(coord[0]), "r"(coord[1]), "r"(coord[2]), "r"(smem_ptr) + : "memory"); +#endif +} + // 4D, TILED template <> inline __device__ void utmastg<4, fmha::cudaTmaDescType::TILED>( diff --git a/cpp/kernels/fmha_v2/src/fmha/warpspec/compute.h b/cpp/kernels/fmha_v2/src/fmha/warpspec/compute.h index 1df784d3ed13..b95316e18485 100644 --- a/cpp/kernels/fmha_v2/src/fmha/warpspec/compute.h +++ b/cpp/kernels/fmha_v2/src/fmha/warpspec/compute.h @@ -173,7 +173,7 @@ struct Compute enum { - TILE_SIZE_V = STEP_KV * Kernel_traits::D + TILE_SIZE_V = STEP_KV * Kernel_traits::DV }; enum diff --git a/cpp/kernels/fmha_v2/src/fmha/warpspec/dma.h b/cpp/kernels/fmha_v2/src/fmha/warpspec/dma.h index cdea9428858b..42d766bfc91b 100644 --- a/cpp/kernels/fmha_v2/src/fmha/warpspec/dma.h +++ b/cpp/kernels/fmha_v2/src/fmha/warpspec/dma.h @@ -76,7 +76,7 @@ struct DMA // The tile size of V. enum { - TILE_SIZE_V = TILE_SIZE_K + TILE_SIZE_V = STEP_KV * Kernel_traits::DV }; // The tile size of V after head_dimension split. @@ -171,8 +171,6 @@ struct DMA int sum_s_q_; // The sum_s for kv. int sum_s_kv_; - // multi_query_attention (multiple heads share the same key/value). - bool multi_query_attention_; // Tile id for q tile scheduling uint32_t tile_id_; @@ -242,9 +240,6 @@ struct DMA auto headinfo_tracker0 = shared->head_info_tracker[0].createWriter(); auto headinfo_tracker1 = shared->head_info_tracker[1].createWriter(); - // When compiled for TRT-LLLM (heads_interleaved = false), this flag won't make a difference. - multi_query_attention_ = params.h_kv < params.h; - while (tile_id_ < params.num_tiles) { // If we do bidh = next_head % h, we'd guarantee b to be spread across CTAs. @@ -279,7 +274,8 @@ struct DMA } cudaTmaDesc const* desc_q = ¶ms.tma_desc_q; - cudaTmaDesc const* desc_kv = ¶ms.tma_desc_kv; + cudaTmaDesc const* desc_k = ¶ms.tma_desc_k; + cudaTmaDesc const* desc_v = ¶ms.tma_desc_v; int actual_seqlen; if (params.is_s_padded) { @@ -291,6 +287,7 @@ struct DMA sum_s_q_ = params.cu_q_seqlens[bidb]; actual_seqlen = params.cu_q_seqlens[bidb + 1] - sum_s_q_; } + sum_s_kv_ = sum_s_q_; // The cumulative packed_mask seqlens. // Each sequence length in the batch has to be padded to multiple of 128. @@ -326,11 +323,10 @@ struct DMA // Split work across N. int const kv_steps = (actual_seqlen + STEP_KV - 1) / STEP_KV; - for (int q_step_idx = 0; q_step_idx < q_steps; q_step_idx += 2) { - load_q(bidh, q_step_idx + 0 + q_step_offset, desc_q, shared->smem_q[0], cbw0); - load_q(bidh, q_step_idx + 1 + q_step_offset, desc_q, shared->smem_q[1], cbw1); + load_q(bidh, (q_step_idx + 0 + q_step_offset) * STEP_Q, desc_q, shared->smem_q[0], cbw0); + load_q(bidh, (q_step_idx + 1 + q_step_offset) * STEP_Q, desc_q, shared->smem_q[1], cbw1); // Q step bound is 2 tiles away at this moment because of 2x1 math warpgroup int const q_step_end = (q_step_idx + q_step_offset + 2) * STEP_Q - 1; @@ -342,8 +338,8 @@ struct DMA // Iterate over the kv tiles for this q step. for (int kv_step_idx = kv_idx_start; kv_step_idx < kv_idx_end; kv_step_idx++) { - int bar_id = load_kv(bidh, params.h, params.h_kv, kv_step_idx, desc_kv, shared, cbw_k, cbw_v, - cbw_v_scratch, cbr_v_scratch); + int bar_id = load_kv(bidh / params.h_q_per_kv, kv_step_idx * STEP_KV, desc_k, desc_v, shared, + cbw_k, cbw_v, cbw_v_scratch); // Opportunistically hide headinfo in the shadow of UTMALDGs of the QKV tensor if (q_step_idx == 0 && kv_step_idx == kv_idx_start) @@ -354,12 +350,12 @@ struct DMA q_tile_offset, USE_CUSTOM_MASK ? sum_mask_s : q_tile_offset, kv_steps, // q, and kv have the same length. actual_seqlen, actual_seqlen, sum_s_q_ * params.h + bidh, bidh, bidb}; - // NOTE: The need for the sync after consumer bar wait is to avoid a deadlock hazard - // when DMA thread 0 is ahead of other DMA threads. For example: - // DMA thread 0 have finished consumer bar wait phase 0 and producer bar arrive phase 0, and - // then MMA warps have finished producer bar wait phase 0 and consumer bar arrive phase 1. - // At this time other DMA threads start consumer bar wait phase 0. It will never become - // ready. DMA warps then fail to continue to the next loop. + // NOTE(tizheng): The need for the sync after consumer bar wait is to avoid a deadlock + // hazard when DMA thread 0 is ahead of other DMA threads. For example: DMA thread 0 have + // finished consumer bar wait phase 0 and producer bar arrive phase 0, and then MMA warps + // have finished producer bar wait phase 0 and consumer bar arrive phase 1. At this time + // other DMA threads start consumer bar wait phase 0. It will never become ready. DMA warps + // then fail to continue to the next loop. // // It is the same consideration for the sync after tmaReserve in load_q and load_kv // implementation below. @@ -508,9 +504,11 @@ struct DMA // Prepare the tma descriptors. cudaTmaDesc const* desc_q = ¶ms.tma_desc_q; + cudaTmaDesc const* desc_k = ¶ms.tma_desc_k; + cudaTmaDesc const* desc_v = ¶ms.tma_desc_v; + int32_t const* paged_block_offsets = params.paged_kv_cache.mBlockOffsets + bidb * 2 * params.paged_kv_cache.mMaxBlocksPerSeq; - cudaTmaDesc const* desc_kv = ¶ms.tma_desc_kv; if (SCHEDULING_MODE == 0) { @@ -549,9 +547,8 @@ struct DMA for (int q_step_idx = 0; q_step_idx < q_steps; q_step_idx += 2) { - load_separate_q(bidh, q_step_idx * STEP_Q + local_q_tile_offset, desc_q, shared->smem_q[0], cbw0); - load_separate_q( - bidh, (q_step_idx + 1) * STEP_Q + local_q_tile_offset, desc_q, shared->smem_q[1], cbw1); + load_q(bidh, q_step_idx * STEP_Q + local_q_tile_offset, desc_q, shared->smem_q[0], cbw0); + load_q(bidh, (q_step_idx + 1) * STEP_Q + local_q_tile_offset, desc_q, shared->smem_q[1], cbw1); // Q step end is 2 tiles away at this moment because of 2x1 math warpgroup int const q_step_end = (q_step_idx + 2) * STEP_Q - 1 + q_tile_offset; @@ -575,12 +572,12 @@ struct DMA bar_id = load_paged_kv(bidh_kv, remapped_kv_step_idx * STEP_KV, num_valid_kv_blocks, params.paged_kv_cache.mTokensPerBlockLog2, params.blocks_per_tma_load, params.blocks_per_tma_load_log2, params.paged_kv_cache.mMaxBlocksPerSeq, - paged_block_offsets, desc_kv, shared, cbw_k, cbw_v, cbw_v_scratch, cbr_v_scratch); + paged_block_offsets, desc_k, desc_v, shared, cbw_k, cbw_v, cbw_v_scratch); } else { - bar_id = load_contiguous_kv(bidh, params.h, params.h_kv, remapped_kv_step_idx, desc_kv, - shared, cbw_k, cbw_v, cbw_v_scratch, cbr_v_scratch); + bar_id = load_kv(bidh_kv, remapped_kv_step_idx * STEP_KV, desc_k, desc_v, shared, cbw_k, + cbw_v, cbw_v_scratch); } // Opportunistically hide headinfo in the shadow of UTMALDGs of the QKV tensor @@ -622,141 +619,90 @@ struct DMA // Load q tiles from gmem to smem by TMA. template inline __device__ void load_q( - int bidh, int q_step_idx, cudaTmaDesc const* desc_q, Smem_q& smem_q, BufferWriter& cbw) + int bidh, int q_tile_start_offset, cudaTmaDesc const* desc_q, Smem_q& smem_q, BufferWriter& cbw) { int barrier_id = cbw.tmaReserve(elect_one_, TILE_SIZE_Q * Kernel_traits::ELEMENT_BYTES); named_barrier_wait(SYNC_BARRIER, NUM_THREADS_IN_DMA_GROUP); - // coordinates: d, 3, h, s // split D into multiple groups in order to satisfy the TMA 128B sizzle mode - int32_t const q_coord_dim1 = !HEADS_INTERLEAVED || multi_query_attention_ ? bidh : 0; - int32_t const q_coord_dim2 = !HEADS_INTERLEAVED || multi_query_attention_ ? 0 : bidh; #pragma unroll for (int di = 0; di < Kernel_traits::D_GROUPS; ++di) { - int32_t const coords[4] - = {di * Kernel_traits::D_PER_GROUP, q_coord_dim1, q_coord_dim2, sum_s_q_ + q_step_idx * STEP_Q}; - fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_q, + const int32_t coords[3] = {di * Kernel_traits::D_PER_GROUP, bidh, sum_s_q_ + q_tile_start_offset}; + fmha::utmaldg<3, fmha::cudaTmaDescType::TILED, false>(desc_q, __cvta_generic_to_shared(&smem_q[barrier_id * TILE_SIZE_Q + di * TILE_SIZE_Q_PER_D_GROUP]), __cvta_generic_to_shared(cbw.barrier_ptr(barrier_id)), coords, elect_one_); } } - // Load q tiles from gmem to smem by TMA. - // Only has q tiles in this buffer, kv tiles are read from paged kv buffers. - template - inline __device__ void load_separate_q( - int bidh, int q_tile_start_offset, cudaTmaDesc const* desc_q, Smem_q& smem_q, BufferWriter& cbw) - { - - int barrier_id = cbw.tmaReserve(elect_one_, TILE_SIZE_Q * Kernel_traits::ELEMENT_BYTES); - - named_barrier_wait(SYNC_BARRIER, NUM_THREADS_IN_DMA_GROUP); - -// coordinates: d, h, 1, s -// split D into multiple groups in order to satisfy the TMA 128B sizzle mode -#pragma unroll - for (int di = 0; di < Kernel_traits::D_GROUPS; ++di) - { - int32_t const coords[4] = {di * Kernel_traits::D_PER_GROUP, bidh, 0, sum_s_q_ + q_tile_start_offset}; - fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_q, - __cvta_generic_to_shared(&smem_q[barrier_id * TILE_SIZE_Q + di * TILE_SIZE_Q_PER_D_GROUP]), - __cvta_generic_to_shared(cbw.barrier_ptr(barrier_id)), coords, elect_one_); - } - } +#define PREPARE_KV_BUFFER() \ + int k_barrier_id = cbw_k.tmaReserve(elect_one_, (TILE_SIZE_K) *Kernel_traits::ELEMENT_BYTES); \ + \ + int v_barrier_id; \ + void* v_barrier_ptr; \ + typename Kernel_traits::Element_data_type* v_smem; \ + \ + if constexpr (DMA_GROUP_TRANSPOSE_V) \ + { \ + v_barrier_id = cbw_v_scratch.tmaReserve(elect_one_, (TILE_SIZE_V) *Kernel_traits::ELEMENT_BYTES); \ + v_barrier_ptr = cbw_v_scratch.barrier_ptr(v_barrier_id); \ + v_smem = shared->smem_v_scratch.data(); \ + } \ + else \ + { \ + v_barrier_id = cbw_v.tmaReserve(elect_one_, (TILE_SIZE_V) *Kernel_traits::ELEMENT_BYTES); \ + v_barrier_ptr = cbw_v.barrier_ptr(v_barrier_id); \ + v_smem = shared->smem_v.data(); \ + } \ + \ + named_barrier_wait(SYNC_BARRIER, NUM_THREADS_IN_DMA_GROUP); // Load k,v tiles from gmem to smem by TMA. - template - inline __device__ void load_kv_impl(int bidh, int h, int h_kv, int kv_step_idx, cudaTmaDesc const* desc_kv, - Shared* shared, BufferWriter& cbw_k, BufferWriter& cbw_v) + template + inline __device__ int load_kv(int bidh_kv, int kv_tile_start_offset, cudaTmaDesc const* desc_k, + cudaTmaDesc const* desc_v, Shared* shared, BufferWriter& cbw_k, BufferWriter& cbw_v, + BufferWriterScratch& cbw_v_scratch) { + PREPARE_KV_BUFFER() - int k_barrier_id = cbw_k.tmaReserve(elect_one_, (TILE_SIZE_K) *Kernel_traits::ELEMENT_BYTES); - - int v_barrier_id = cbw_v.tmaReserve(elect_one_, (TILE_SIZE_V) *Kernel_traits::ELEMENT_BYTES); - - named_barrier_wait(SYNC_BARRIER, NUM_THREADS_IN_DMA_GROUP); - - // Coordinates: - // [d, 3, h, s] for head_interleaved, otherwise [d, h, 3, s] - // for multi_query attention, it will be [d, h + 2, 1, s] // split D into multiple groups in order to satisfy the TMA 128B sizzle mode - int32_t const k_coord_dim1 = HEADS_INTERLEAVED ? 1 : bidh; - int32_t const k_coord_dim2 = HEADS_INTERLEAVED ? bidh : 1; - int32_t const v_coord_dim1 = HEADS_INTERLEAVED ? 2 : bidh; - int32_t const v_coord_dim2 = HEADS_INTERLEAVED ? bidh : 2; - #pragma unroll for (int di = 0; di < Kernel_traits::D_GROUPS; ++di) { - int32_t const k_coords[4] - = {di * Kernel_traits::D_PER_GROUP, multi_query_attention_ ? h + bidh / (h / h_kv) : k_coord_dim1, - multi_query_attention_ ? 0 : k_coord_dim2, sum_s_q_ + kv_step_idx * STEP_KV}; + const int32_t k_coords[3] + = {di * Kernel_traits::D_PER_GROUP, bidh_kv, sum_s_kv_ + kv_tile_start_offset}; - fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_kv, + fmha::utmaldg<3, fmha::cudaTmaDescType::TILED, false>(desc_k, __cvta_generic_to_shared( &shared->smem_k[k_barrier_id * TILE_SIZE_K + di * TILE_SIZE_K_PER_D_GROUP]), __cvta_generic_to_shared(cbw_k.barrier_ptr(k_barrier_id)), k_coords, elect_one_); - - int32_t const v_coords[4] = {di * Kernel_traits::D_PER_GROUP, - multi_query_attention_ ? h + h_kv + bidh / (h / h_kv) : v_coord_dim1, - multi_query_attention_ ? 0 : v_coord_dim2, sum_s_q_ + kv_step_idx * STEP_KV}; - - fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_kv, - __cvta_generic_to_shared( - &shared->smem_v[v_barrier_id * TILE_SIZE_V + di * TILE_SIZE_V_PER_D_GROUP]), - __cvta_generic_to_shared(cbw_v.barrier_ptr(v_barrier_id)), v_coords, elect_one_); } - } - - // Load contiguous kv tiles [B, S, 2, H, D] from gmem to smem by TMA. - template - inline __device__ void load_contiguous_kv_impl(int bidh, int h, int h_kv, int kv_step_idx, - cudaTmaDesc const* desc_kv, Shared* shared, BufferWriter& cbw_k, BufferWriter& cbw_v) - { - - int k_barrier_id = cbw_k.tmaReserve(elect_one_, (TILE_SIZE_K) *Kernel_traits::ELEMENT_BYTES); - - int v_barrier_id = cbw_v.tmaReserve(elect_one_, (TILE_SIZE_V) *Kernel_traits::ELEMENT_BYTES); - - named_barrier_wait(SYNC_BARRIER, NUM_THREADS_IN_DMA_GROUP); #pragma unroll - for (int di = 0; di < Kernel_traits::D_GROUPS; ++di) + for (int di = 0; di < Kernel_traits::DV_GROUPS; ++di) { - int32_t const k_coords[4] - = {di * Kernel_traits::D_PER_GROUP, bidh / (h / h_kv), 0, sum_s_kv_ + kv_step_idx * STEP_KV}; - - fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_kv, - __cvta_generic_to_shared( - &shared->smem_k[k_barrier_id * TILE_SIZE_K + di * TILE_SIZE_K_PER_D_GROUP]), - __cvta_generic_to_shared(cbw_k.barrier_ptr(k_barrier_id)), k_coords, elect_one_); - - int32_t const v_coords[4] - = {di * Kernel_traits::D_PER_GROUP, bidh / (h / h_kv), 1, sum_s_kv_ + kv_step_idx * STEP_KV}; + const int32_t v_coords[3] + = {di * Kernel_traits::D_PER_GROUP, bidh_kv, sum_s_kv_ + kv_tile_start_offset}; - fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_kv, - __cvta_generic_to_shared( - &shared->smem_v[v_barrier_id * TILE_SIZE_V + di * TILE_SIZE_V_PER_D_GROUP]), - __cvta_generic_to_shared(cbw_v.barrier_ptr(v_barrier_id)), v_coords, elect_one_); + fmha::utmaldg<3, fmha::cudaTmaDescType::TILED, false>(desc_v, + __cvta_generic_to_shared(&v_smem[v_barrier_id * TILE_SIZE_V + di * TILE_SIZE_V_PER_D_GROUP]), + __cvta_generic_to_shared(v_barrier_ptr), v_coords, elect_one_); } + + return v_barrier_id; } - // Load k,v tiles from gmem to smem by TMA. - template - inline __device__ void load_paged_kv_impl(int bidh, int kv_tile_start_offset, int num_valid_kv_blocks, + // Load paged k,v tiles from gmem to smem by TMA. + template + inline __device__ int load_paged_kv(int bidh_kv, int kv_tile_start_offset, int num_valid_kv_blocks, int tokens_per_block_log2, int blocks_per_tma_load, int blocks_per_tma_load_log2, - int max_blocks_per_sequence, int32_t const* paged_block_offsets, cudaTmaDesc const* desc_kv, Shared* shared, - BufferWriter& cbw_k, BufferWriter& cbw_v) + int max_blocks_per_sequence, int32_t const* paged_block_offsets, cudaTmaDesc const* desc_k, + cudaTmaDesc const* desc_v, Shared* shared, BufferWriter& cbw_k, BufferWriter& cbw_v, + BufferWriterScratch& cbw_v_scratch) { - - int k_barrier_id = cbw_k.tmaReserve(elect_one_, (TILE_SIZE_K) *Kernel_traits::ELEMENT_BYTES); - - int v_barrier_id = cbw_v.tmaReserve(elect_one_, (TILE_SIZE_V) *Kernel_traits::ELEMENT_BYTES); - - named_barrier_wait(SYNC_BARRIER, NUM_THREADS_IN_DMA_GROUP); + PREPARE_KV_BUFFER() // Paged KV cache block idx. int paged_kv_block_idx = kv_tile_start_offset >> tokens_per_block_log2; @@ -770,29 +716,35 @@ struct DMA { int const bounded_block_idx = min(num_valid_kv_blocks - 1, paged_kv_block_idx + bi); - int32_t const k_paged_block_offset = paged_block_offsets[bounded_block_idx]; - int32_t const v_paged_block_offset = paged_block_offsets[max_blocks_per_sequence + bounded_block_idx]; + const int32_t k_paged_block_offset = paged_block_offsets[bounded_block_idx]; + const int32_t v_paged_block_offset = paged_block_offsets[max_blocks_per_sequence + bounded_block_idx]; #pragma unroll for (int di = 0; di < Kernel_traits::D_GROUPS; ++di) { - int32_t const k_coords[4] - = {di * Kernel_traits::D_PER_GROUP, kv_offset_in_block, bidh, k_paged_block_offset}; + const int32_t k_coords[4] + = {di * Kernel_traits::D_PER_GROUP, kv_offset_in_block, bidh_kv, k_paged_block_offset}; - fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_kv, + fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_k, __cvta_generic_to_shared(&shared->smem_k[k_barrier_id * TILE_SIZE_K + di * TILE_SIZE_K_PER_D_GROUP + bi * tile_size_k_per_block]), __cvta_generic_to_shared(cbw_k.barrier_ptr(k_barrier_id)), k_coords, elect_one_); + } - int32_t const v_coords[4] - = {di * Kernel_traits::D_PER_GROUP, kv_offset_in_block, bidh, v_paged_block_offset}; +#pragma unroll + for (int di = 0; di < Kernel_traits::DV_GROUPS; ++di) + { + const int32_t v_coords[4] + = {di * Kernel_traits::D_PER_GROUP, kv_offset_in_block, bidh_kv, v_paged_block_offset}; - fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_kv, - __cvta_generic_to_shared(&shared->smem_v[v_barrier_id * TILE_SIZE_V - + di * TILE_SIZE_V_PER_D_GROUP + bi * tile_size_k_per_block]), - __cvta_generic_to_shared(cbw_v.barrier_ptr(v_barrier_id)), v_coords, elect_one_); + fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_v, + __cvta_generic_to_shared(&v_smem[v_barrier_id * TILE_SIZE_V + di * TILE_SIZE_V_PER_D_GROUP + + bi * tile_size_k_per_block]), + __cvta_generic_to_shared(v_barrier_ptr), v_coords, elect_one_); } } + + return v_barrier_id; } template @@ -874,225 +826,6 @@ struct DMA cbr_v_scratch.pop(elect_one_); // Advance to next phase } - // Load k,v tiles from gmem to smem by TMA. - template - inline __device__ int load_kv_transpose_v_impl(int bidh, int h, int h_kv, int kv_step_idx, - cudaTmaDesc const* desc_kv, Shared* shared, BufferWriter& cbw_k, BufferWriter& cbw_v, - BufferWriterScratch& cbw_v_scratch, BufferReaderScratch& cbr_v_scratch) - { - int k_barrier_id = cbw_k.tmaReserve(elect_one_, (TILE_SIZE_K) *Kernel_traits::ELEMENT_BYTES); - - named_barrier_wait(SYNC_BARRIER, NUM_THREADS_IN_DMA_GROUP); - - // Coordinates: - // [d, 3, h, s] for head_interleaved, otherwise [d, h, 3, s] - // for multi_query attention, it will be [d, h + 2, 1, s] - // split D into multiple groups in order to satisfy the TMA 128B sizzle mode - int32_t const k_coord_dim1 = HEADS_INTERLEAVED ? 1 : bidh; - int32_t const k_coord_dim2 = HEADS_INTERLEAVED ? bidh : 1; - int32_t const v_coord_dim1 = HEADS_INTERLEAVED ? 2 : bidh; - int32_t const v_coord_dim2 = HEADS_INTERLEAVED ? bidh : 2; - -#pragma unroll - for (int di = 0; di < Kernel_traits::D_GROUPS; ++di) - { - int32_t const k_coords[4] - = {di * Kernel_traits::D_PER_GROUP, multi_query_attention_ ? h + bidh / (h / h_kv) : k_coord_dim1, - multi_query_attention_ ? 0 : k_coord_dim2, sum_s_q_ + kv_step_idx * STEP_KV}; - - fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_kv, - __cvta_generic_to_shared( - &shared->smem_k[k_barrier_id * TILE_SIZE_K + di * TILE_SIZE_K_PER_D_GROUP]), - __cvta_generic_to_shared(cbw_k.barrier_ptr(k_barrier_id)), k_coords, elect_one_); - } - - int v_scratch_barrier_id - = cbw_v_scratch.tmaReserve(elect_one_, (TILE_SIZE_V) *Kernel_traits::ELEMENT_BYTES); - -#pragma unroll - for (int di = 0; di < Kernel_traits::D_GROUPS; ++di) - { - int32_t const v_coords[4] = {di * Kernel_traits::D_PER_GROUP, - multi_query_attention_ ? h + h_kv + bidh / (h / h_kv) : v_coord_dim1, - multi_query_attention_ ? 0 : v_coord_dim2, sum_s_q_ + kv_step_idx * STEP_KV}; - - fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_kv, - __cvta_generic_to_shared( - &shared->smem_v_scratch[v_scratch_barrier_id * TILE_SIZE_V + di * TILE_SIZE_V_PER_D_GROUP]), - __cvta_generic_to_shared(cbw_v_scratch.barrier_ptr(v_scratch_barrier_id)), v_coords, elect_one_); - } - - // Do we really need this as we only have one buffer ? - return v_scratch_barrier_id; - } - - // Load contiguous kv tiles [B, S, 2, H, D] from gmem to smem by TMA. - template - inline __device__ int load_contiguous_kv_transpose_v_impl(int bidh, int h, int h_kv, int kv_step_idx, - cudaTmaDesc const* desc_kv, Shared* shared, BufferWriter& cbw_k, BufferWriter& cbw_v, - BufferWriterScratch& cbw_v_scratch, BufferReaderScratch& cbr_v_scratch) - { - int k_barrier_id = cbw_k.tmaReserve(elect_one_, (TILE_SIZE_K) *Kernel_traits::ELEMENT_BYTES); - - named_barrier_wait(SYNC_BARRIER, NUM_THREADS_IN_DMA_GROUP); - -#pragma unroll - for (int di = 0; di < Kernel_traits::D_GROUPS; ++di) - { - int32_t const k_coords[4] - = {di * Kernel_traits::D_PER_GROUP, bidh / (h / h_kv), 0, sum_s_kv_ + kv_step_idx * STEP_KV}; - - fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_kv, - __cvta_generic_to_shared( - &shared->smem_k[k_barrier_id * TILE_SIZE_K + di * TILE_SIZE_K_PER_D_GROUP]), - __cvta_generic_to_shared(cbw_k.barrier_ptr(k_barrier_id)), k_coords, elect_one_); - } - - int v_scratch_barrier_id - = cbw_v_scratch.tmaReserve(elect_one_, (TILE_SIZE_V) *Kernel_traits::ELEMENT_BYTES); - -#pragma unroll - for (int di = 0; di < Kernel_traits::D_GROUPS; ++di) - { - int32_t const v_coords[4] - = {di * Kernel_traits::D_PER_GROUP, bidh / (h / h_kv), 1, sum_s_kv_ + kv_step_idx * STEP_KV}; - - fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_kv, - __cvta_generic_to_shared( - &shared->smem_v_scratch[v_scratch_barrier_id * TILE_SIZE_V + di * TILE_SIZE_V_PER_D_GROUP]), - __cvta_generic_to_shared(cbw_v_scratch.barrier_ptr(v_scratch_barrier_id)), v_coords, elect_one_); - } - - // Do we really need this as we only have one buffer ? - return v_scratch_barrier_id; - } - - // Load paged k,v tiles from gmem to smem by TMA. - template - inline __device__ int load_paged_kv_transpose_v_impl(int bidh, int kv_tile_start_offset, - int num_valid_kv_blocks, int tokens_per_block_log2, int blocks_per_tma_load, int blocks_per_tma_load_log2, - int max_blocks_per_sequence, int32_t const* paged_block_offsets, cudaTmaDesc const* desc_kv, Shared* shared, - BufferWriter& cbw_k, BufferWriter& cbw_v, BufferWriterScratch& cbw_v_scratch, - BufferReaderScratch& cbr_v_scratch) - { - int k_barrier_id = cbw_k.tmaReserve(elect_one_, (TILE_SIZE_K) *Kernel_traits::ELEMENT_BYTES); - - int v_scratch_barrier_id - = cbw_v_scratch.tmaReserve(elect_one_, (TILE_SIZE_V) *Kernel_traits::ELEMENT_BYTES); - - named_barrier_wait(SYNC_BARRIER, NUM_THREADS_IN_DMA_GROUP); - - // Paged KV cache block idx. - int paged_kv_block_idx = kv_tile_start_offset >> tokens_per_block_log2; - int kv_offset_in_block = kv_tile_start_offset & ((1 << tokens_per_block_log2) - 1); - - // coordinates: d, s, h, 1 - int const tile_size_k_per_block = TILE_SIZE_K_PER_D_GROUP >> blocks_per_tma_load_log2; - static_assert( - TILE_SIZE_V_PER_D_GROUP == TILE_SIZE_K_PER_D_GROUP, "KV tile should have the same tensor size."); - for (int bi = 0; bi < blocks_per_tma_load; ++bi) - { - int const bounded_block_idx = min(num_valid_kv_blocks - 1, paged_kv_block_idx + bi); - - int32_t const k_paged_block_offset = paged_block_offsets[bounded_block_idx]; - int32_t const v_paged_block_offset = paged_block_offsets[max_blocks_per_sequence + bounded_block_idx]; - -#pragma unroll - for (int di = 0; di < Kernel_traits::D_GROUPS; ++di) - { - int32_t const k_coords[4] - = {di * Kernel_traits::D_PER_GROUP, kv_offset_in_block, bidh, k_paged_block_offset}; - - fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_kv, - __cvta_generic_to_shared(&shared->smem_k[k_barrier_id * TILE_SIZE_K - + di * TILE_SIZE_K_PER_D_GROUP + bi * tile_size_k_per_block]), - __cvta_generic_to_shared(cbw_k.barrier_ptr(k_barrier_id)), k_coords, elect_one_); - } - -#pragma unroll - for (int di = 0; di < Kernel_traits::D_GROUPS; ++di) - { - int32_t const v_coords[4] - = {di * Kernel_traits::D_PER_GROUP, kv_offset_in_block, bidh, v_paged_block_offset}; - - fmha::utmaldg<4, fmha::cudaTmaDescType::TILED, false>(desc_kv, - __cvta_generic_to_shared(&shared->smem_v_scratch[v_scratch_barrier_id * TILE_SIZE_V - + di * TILE_SIZE_V_PER_D_GROUP + bi * tile_size_k_per_block]), - __cvta_generic_to_shared(cbw_v_scratch.barrier_ptr(v_scratch_barrier_id)), v_coords, - elect_one_); - } - } - - // Do we really need this as we only have one buffer ? - return v_scratch_barrier_id; - } - - // Load k,v tiles from gmem to smem by TMA. - template - inline __device__ int load_kv(int bidh, int h, int h_kv, int kv_step_idx, cudaTmaDesc const* desc_kv, - Shared* shared, BufferWriter& cbw_k, BufferWriter& cbw_v, BufferWriterScratch& cbw_v_scratch, - BufferReaderScratch& cbr_v_scratch) - { - - if constexpr (DMA_GROUP_TRANSPOSE_V) - { - int v_scratch_barrier_id = load_kv_transpose_v_impl( - bidh, h, h_kv, kv_step_idx, desc_kv, shared, cbw_k, cbw_v, cbw_v_scratch, cbr_v_scratch); - return v_scratch_barrier_id; - } - else - { - load_kv_impl(bidh, h, h_kv, kv_step_idx, desc_kv, shared, cbw_k, cbw_v); - return 0; - } - } - - // Load contiguous kv tiles [B, S, 2, H, D] from gmem to smem by TMA. - template - inline __device__ int load_contiguous_kv(int bidh, int h, int h_kv, int kv_step_idx, cudaTmaDesc const* desc_kv, - Shared* shared, BufferWriter& cbw_k, BufferWriter& cbw_v, BufferWriterScratch& cbw_v_scratch, - BufferReaderScratch& cbr_v_scratch) - { - - if constexpr (DMA_GROUP_TRANSPOSE_V) - { - int v_scratch_barrier_id = load_contiguous_kv_transpose_v_impl( - bidh, h, h_kv, kv_step_idx, desc_kv, shared, cbw_k, cbw_v, cbw_v_scratch, cbr_v_scratch); - return v_scratch_barrier_id; - } - else - { - load_contiguous_kv_impl(bidh, h, h_kv, kv_step_idx, desc_kv, shared, cbw_k, cbw_v); - return 0; - } - } - - // Load paged k,v tiles from gmem to smem by TMA. - template - inline __device__ int load_paged_kv(int bidh, int kv_tile_start_offset, int num_valid_kv_blocks, - int tokens_per_block_log2, int blocks_per_tma_load, int blocks_per_tma_load_log2, - int max_blocks_per_sequence, int32_t const* paged_block_offsets, cudaTmaDesc const* desc_kv, Shared* shared, - BufferWriter& cbw_k, BufferWriter& cbw_v, BufferWriterScratch& cbw_v_scratch, - BufferReaderScratch& cbr_v_scratch) - { - - if constexpr (DMA_GROUP_TRANSPOSE_V) - { - int v_scratch_barrier_id - = load_paged_kv_transpose_v_impl(bidh, kv_tile_start_offset, num_valid_kv_blocks, - tokens_per_block_log2, blocks_per_tma_load, blocks_per_tma_load_log2, max_blocks_per_sequence, - paged_block_offsets, desc_kv, shared, cbw_k, cbw_v, cbw_v_scratch, cbr_v_scratch); - return v_scratch_barrier_id; - } - else - { - load_paged_kv_impl(bidh, kv_tile_start_offset, num_valid_kv_blocks, tokens_per_block_log2, - blocks_per_tma_load, blocks_per_tma_load_log2, max_blocks_per_sequence, paged_block_offsets, - desc_kv, shared, cbw_k, cbw_v); - return 0; - } - } - inline __device__ void get_next_tile_id( int local_wid, int tiw, uint32_t smem_tile_id, uint32_t* tile_id_counter_ptr) { @@ -1134,255 +867,173 @@ struct DMA void init_params(bert::Fused_multihead_attention_params_v2& params, bert::Fused_multihead_attention_launch_params const& launch_params, cudaStream_t stream) const { - if (launch_params.attention_input_layout == fmha::Attention_input_layout::PACKED_QKV) - { - // Packed qkv tma descriptors (continuous buffer). - fmha::Multiple_tma_descriptor<4> qkv_tma_descriptor; - - // Per batch tensor size. - uint32_t tensor_size_qkv[4]; - // Total sequence length. - int const total_seqlen = params.is_s_padded ? (params.b * params.s) : launch_params.total_q_seqlen; - tensor_size_qkv[3] = total_seqlen; - if (params.h_kv < params.h) - { - // Take MQA as non-heads-interleaved. - tensor_size_qkv[2] = 1; - tensor_size_qkv[1] = (params.h + 2 * params.h_kv); - tensor_size_qkv[0] = params.d; // params.d; - } - else if (HEADS_INTERLEAVED) - { - tensor_size_qkv[2] = params.h; - tensor_size_qkv[1] = 3; - tensor_size_qkv[0] = params.d; // params.d; - } - else - { - tensor_size_qkv[2] = 3; - tensor_size_qkv[1] = params.h; - tensor_size_qkv[0] = params.d; // params.d; - } + const uint32_t d = params.d; + const uint32_t dv = params.dv; + const uint32_t h = params.h; + const uint32_t h_kv = params.h_kv; - // O : [TOTAL, 1, h, d] - uint32_t tensor_size_o[4]; - tensor_size_o[0] = params.d; - tensor_size_o[1] = params.h; - tensor_size_o[2] = 1; - tensor_size_o[3] = total_seqlen; - - // Box size for k and v. - uint32_t box_size[4]; - // Update this on device? - box_size[2] = 1; - box_size[1] = 1; - box_size[0] = Kernel_traits::D_PER_GROUP; - - // Stride size in bytes. Assumes least significant dim is 1 (?) - uint64_t tensor_stride_qkv[3]; - tensor_stride_qkv[0] = tensor_size_qkv[0] * Kernel_traits::ELEMENT_BYTES; // d - tensor_stride_qkv[1] = tensor_size_qkv[1] * tensor_stride_qkv[0]; // d*h - tensor_stride_qkv[2] = tensor_size_qkv[2] * tensor_stride_qkv[1]; // d*h*3 - - uint64_t tensor_stride_o[3]; - tensor_stride_o[0] = tensor_size_o[0] * Kernel_traits::ELEMENT_BYTES; // d - tensor_stride_o[1] = tensor_size_o[1] * tensor_stride_o[0]; // d*h - tensor_stride_o[2] = tensor_size_o[2] * tensor_stride_o[1]; // d*h*1 - - // Traversal stride. - uint32_t traversal_stride_qkv[4] = {1, 1, 1, 1}; - uint32_t traversal_stride_o[4] = {1, 1, 1, 1}; - - // OOB fill zeros. - uint32_t oob_fill = 0; - - // FP32 to TF32 conversion disabled. - uint32_t fp32_to_tf32 = 0; - - // GMMA descriptor mode. - static constexpr int D_BYTES_PER_GROUP = Kernel_traits::D_BYTES_PER_GROUP; - static constexpr fmha::cudaTmaDescSwizzle swizzle_mode - = (D_BYTES_PER_GROUP > 64 ? fmha::cudaTmaDescSwizzle::SWIZZLE_128B - : D_BYTES_PER_GROUP > 32 ? fmha::cudaTmaDescSwizzle::SWIZZLE_64B - : fmha::cudaTmaDescSwizzle::SWIZZLE_32B); - - static_assert(STEP_KV <= 256 && STEP_Q <= 256, "max box size is 256"); - - // QKV [TOTAL, 3, h, d]. - tensor_size_qkv[3] = params.is_s_padded ? (params.b * params.s) : launch_params.total_q_seqlen; - tensor_size_o[3] = tensor_size_qkv[3]; - - // QKV ptr. - char* qkv_ptr = reinterpret_cast(params.qkv_ptr); - char* o_ptr = reinterpret_cast(params.o_ptr); - - // Desc Format (data type). - static constexpr fmha::cudaTmaDescFormat desc_format = (Kernel_traits::ELEMENT_BYTES == 1) - ? fmha::cudaTmaDescFormat::U8 - : fmha::cudaTmaDescFormat::F16_RN; - - // Q: STEP_Q. - box_size[3] = STEP_Q; - qkv_tma_descriptor.set_tma_desctriptor(qkv_ptr, desc_format, - fmha::cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, - fmha::cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_qkv, tensor_stride_qkv, - traversal_stride_qkv, box_size, oob_fill, fp32_to_tf32, ¶ms.tma_desc_q); + // Total sequence length. + const uint32_t total_seqlen = params.is_s_padded ? (params.b * params.s) : launch_params.total_q_seqlen; - // O: 16 - box_size[3] = 16; - if (Kernel_traits::USE_TMA_STORE) - { - qkv_tma_descriptor.set_tma_desctriptor(o_ptr, desc_format, - fmha::cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, - fmha::cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_o, tensor_stride_o, - traversal_stride_o, box_size, oob_fill, fp32_to_tf32, ¶ms.tma_desc_o); - } + // O Layout: [total_seqlen, H, DV] + // Per batch tensor size. + uint32_t tensor_size_o[3] = {dv, h, total_seqlen}; + + // Stride size in bytes. Assumes least significant dim is 1 + uint64_t tensor_stride_o[2] = {dv * Kernel_traits::ELEMENT_BYTES, uint64_t(params.o_stride_in_bytes)}; + + // Starting memory address + char* o_ptr = reinterpret_cast(params.o_ptr); + + // Box size of TMA + uint32_t box_size_o[3] = {Kernel_traits::D_PER_GROUP, 1, 16}; + + // Traversal stride. + uint32_t traversal_stride[3] = {1, 1, 1}; + + // OOB fill zeros. + uint32_t oob_fill = 0; + + // FP32 to TF32 conversion disabled. + uint32_t fp32_to_tf32 = 0; + + // GMMA descriptor mode. + static constexpr int D_BYTES_PER_GROUP = Kernel_traits::D_BYTES_PER_GROUP; + static constexpr fmha::cudaTmaDescSwizzle swizzle_mode + = (D_BYTES_PER_GROUP > 64 ? fmha::cudaTmaDescSwizzle::SWIZZLE_128B + : D_BYTES_PER_GROUP > 32 ? fmha::cudaTmaDescSwizzle::SWIZZLE_64B + : fmha::cudaTmaDescSwizzle::SWIZZLE_32B); - // K: STEP_KV. - box_size[3] = STEP_KV; - qkv_tma_descriptor.set_tma_desctriptor(qkv_ptr, desc_format, + static_assert(STEP_KV <= 256 && STEP_Q <= 256, "max box size is 256"); + + // Desc Format (data type). + static constexpr fmha::cudaTmaDescFormat desc_format + = (Kernel_traits::ELEMENT_BYTES == 1) ? fmha::cudaTmaDescFormat::U8 : fmha::cudaTmaDescFormat::F16_RN; + + fmha::Multiple_tma_descriptor<3> qo_tma_descriptor; + + // TMA O + if (Kernel_traits::USE_TMA_STORE) + { + qo_tma_descriptor.set_tma_desctriptor(o_ptr, desc_format, fmha::cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, - fmha::cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_qkv, tensor_stride_qkv, - traversal_stride_qkv, box_size, oob_fill, fp32_to_tf32, ¶ms.tma_desc_kv); + fmha::cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_o, tensor_stride_o, traversal_stride, + box_size_o, oob_fill, fp32_to_tf32, ¶ms.tma_desc_o); } - else - { - // Separate contiguous q, contiguous kv, and paged kv tma descriptors. - fmha::Multiple_tma_descriptor<4> qo_tma_descriptor; - fmha::Multiple_tma_descriptor<4> contiguous_kv_tma_descriptor; - fmha::Multiple_tma_descriptor<4> paged_kv_tma_descriptor; - // params.b * 2 * params.paged_kv_cache.mMaxBlocksPerSeq - // Per batch tensor size. - uint32_t tensor_size_qo[4]; - tensor_size_qo[3] = params.is_s_padded ? params.b * params.s : launch_params.total_q_seqlen; - tensor_size_qo[2] = 1; - tensor_size_qo[1] = params.h; - tensor_size_qo[0] = params.d; // params.d; - - // Box size for q and o. - uint32_t box_size_qo[4]; - box_size_qo[3] = STEP_Q; - box_size_qo[2] = 1; - box_size_qo[1] = 1; - box_size_qo[0] = Kernel_traits::D_PER_GROUP; - - // Stride size in bytes. Assumes least significant dim is 1 (?) - uint64_t tensor_stride_qo[3]; - tensor_stride_qo[0] = tensor_size_qo[0] * Kernel_traits::ELEMENT_BYTES; // d - tensor_stride_qo[1] = tensor_size_qo[1] * tensor_stride_qo[0]; // d*h - tensor_stride_qo[2] = tensor_size_qo[2] * tensor_stride_qo[1]; // d*h*3 - - // Traversal stride. - uint32_t traversal_stride[4] = {1, 1, 1, 1}; - // OOB fill zeros. - uint32_t oob_fill = 0; + auto const layout = launch_params.attention_input_layout; - // FP32 to TF32 conversion disabled. - uint32_t fp32_to_tf32 = 0; + // Q always uses 3D tensor + uint32_t tensor_size_q[3] = {d, h, total_seqlen}; - // GMMA descriptor mode. - static constexpr int D_BYTES_PER_GROUP = Kernel_traits::D_BYTES_PER_GROUP; - static constexpr fmha::cudaTmaDescSwizzle swizzle_mode - = (D_BYTES_PER_GROUP > 64 ? fmha::cudaTmaDescSwizzle::SWIZZLE_128B - : D_BYTES_PER_GROUP > 32 ? fmha::cudaTmaDescSwizzle::SWIZZLE_64B - : fmha::cudaTmaDescSwizzle::SWIZZLE_32B); + uint64_t tensor_stride_q[2] = {d * Kernel_traits::ELEMENT_BYTES, uint64_t(params.q_stride_in_bytes)}; - static_assert(STEP_KV <= 256 && STEP_Q <= 256, "max box size is 256"); + char* q_ptr = reinterpret_cast( + layout == fmha::Attention_input_layout::PACKED_QKV ? params.qkv_ptr : params.q_ptr); - // Q ptr. - char* q_ptr = reinterpret_cast(params.q_ptr); + uint32_t box_size_q[3] = {Kernel_traits::D_PER_GROUP, 1, STEP_Q}; - // Desc Format (data type). - static constexpr fmha::cudaTmaDescFormat desc_format = (Kernel_traits::ELEMENT_BYTES == 1) - ? fmha::cudaTmaDescFormat::U8 - : fmha::cudaTmaDescFormat::F16_RN; + if (layout == fmha::Attention_input_layout::Q_PAGED_KV) + { + // KV in q_paged_kv uses 4D tensor + // Layout: [INT32_MAX, H_KV, TokensPerBlock, D] + const uint32_t tokens_per_block = params.paged_kv_cache.mTokensPerBlock; + uint32_t tensor_size_k[4] = {d, tokens_per_block, h_kv, INT_MAX}; + uint32_t tensor_size_v[4] = {dv, tokens_per_block, h_kv, INT_MAX}; + + uint64_t tensor_stride_k[3]; + tensor_stride_k[0] = params.k_stride_in_bytes / tokens_per_block; // d + tensor_stride_k[1] = params.k_stride_in_bytes; // d * 64 + tensor_stride_k[2] = params.paged_kv_cache.mBytesPerBlock; + uint64_t tensor_stride_v[3]; + // we cannot use dv * Kernel_traits::ELEMENT_BYTES because V may be padded (MLA) + tensor_stride_v[0] = params.v_stride_in_bytes / tokens_per_block; // dv + tensor_stride_v[1] = params.v_stride_in_bytes; // dv * 64 + tensor_stride_v[2] = params.paged_kv_cache.mBytesPerBlock; + + char* kv_ptr = reinterpret_cast(params.paged_kv_cache.mPoolPtr); + + uint32_t box_size_kv[4] + = {Kernel_traits::D_PER_GROUP, std::min(tokens_per_block, STEP_KV), 1, 1}; + + assert(STEP_KV % tokens_per_block == 0 || tokens_per_block % STEP_KV == 0); + params.blocks_per_tma_load = std::max(1, STEP_KV / tokens_per_block); + params.blocks_per_tma_load_log2 = log2(params.blocks_per_tma_load); + + uint32_t traversal_stride[4] = {1, 1, 1, 1}; - // Q: STEP_Q. - qo_tma_descriptor.set_tma_desctriptor(q_ptr, desc_format, + fmha::Multiple_tma_descriptor<4> kv_tma_descriptor; + // K + kv_tma_descriptor.set_tma_desctriptor(kv_ptr, desc_format, + fmha::cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, + fmha::cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_k, tensor_stride_k, traversal_stride, + box_size_kv, oob_fill, fp32_to_tf32, ¶ms.tma_desc_k); + // V + kv_tma_descriptor.set_tma_desctriptor(kv_ptr, desc_format, fmha::cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, - fmha::cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_qo, tensor_stride_qo, traversal_stride, - box_size_qo, oob_fill, fp32_to_tf32, ¶ms.tma_desc_q); + fmha::cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_v, tensor_stride_v, traversal_stride, + box_size_kv, oob_fill, fp32_to_tf32, ¶ms.tma_desc_v); + } + else + { + // Otherwise KV uses 3D tensor + uint32_t tensor_size_k[3] = {d, h_kv, total_seqlen}; + uint32_t tensor_size_v[3] = {dv, h_kv, total_seqlen}; - // O ptr. - char* o_ptr = reinterpret_cast(params.o_ptr); + uint64_t tensor_stride_k[2] = {d * Kernel_traits::ELEMENT_BYTES, uint64_t(params.k_stride_in_bytes)}; + uint64_t tensor_stride_v[2] = {dv * Kernel_traits::ELEMENT_BYTES, uint64_t(params.v_stride_in_bytes)}; - // O: 16 - box_size_qo[3] = 16; - if (Kernel_traits::USE_TMA_STORE) + uint32_t box_size_kv[3] = {Kernel_traits::D_PER_GROUP, 1, STEP_KV}; + + char *k_ptr, *v_ptr; + + if (layout == fmha::Attention_input_layout::PACKED_QKV) { - qo_tma_descriptor.set_tma_desctriptor(o_ptr, desc_format, - fmha::cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, - fmha::cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_qo, tensor_stride_qo, - traversal_stride, box_size_qo, oob_fill, fp32_to_tf32, ¶ms.tma_desc_o); + if (!HEADS_INTERLEAVED || h != h_kv) + { + // Layout: [total_seqlen, (H, D) + (H_KV, D) + (H_KV, DV)] + // All of MHA in TRTLLM is in this layout, + // and MQA/GQA must use this layout. + k_ptr = q_ptr + h * d * Kernel_traits::ELEMENT_BYTES; + v_ptr = k_ptr + h_kv * d * Kernel_traits::ELEMENT_BYTES; + } + else + { + // Layout: [total_seqlen, H, D + D + DV] + // Currently only used in MHA in fmha_v2 tests. + tensor_stride_q[0] = tensor_stride_k[0] = tensor_stride_v[0] + = (2 * d + dv) * Kernel_traits::ELEMENT_BYTES; + k_ptr = q_ptr + d * Kernel_traits::ELEMENT_BYTES; + v_ptr = k_ptr + d * Kernel_traits::ELEMENT_BYTES; + } } - - // Contiguous KV: [B, S, 2, H, D]. - if (launch_params.attention_input_layout == fmha::Attention_input_layout::CONTIGUOUS_Q_KV) + else if (layout == fmha::Attention_input_layout::CONTIGUOUS_Q_KV) { - - // Total sequence length. - int const total_seqlen = params.is_s_padded ? (params.b * params.s) : launch_params.total_kv_seqlen; - uint32_t tensor_size_kv[4]; - tensor_size_kv[3] = total_seqlen; - tensor_size_kv[2] = 2; - tensor_size_kv[1] = params.h_kv; - tensor_size_kv[0] = params.d; - - // Box size for k and v. - uint32_t box_size_kv[4]; - box_size_kv[3] = int32_t(STEP_KV); - box_size_kv[2] = 1; - box_size_kv[1] = 1; - box_size_kv[0] = Kernel_traits::D_PER_GROUP; - - // Stride size in bytes. Assumes least significant dim is 1 (?) - uint64_t tensor_stride_kv[3]; - tensor_stride_kv[0] = tensor_size_kv[0] * Kernel_traits::ELEMENT_BYTES; // d - tensor_stride_kv[1] = tensor_size_kv[1] * tensor_stride_kv[0]; // d*h_kv - tensor_stride_kv[2] = tensor_size_kv[2] * tensor_stride_kv[1]; // d*h_kv*2 - - // Contiguous KV pool tma descriptors. - contiguous_kv_tma_descriptor.set_tma_desctriptor(reinterpret_cast(params.kv_ptr), - desc_format, fmha::cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, - fmha::cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_kv, tensor_stride_kv, - traversal_stride, box_size_kv, oob_fill, fp32_to_tf32, ¶ms.tma_desc_kv); + k_ptr = reinterpret_cast(params.kv_ptr); + v_ptr = k_ptr + h_kv * d * Kernel_traits::ELEMENT_BYTES; } - else + else if (layout == fmha::Attention_input_layout::SEPARATE_Q_K_V) { - // Paged KV: [UINT32_MAX, H, TokensPerBlock, D] - // Per batch tensor size. - uint32_t tensor_size_kv[4]; - tensor_size_kv[3] = params.b * 2 * params.paged_kv_cache.mMaxBlocksPerSeq; - tensor_size_kv[2] = params.h_kv; - tensor_size_kv[1] = params.paged_kv_cache.mTokensPerBlock; - tensor_size_kv[0] = params.d; // params.d; - - // Box size for k and v. - uint32_t box_size_kv[4]; - box_size_kv[3] = 1; - box_size_kv[2] = 1; - box_size_kv[1] = std::min(params.paged_kv_cache.mTokensPerBlock, int32_t(STEP_KV)); - box_size_kv[0] = Kernel_traits::D_PER_GROUP; - - assert(int32_t(STEP_KV) % params.paged_kv_cache.mTokensPerBlock == 0 - || params.paged_kv_cache.mTokensPerBlock % int32_t(STEP_KV) == 0); - params.blocks_per_tma_load = std::max(1, int32_t(STEP_KV) / params.paged_kv_cache.mTokensPerBlock); - params.blocks_per_tma_load_log2 = log2(params.blocks_per_tma_load); - - // Stride size in bytes. Assumes least significant dim is 1 (?) - uint64_t tensor_stride_kv[3]; - tensor_stride_kv[0] = tensor_size_kv[0] * Kernel_traits::ELEMENT_BYTES; // d - tensor_stride_kv[1] = tensor_size_kv[1] * tensor_stride_kv[0]; // d*h - tensor_stride_kv[2] = tensor_size_kv[2] * tensor_stride_kv[1]; // d*h*3 - - // Paged KV pool tma descriptors. - paged_kv_tma_descriptor.set_tma_desctriptor(reinterpret_cast(params.paged_kv_cache.mPoolPtr), - desc_format, fmha::cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, - fmha::cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_kv, tensor_stride_kv, - traversal_stride, box_size_kv, oob_fill, fp32_to_tf32, ¶ms.tma_desc_kv); + k_ptr = reinterpret_cast(params.k_ptr); + v_ptr = reinterpret_cast(params.v_ptr); } + + fmha::Multiple_tma_descriptor<3> kv_tma_descriptor; + // K + kv_tma_descriptor.set_tma_desctriptor(k_ptr, desc_format, + fmha::cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, + fmha::cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_k, tensor_stride_k, traversal_stride, + box_size_kv, oob_fill, fp32_to_tf32, ¶ms.tma_desc_k); + // V + kv_tma_descriptor.set_tma_desctriptor(v_ptr, desc_format, + fmha::cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, + fmha::cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_v, tensor_stride_v, traversal_stride, + box_size_kv, oob_fill, fp32_to_tf32, ¶ms.tma_desc_v); } + // Q + qo_tma_descriptor.set_tma_desctriptor(q_ptr, desc_format, fmha::cudaTmaDescInterleave::INTERLEAVE_DISABLED, + swizzle_mode, fmha::cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_q, tensor_stride_q, + traversal_stride, box_size_q, oob_fill, fp32_to_tf32, ¶ms.tma_desc_q); } }; }; diff --git a/cpp/kernels/fmha_v2/src/fmha/warpspec/kernel_traits.h b/cpp/kernels/fmha_v2/src/fmha/warpspec/kernel_traits.h index 0e5c208b71f4..8c93ce8a9885 100644 --- a/cpp/kernels/fmha_v2/src/fmha/warpspec/kernel_traits.h +++ b/cpp/kernels/fmha_v2/src/fmha/warpspec/kernel_traits.h @@ -36,6 +36,8 @@ template < int STEP_KV_, // The head dimension. int D_, + // The head dimension of V. + int DV_, // The number of smem buffers for Q tiles. int Q_BUFFERS_, // The number of smem buffers for K, and V tiles. @@ -83,16 +85,15 @@ struct Kernel_traits STEP_KV = STEP_KV_ }; - // The padded head dimension. + // The valid head dimension. enum { - D = Next_power_of_two::VALUE + VALID_D = D_ }; - // The valid head dimension. enum { - VALID_D = D_ + VALID_DV = (DV_ == 0 ? D_ : DV_) }; // Bootstrap GMMA_K from dummy Instruction_traits where FP16/BF16 K = 16, FP8 K = 32. @@ -113,6 +114,17 @@ struct Kernel_traits ELEMENT_BYTES = sizeof(Element_data_type) }; + // The padded head dimension. + enum + { + D = std::min(Round_up::VALUE, Next_power_of_two::VALUE) + }; + + enum + { + DV = std::min(Round_up::VALUE, Next_power_of_two::VALUE) + }; + // The number of smem buffers for Q tiles. enum { @@ -326,6 +338,18 @@ struct Kernel_traits D_BYTES_PER_GROUP = D_BYTES / D_GROUPS }; + // The bytes of head dimension of V. + enum + { + DV_BYTES = DV * ELEMENT_BYTES + }; + + // The number of head_dimension groups of V. + enum + { + DV_GROUPS = fmha::Div_up::VALUE + }; + // QGMMA: BMM2 will be split into multiple K groups as we explicitly transpose v (128 * D) in the smem. // HGMMA: BMM2 will load from row-major (K * N) smem_v, so we don't need to explicitly split K. static constexpr auto BMM2_LEADING_DIM_BYTES = ELEMENT_BYTES == 1 ? 128 : STEP_KV * ELEMENT_BYTES; @@ -364,7 +388,7 @@ struct Kernel_traits // The instruction traits for the BMM2. // FP16/BF16 K = 16, FP8 K = 32. - using Traits_o = Instruction_traits; + using Traits_o = Instruction_traits; // The CTA description for BMM1. using Cta_tile_p = @@ -375,7 +399,7 @@ struct Kernel_traits typename Traits_p::template Cta_tile; // The CTA description for BMM2. - using Cta_tile_o = typename Traits_o::template Cta_padded_tile; // The MMA tile for the 1st GEMM. @@ -415,9 +439,9 @@ struct Kernel_traits // The q, k, v tile buffer. using Buffer_q_t = cuda::std::array; using Buffer_k_t = cuda::std::array; - using Buffer_v_t = cuda::std::array; + using Buffer_v_t = cuda::std::array; // We need one kv buffer to explicitly transose fp8 smem_tile. - using Buffer_v_scratch_t = cuda::std::array; + using Buffer_v_scratch_t = cuda::std::array; // The smem bytes of q, k, v tiles. enum @@ -521,6 +545,8 @@ template < // The step size in query sequence dimension (M of BMM1 and BMM2). int STEP_KV_, // The head dimension. int D_, + // The head dimension of V. + int DV_, // The number of smem buffers for Q tiles. int Q_BUFFERS_, // The number of smem buffers for K, and V tiles. @@ -554,14 +580,14 @@ template < // The step size in query sequence dimension (M of BMM1 and BMM2). // The sage attention block size for Q, K and V int SAGE_BLOCK_SIZE_Q_ = 0, int SAGE_BLOCK_SIZE_K_ = 0, int SAGE_BLOCK_SIZE_V_ = 0> struct Kernel_traits_Hopper_qgmma_e4m3_fp32 - : public Kernel_traits { // Base class. - using Base = Kernel_traits; @@ -601,7 +627,7 @@ struct Kernel_traits_Hopper_qgmma_e4m3_fp32 using Buffer_v_scratch_t = typename Base::Buffer_v_scratch_t; // Extra O buffer if TMA is used for epilogue using Element_data_type = typename Base::Element_data_type; - using Buffer_o_t = cuda::std::array; + using Buffer_o_t = cuda::std::array; // The struct of shared memory buffers. struct __align__(128) Shared diff --git a/cpp/kernels/fmha_v2/src/fused_multihead_attention.cpp b/cpp/kernels/fmha_v2/src/fused_multihead_attention.cpp index 182df74d2e59..e2640241db48 100644 --- a/cpp/kernels/fmha_v2/src/fused_multihead_attention.cpp +++ b/cpp/kernels/fmha_v2/src/fused_multihead_attention.cpp @@ -250,6 +250,10 @@ static inline void set_params(bert::Fused_multihead_attention_params_v2& params, void* qkv_packed_d, // contiguous q. void* q_d, + // separate k. + void* k_d, + // separate v. + void* v_d, // contiguous kv. void* kv_d, // start address of the paged kv pool. @@ -267,42 +271,57 @@ static inline void set_params(bert::Fused_multihead_attention_params_v2& params, memset(¶ms, 0, sizeof(params)); - // Set the pointers. - params.qkv_ptr = qkv_packed_d; - // For grouped- or multi-query attention (h denotes num_q_heads; h' denotes h_kv): - // qkv_layout = [b, s, [q_hd, k_h'd, v_h'd]] - // qkv_stride = (h+2*h')d * bytes_per_elt - // Otherwise: - // qkv_layout = [b, s, 3, h, d] or [b, s, h, 3, d] - // qkv_stride = 3hd * bytes_per_elt - params.qkv_stride_in_bytes = get_size_in_bytes(h * d + h_kv * d + h_kv * dv, data_type); params.o_ptr = o_packed_d; params.o_stride_in_bytes = get_size_in_bytes(h * dv, output_dtype); if (interleaved) { - params.qkv_stride_in_bytes = total; + params.q_stride_in_bytes = total; params.o_stride_in_bytes = total; } - // Contiguous q + Paged kv cache. - int max_blocks_per_sequence = (s_kv + tokens_per_block - 1) / tokens_per_block; - params.paged_kv_cache = Kv_block_array(b, max_blocks_per_sequence, tokens_per_block, - get_size_in_bytes(tokens_per_block * h_kv * std::gcd(d, dv), data_type), paged_kv_pool_ptr); - params.paged_kv_cache.mBlockOffsets = paged_block_offsets; - params.q_stride_in_bytes = get_size_in_bytes(h * d, data_type); - // Layout [B, S, H, D]. - params.q_ptr = q_d; - // Layout [B, S, 2, H, D]. - params.kv_ptr = kv_d; - if (input_layout == Attention_input_layout::Q_PAGED_KV) + if (input_layout == Attention_input_layout::PACKED_QKV) { - params.kv_stride_in_bytes = get_size_in_bytes(tokens_per_block * d, data_type); - params.v_stride_in_bytes = get_size_in_bytes(tokens_per_block * dv, data_type); + // For grouped- or multi-query attention (h denotes num_q_heads; h' denotes h_kv): + // qkv_layout = [b, s, [q_hd, k_h'd, v_h'd]] + // qkv_stride = (h+2*h')d * bytes_per_elt + // Otherwise: + // qkv_layout = [b, s, 3, h, d] or [b, s, h, 3, d] + // qkv_stride = 3hd * bytes_per_elt + params.qkv_ptr = qkv_packed_d; + params.q_stride_in_bytes = params.k_stride_in_bytes = params.v_stride_in_bytes + = get_size_in_bytes(h * d + h_kv * d + h_kv * dv, data_type); } else { - params.kv_stride_in_bytes = get_size_in_bytes(2 * h_kv * d, data_type); + // Layout [B, S, H, D]. + params.q_ptr = q_d; + params.q_stride_in_bytes = get_size_in_bytes(h * d, data_type); + + if (input_layout == Attention_input_layout::CONTIGUOUS_Q_KV) + { + // Layout [B, S, 2, H, D]. + params.kv_ptr = kv_d; + params.k_stride_in_bytes = params.v_stride_in_bytes = get_size_in_bytes(h_kv * (d + dv), data_type); + } + else if (input_layout == Attention_input_layout::Q_PAGED_KV) + { + int max_blocks_per_sequence = (s_kv + tokens_per_block - 1) / tokens_per_block; + params.paged_kv_cache = Kv_block_array(b, max_blocks_per_sequence, tokens_per_block, + get_size_in_bytes(tokens_per_block * h_kv * std::gcd(d, dv), data_type), paged_kv_pool_ptr); + params.paged_kv_cache.mBlockOffsets = paged_block_offsets; + params.k_stride_in_bytes = get_size_in_bytes(tokens_per_block * d, data_type); + params.v_stride_in_bytes = get_size_in_bytes(tokens_per_block * dv, data_type); + } + else if (input_layout == Attention_input_layout::SEPARATE_Q_K_V) + { + // Layout [B, S, H_kv, D]. + params.k_ptr = k_d; + // Layout [B, S, H_kv, Dv]. + params.v_ptr = v_d; + params.k_stride_in_bytes = get_size_in_bytes(h_kv * d, data_type); + params.v_stride_in_bytes = get_size_in_bytes(h_kv * dv, data_type); + } } // Packed mask. @@ -756,6 +775,10 @@ int main(int argc, char** argv) { input_layout = Attention_input_layout::Q_PAGED_KV; } + else if (!strcmp(argv[ii], "-separate-q-k-v")) + { + input_layout = Attention_input_layout::SEPARATE_Q_K_V; + } else if (!strcmp(argv[ii], "-tokens-per-block") && ++ii < argc) { tokens_per_block = strtol(argv[ii], nullptr, 10); @@ -1032,7 +1055,7 @@ int main(int argc, char** argv) // Contiguous KV cache buffer. // The shape is [B, 2, S, H, D]. - size_t const kv_size = b * 2 * s * h_kv * d; + const size_t kv_size = b * s * h_kv * (d + dv); // The size in bytes. size_t const kv_size_in_bytes = get_size_in_bytes(kv_size, data_type); // Allocate on the host. @@ -1084,6 +1107,16 @@ int main(int argc, char** argv) size_t const q_size = s * b * h * d; FMHA_CHECK_CUDA(cudaMalloc(&q_d, get_size_in_bytes(q_size, data_type))); + // K has [B, S, H_kv, D] with separate kv cache. + void* k_d; + const size_t k_size = s * b * h_kv * d; + FMHA_CHECK_CUDA(cudaMalloc(&k_d, get_size_in_bytes(k_size, data_type))); + + // V has [B, S, H_kv, Dv] with separate kv cache. + void* v_d; + const size_t v_size = s * b * h_kv * dv; + FMHA_CHECK_CUDA(cudaMalloc(&v_d, get_size_in_bytes(v_size, data_type))); + // Scale bmm2 (per-tensor). void* scale_bmm2_d; FMHA_CHECK_CUDA(cudaMalloc(&scale_bmm2_d, sizeof(uint32_t))); @@ -1499,8 +1532,8 @@ int main(int argc, char** argv) // "Padded MQA V[b, s, h_kv*d]"); // } - // Contiguous KV Cache. - store_q_and_contiguous_kv_cache(q_d, contiguous_kv_h, contiguous_kv_d, + // Contiguous KV Cache and Separate KV Cache. + store_q_and_contiguous_kv_cache(q_d, k_d, v_d, contiguous_kv_h, contiguous_kv_d, reinterpret_cast(qkv_packed_h.data()), reinterpret_cast(cu_seqlens.data()), reinterpret_cast(cu_q_seqlens.data()), b, s, h, h_kv, d, dv, data_type); @@ -1642,9 +1675,10 @@ int main(int argc, char** argv) set_params(params_v2, launch_params, data_type, acc_type, output_dtype, input_layout, b, s_q, s, h, h_kv, d, dv, total, num_grouped_heads, sliding_window_size, chunked_attention_size, // Paged kv cache. - tokens_per_block, qkv_d_view, q_d, contiguous_kv_d, kv_cache_pool_ptr, kv_cache_block_offsets_d, packed_mask_d, - cu_mask_rows_d, cu_seqlens_d, cu_q_seqlens_d, o_d_view, p_d, s_d, softmax_stats_ptr, scale_bmm2_d, scale_bmm1, - scale_softmax, scale_bmm2, softcapping_scale_bmm1, use_int8_scale_max, interleaved, is_s_padded, has_alibi); + tokens_per_block, qkv_d_view, q_d, k_d, v_d, contiguous_kv_d, kv_cache_pool_ptr, kv_cache_block_offsets_d, + packed_mask_d, cu_mask_rows_d, cu_seqlens_d, cu_q_seqlens_d, o_d_view, p_d, s_d, softmax_stats_ptr, + scale_bmm2_d, scale_bmm1, scale_softmax, scale_bmm2, softcapping_scale_bmm1, use_int8_scale_max, interleaved, + is_s_padded, has_alibi); // total number of tokens is needed to set TMA desc on the host. launch_params.total_q_seqlen = q_seqlens[b]; @@ -1753,10 +1787,12 @@ int main(int argc, char** argv) #else { // use external quant kernel - int const stride_qkv = params_v2.qkv_stride_in_bytes; run_sage_quant(b, h, d, s, params_v2.qkv_ptr, (char*) params_v2.qkv_ptr + get_size_in_bytes(h * d, data_type), - (char*) params_v2.qkv_ptr + get_size_in_bytes(2 * h * d, data_type), stride_qkv, stride_qkv, stride_qkv, + (char*) params_v2.qkv_ptr + get_size_in_bytes(2 * h * d, data_type, + params_v2.q_stride_in_bytes, + params_v2.k_stride_in_bytes, + params_v2.v_stride_in_bytes, params_v2.cu_q_seqlens, params_v2.cu_kv_seqlens, sage_block_size_q, sage_block_size_k, sage_block_size_v, quant_qkv, quant_qkv + h * d, quant_qkv + 2 * h * d, params_v2.sage.q.scales, params_v2.sage.k.scales, params_v2.sage.v.scales); @@ -1764,7 +1800,8 @@ int main(int argc, char** argv) #endif // no need to free old params_v2.qkv_ptr, it will be released in the end params_v2.qkv_ptr = quant_qkv; - params_v2.qkv_stride_in_bytes = get_size_in_bytes((h + 2 * h_kv) * d, DATA_TYPE_E4M3); + params_v2.q_stride_in_bytes = params_v2.k_stride_in_bytes = params_v2.v_stride_in_bytes + = get_size_in_bytes((h + 2 * h_kv) * d, DATA_TYPE_E4M3); } #if defined(DEBUG_HAS_PRINT_BUFFER) @@ -2052,6 +2089,9 @@ int main(int argc, char** argv) FMHA_CHECK_CUDA(cudaFree(qkv_bsh3d_d)); FMHA_CHECK_CUDA(cudaFree(mask_d)); FMHA_CHECK_CUDA(cudaFree(packed_mask_d)); + FMHA_CHECK_CUDA(cudaFree(q_d)); + FMHA_CHECK_CUDA(cudaFree(k_d)); + FMHA_CHECK_CUDA(cudaFree(v_d)); FMHA_CHECK_CUDA(cudaFree(p_d)); FMHA_CHECK_CUDA(cudaFree(s_d)); FMHA_CHECK_CUDA(cudaFree(o_d)); diff --git a/cpp/kernels/fmha_v2/src/fused_multihead_attention.h b/cpp/kernels/fmha_v2/src/fused_multihead_attention.h index 33610dca7812..f77e3f14d0c4 100644 --- a/cpp/kernels/fmha_v2/src/fused_multihead_attention.h +++ b/cpp/kernels/fmha_v2/src/fused_multihead_attention.h @@ -74,6 +74,10 @@ enum class Attention_input_layout // of [B, 2, Blocks_per_Seq], and the indice indicates the block distance to the pool ptr in // global memory. Q_PAGED_KV, + // Q has [B, S, H, D] layout, + // K has [B, S, H_kv, D] layout, + // V has [B, S, H_kv, Dv] layout, + SEPARATE_Q_K_V, }; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -85,6 +89,7 @@ static inline std::string attention_input_layout_to_string(Attention_input_layou case Attention_input_layout::PACKED_QKV: return "packed_qkv"; case Attention_input_layout::CONTIGUOUS_Q_KV: return "contiguous_q_kv"; case Attention_input_layout::Q_PAGED_KV: return "contiguous_q_paged_kv"; + case Attention_input_layout::SEPARATE_Q_K_V: return "separate_q_k_v"; default: assert(false); return ""; } } @@ -114,8 +119,6 @@ struct Fused_multihead_attention_params_base // The O matrix (output). void* o_ptr; - // The stride between rows of the Q, K and V matrices. - int64_t qkv_stride_in_bytes; // The stride between rows of O. int64_t o_stride_in_bytes; @@ -169,6 +172,8 @@ struct Fused_multihead_attention_params_base struct Fused_multihead_attention_params_v1 : Fused_multihead_attention_params_base { + // The stride between rows of the Q, K and V matrices. + int64_t qkv_stride_in_bytes; // The mask to implement drop-out. void* packed_mask_ptr; @@ -207,20 +212,25 @@ struct Fused_multihead_attention_params_v2 : Fused_multihead_attention_params_ba // Kv in packed qkv layout: [B, S, 3, H, D] // Contiguous kv layout: [B, 2, H, S, D]. // Paged kv layout: [UINT32_MAX, H, Tokens_per_block, D]. - fmha::cudaTmaDesc tma_desc_kv; + fmha::cudaTmaDesc tma_desc_k; + fmha::cudaTmaDesc tma_desc_v; // Tma descriptor for o fmha::cudaTmaDesc tma_desc_o; // Contiguous Q buffer pointer [B, S, H, D]. void* q_ptr; + // The separate K matrice. + void* k_ptr; + // The separate V matrice. + void* v_ptr; // Contiguous KV buffer pointer [B, 2, H, S, D]. void* kv_ptr; // Paged KV Cache buffer. fmha::Kv_block_array paged_kv_cache; // Q and KV stride (used by LDGSTS). int64_t q_stride_in_bytes; - int64_t kv_stride_in_bytes; - int64_t v_stride_in_bytes = 0; + int64_t k_stride_in_bytes; + int64_t v_stride_in_bytes; // Paged KV load. int blocks_per_tma_load; diff --git a/cpp/kernels/fmha_v2/src/fused_multihead_attention_demo_bert_params.h b/cpp/kernels/fmha_v2/src/fused_multihead_attention_demo_bert_params.h index ce8522b52f90..76670971e578 100644 --- a/cpp/kernels/fmha_v2/src/fused_multihead_attention_demo_bert_params.h +++ b/cpp/kernels/fmha_v2/src/fused_multihead_attention_demo_bert_params.h @@ -73,11 +73,15 @@ struct Fused_multihead_attention_params_v1 struct Fused_multihead_attention_params_v2 { - // The QKV matrices. + // The packed QKV matrices. void* qkv_ptr; // The separate Q matrice. void* q_ptr; - // The separate KV matrice. + // The separate K matrice. + void* k_ptr; + // The separate V matrice. + void* v_ptr; + // The separate KV matrice (contiguous KV). void* kv_ptr; // The separate paged kv cache. fmha::Kv_block_array paged_kv_cache; @@ -88,14 +92,12 @@ struct Fused_multihead_attention_params_v2 // The Softmax stats vector of layout [2, B, S, H], including softmax_sum and softmax_max void* softmax_stats_ptr; - // The stride between rows of the Q, K and V matrices. - int64_t qkv_stride_in_bytes; - // The stride between rows of the separate Q matrice. + // The stride between rows of Q. int64_t q_stride_in_bytes; - // The stride between rows of the separate KV matrice. - int64_t kv_stride_in_bytes; - // The stride between rows of the separate V matrice, set if it is not same as that of K. - int64_t v_stride_in_bytes = 0; + // The stride between rows of K. + int64_t k_stride_in_bytes; + // The stride between rows of V. + int64_t v_stride_in_bytes; // The stride between matrices of packed mask. int64_t packed_mask_stride_in_bytes; // The stride between rows of O. @@ -110,7 +112,8 @@ struct Fused_multihead_attention_params_v2 // Kv in packed qkv layout: [B, S, 3, H, D] // Contiguous kv layout: [B, 2, H, S, D]. // Paged kv layout: [UINT32_MAX, H, Tokens_per_block, D]. - fmha::cudaTmaDesc tma_desc_kv; + fmha::cudaTmaDesc tma_desc_k; + fmha::cudaTmaDesc tma_desc_v; // Tma descriptor for o fmha::cudaTmaDesc tma_desc_o; diff --git a/cpp/kernels/fmha_v2/src/fused_multihead_attention_utils.h b/cpp/kernels/fmha_v2/src/fused_multihead_attention_utils.h index ff517df9d75a..245adc65a8a3 100644 --- a/cpp/kernels/fmha_v2/src/fused_multihead_attention_utils.h +++ b/cpp/kernels/fmha_v2/src/fused_multihead_attention_utils.h @@ -441,6 +441,8 @@ static inline void extract_and_transpose_output(void* dst_, void* src_, std::vec //////////////////////////////////////////////////////////////////////////////////////////////////// static inline void store_q_and_contiguous_kv_cache(void* q_d, // [B, S, H, D] + void* k_d, // [B, S, H_kv, D] + void* v_d, // [B, S, H_kv, Dv] void* contiguous_kv_h, // [B, S, 2, H, D] void* contiguous_kv_d, // [B, S, 2, H, D] float const* qkv_packed_src, // [B, S, H, 3, D] @@ -485,19 +487,21 @@ static inline void store_q_and_contiguous_kv_cache(void* q_d, // [B, S, H, D] } } FMHA_CHECK_CUDA(cudaMemcpy(q_d, q_tmp, q_sz, cudaMemcpyDefault)); + free(q_tmp); - // DeepSeek MLA only use paged kv for now, will enable it in the future - if (d != dv) - { - return; - } // Handle contiguous KV [B, S, 2, H, D]. // Group head size. int h_q_per_kv = h_q / h_kv; // The total number of kv tokens. size_t const total_kv_tokens = cu_kv_seqlens[b]; // The kv cache size in bytes. - size_t const kv_size_in_bytes = get_size_in_bytes(total_kv_tokens * 2 * h_kv * d, dtype); + size_t const kv_size_in_bytes = get_size_in_bytes(total_kv_tokens * h_kv * (d + dv), dtype); + // Handle Separate K and V. + size_t k_size_in_bytes = get_size_in_bytes(total_kv_tokens * h_kv * d, dtype); + void* k_h = (void*) malloc(k_size_in_bytes); + size_t v_size_in_bytes = get_size_in_bytes(total_kv_tokens * h_kv * dv, dtype); + void* v_h = (void*) malloc(v_size_in_bytes); + // Batch size. for (size_t bi = 0; bi < b; bi++) { @@ -506,37 +510,61 @@ static inline void store_q_and_contiguous_kv_cache(void* q_d, // [B, S, H, D] // The actual kv sequence length. int const actual_kv_seqlen = cu_kv_seqlens[bi + 1] - cu_kv_seqlens[bi]; // [B, S, H, 3, D] - float const* kv_packed_src = qkv_packed_src + seqlen_offset * h_q * 3 * d; + float const* kv_packed_src = qkv_packed_src + seqlen_offset * h_q * (2 * d + dv); // Head. for (size_t hi = 0; hi < h_kv; hi++) { // Sequence. for (size_t si = 0; si < actual_kv_seqlen; si++) { - // Head size. + // K + size_t dst_k_offset_1 = (seqlen_offset + si) * h_kv * (d + dv) + hi * d; + size_t dst_k_offset_2 = (seqlen_offset + si) * h_kv * d + hi * d; + size_t src_k_offset = (si * h_q + hi * h_q_per_kv) * (2 * d + dv) + d; for (size_t di = 0; di < d; di++) { - size_t dst_k_offset = (seqlen_offset + si) * 2 * h_kv * d + hi * d + di; - size_t dst_v_offset = dst_k_offset + h_kv * d; - size_t src_k_offset = si * h_q * 3 * d + hi * h_q_per_kv * 3 * d + di + d; - size_t src_v_offset = src_k_offset + d; switch (dtype) { case DATA_TYPE_FP16: - reinterpret_cast(contiguous_kv_h)[dst_k_offset] = half(kv_packed_src[src_k_offset]); - reinterpret_cast(contiguous_kv_h)[dst_v_offset] = half(kv_packed_src[src_v_offset]); + reinterpret_cast(contiguous_kv_h)[dst_k_offset_1 + di] + = reinterpret_cast(k_h)[dst_k_offset_2 + di] + = half(kv_packed_src[src_k_offset + di]); + break; + case DATA_TYPE_BF16: + reinterpret_cast<__nv_bfloat16*>(contiguous_kv_h)[dst_k_offset_1 + di] + = reinterpret_cast<__nv_bfloat16*>(k_h)[dst_k_offset_2 + di] + = __float2bfloat16(kv_packed_src[src_k_offset + di]); + break; + case DATA_TYPE_E4M3: + reinterpret_cast<__nv_fp8_e4m3*>(contiguous_kv_h)[dst_k_offset_1 + di] + = reinterpret_cast<__nv_fp8_e4m3*>(k_h)[dst_k_offset_2 + di] + = __nv_fp8_e4m3(kv_packed_src[src_k_offset + di]); + break; + default: assert(false); + } + } + // V + size_t dst_v_offset_1 = (seqlen_offset + si) * h_kv * (d + dv) + h_kv * d + hi * dv; + size_t dst_v_offset_2 = (seqlen_offset + si) * h_kv * dv + hi * dv; + size_t src_v_offset = src_k_offset + d; + for (size_t di = 0; di < dv; di++) + { + switch (dtype) + { + case DATA_TYPE_FP16: + reinterpret_cast(contiguous_kv_h)[dst_v_offset_1 + di] + = reinterpret_cast(v_h)[dst_v_offset_2 + di] + = half(kv_packed_src[src_v_offset + di]); break; case DATA_TYPE_BF16: - reinterpret_cast<__nv_bfloat16*>(contiguous_kv_h)[dst_k_offset] - = __float2bfloat16(kv_packed_src[src_k_offset]); - reinterpret_cast<__nv_bfloat16*>(contiguous_kv_h)[dst_v_offset] - = __float2bfloat16(kv_packed_src[src_v_offset]); + reinterpret_cast<__nv_bfloat16*>(contiguous_kv_h)[dst_v_offset_1 + di] + = reinterpret_cast<__nv_bfloat16*>(v_h)[dst_v_offset_2 + di] + = __float2bfloat16(kv_packed_src[src_v_offset + di]); break; case DATA_TYPE_E4M3: - reinterpret_cast<__nv_fp8_e4m3*>(contiguous_kv_h)[dst_k_offset] - = __nv_fp8_e4m3(kv_packed_src[src_k_offset]); - reinterpret_cast<__nv_fp8_e4m3*>(contiguous_kv_h)[dst_v_offset] - = __nv_fp8_e4m3(kv_packed_src[src_v_offset]); + reinterpret_cast<__nv_fp8_e4m3*>(contiguous_kv_h)[dst_v_offset_1 + di] + = reinterpret_cast<__nv_fp8_e4m3*>(v_h)[dst_v_offset_2 + di] + = __nv_fp8_e4m3(kv_packed_src[src_v_offset + di]); break; default: assert(false); } @@ -546,6 +574,10 @@ static inline void store_q_and_contiguous_kv_cache(void* q_d, // [B, S, H, D] } FMHA_CHECK_CUDA(cudaMemcpy(contiguous_kv_d, contiguous_kv_h, kv_size_in_bytes, cudaMemcpyDefault)); + FMHA_CHECK_CUDA(cudaMemcpy(k_d, k_h, k_size_in_bytes, cudaMemcpyDefault)); + FMHA_CHECK_CUDA(cudaMemcpy(v_d, v_h, v_size_in_bytes, cudaMemcpyDefault)); + free(k_h); + free(v_h); } //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/cpp/kernels/xqa/gen_cubins.py b/cpp/kernels/xqa/gen_cubins.py index a8f0ecece3fd..b520779dab3b 100755 --- a/cpp/kernels/xqa/gen_cubins.py +++ b/cpp/kernels/xqa/gen_cubins.py @@ -117,12 +117,12 @@ }; """ -is_medusa = False +is_spec_dec = False def generate_cubin_meta_info_line(arch: int, compile_macros: List[CompileMacro], function_name: str, cubin_size: int, - is_last: bool, is_medusa: bool): + is_last: bool, is_spec_dec: bool): data_type_str = None kv_data_type_str = None head_dim = None @@ -160,7 +160,7 @@ def generate_cubin_meta_info_line(arch: int, compile_macros: List[CompileMacro], assert (tokens_per_page % 2 == 0) paged_kv_cache = 'true' if tokens_per_page > 0 else 'false' - use_medusa = 'true' if is_medusa else 'false' + use_medusa = 'true' if is_spec_dec else 'false' assert data_type_str is not None assert kv_data_type_str is not None assert head_dim is not None @@ -376,7 +376,7 @@ def generate_compile_arch_macro_list(compile_macro_options: list): option_macro_names, option_short_names, option_combination) ] if arch in (90, ) and option_combination[ - 3] == 2 and option_combination[2] == 1 and not is_medusa: + 3] == 2 and option_combination[2] == 1 and not is_spec_dec: input_file_name = "mha_sm90.cu" else: input_file_name = "mha.cu" @@ -387,7 +387,7 @@ def generate_compile_arch_macro_list(compile_macro_options: list): def generate_header_file_contents( all_arch_macros: List[CompileArchMacrosAndFile], - name_size_list: List[Tuple[str, int]], is_medusa: bool): + name_size_list: List[Tuple[str, int]], is_spec_dec: bool): cubin_data_array = [] cubin_length_array = [] meta_line_array = [] @@ -406,7 +406,7 @@ def generate_header_file_contents( generate_cubin_meta_info_line(arch, macros, function_name, cubin_size, i == len(all_arch_macros) - 1, - is_medusa)) + is_spec_dec)) cubin_data = ''.join(cubin_data_array) cubin_length = ''.join(cubin_length_array) meta_struct = ''.join([ @@ -422,8 +422,8 @@ def generate_header_file_contents( shutil.rmtree(cubin_dir) os.mkdir(cubin_dir) - if len(sys.argv) > 1 and sys.argv[1] == 'medusa': - is_medusa = True + if len(sys.argv) > 1 and sys.argv[1] == 'spec_dec': + is_spec_dec = True nvcc_flags = '-std=c++17 -O3 -cubin -DGENERATE_CUBIN=1 -DNDEBUG -DSPEC_DEC --use_fast_math -Xptxas=-v --allow-unsupported-compiler --expt-relaxed-constexpr -t 0' arch_options = [80, 86, 89, 90] config_list = [[ @@ -444,7 +444,7 @@ def generate_header_file_contents( with multiprocessing.Pool(processes=thread_count) as pool: name_size_list = pool.map(run_cubin_gen, arch_macro_lists) header_file_contents = generate_header_file_contents( - arch_macro_lists, name_size_list, is_medusa) + arch_macro_lists, name_size_list, is_spec_dec) with open(cubin_dir + build_func_name_prefix + '_cubin.h', "w") as f: f.write("".join( diff --git a/cpp/kernels/xqa/mha.cu b/cpp/kernels/xqa/mha.cu index 77a3ca12ee87..1acdb9852f80 100644 --- a/cpp/kernels/xqa/mha.cu +++ b/cpp/kernels/xqa/mha.cu @@ -1609,7 +1609,7 @@ CUBIN_EXPORT __global__ uint32_t const nbSeqIters = useKVCache ? divUp(cacheSeqLen, ctaTile.x) : 0; #if SPEC_DEC - uint32_t const nbSeqItersWithoutMask = (cacheSeqLen - qSeqLen) / ctaTile.x; + uint32_t const nbSeqItersWithoutMask = (cacheSeqLen - actualQSeqLen) / ctaTile.x; #endif uint32_t const seqStrideIters = nbSubSeqPerSeq; diff --git a/cpp/micro_benchmarks/README.md b/cpp/micro_benchmarks/README.md index 39fc5e102c4c..a1504a2dee9a 100644 --- a/cpp/micro_benchmarks/README.md +++ b/cpp/micro_benchmarks/README.md @@ -11,6 +11,9 @@ To build add the `--micro_benchmark` flag to `build_wheel.py` or pass `-DBUILD_M ### Mixture Of Experts Backend Benchmark +> [!CAUTION] +> Disclaimer this benchmark is intended for developers to help evaluating the impact of new optimisations. This benchmark does not meet the same quality standards as other parts of TRT-LLM. Please use with caution + Target `mixtureOfExpertsBackendBenchmark` This benchmark covers the backend used by the `MixtureOfExperts` plugin. It allows you to benchmark different MOE diff --git a/cpp/micro_benchmarks/gen-moe-benchmark-file.py b/cpp/micro_benchmarks/gen-moe-benchmark-file.py index 571edd976da4..c8f72b4ef658 100644 --- a/cpp/micro_benchmarks/gen-moe-benchmark-file.py +++ b/cpp/micro_benchmarks/gen-moe-benchmark-file.py @@ -14,7 +14,8 @@ {dtype_string} {routing_string} {tactic_string} - "bias": 0 + "bias": 0, + "gemm_to_profile": {gemm_to_profile} }}''' @@ -54,39 +55,50 @@ def populate_benchmark_config(**kwargs): # Default Mixtral configurations -num_experts = 256 -k = 8 +num_experts = 8 +k = 2 hidden_size = 4096 -inter_size = 2048 -tp_size = 8 -ep_size = 1 +inter_size = 14336 +# tp_size = 8 +# ep_size = 1 world_rank = 0 act_fn = 3 -dtype_string = make_dtype_string(["fp4", "wfp4afp8"]) # All dtypes -routing_string = make_routing_string( - name="uniform", - is_distribution=True) # Use the default uniform random distribution +dtype_string = make_dtype_string() # All dtypes tactic_id1 = '"auto"' tactic_id2 = '"auto"' +gemms_to_profile = [1, 2, 3] configs = [] -for num_tokens in [1, 8, 64, 2048, 65536]: - configs.append( - populate_benchmark_config( - num_experts=num_experts, - k=k, - hidden_size=hidden_size, - inter_size=inter_size, - tp_size=tp_size, - ep_size=ep_size, - world_rank=world_rank, - num_tokens=num_tokens, - act_fn=act_fn, - dtype_string=dtype_string, - routing_string=routing_string, - tactic_string=make_tactic_string(tactic_id1=tactic_id1, - tactic_id2=tactic_id2), - )) +for ep_size in [1, num_experts]: + for num_tokens in [1, 8, 64, 2048, 16384]: + tp_size = 8 // ep_size + if inter_size % (tp_size * 128) != 0: + continue # Insufficient alignment + if num_tokens <= num_experts: + routing_string = make_routing_string( + name="balanced", + is_distribution=False) # Use the balanced distribution + else: + routing_string = make_routing_string( + name="uniform", is_distribution=True + ) # Use the default uniform random distribution + for gemm_to_profile in gemms_to_profile: + configs.append( + populate_benchmark_config(num_experts=num_experts, + k=k, + hidden_size=hidden_size, + inter_size=inter_size, + tp_size=tp_size, + ep_size=ep_size, + world_rank=world_rank, + num_tokens=num_tokens, + act_fn=act_fn, + dtype_string=dtype_string, + routing_string=routing_string, + tactic_string=make_tactic_string( + tactic_id1=tactic_id1, + tactic_id2=tactic_id2), + gemm_to_profile=gemm_to_profile)) full_string = "[\n" + ",\n".join(configs) + "\n]" diff --git a/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h b/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h index 0790b842d450..565c170e1dfe 100644 --- a/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h +++ b/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h @@ -71,6 +71,13 @@ enum VERBOSE_LEVEL constexpr int LOG_LEVEL = ERROR; +enum class GemmToProfile : int +{ + GEMM_1 = static_cast(GemmProfilerBackend::GemmToProfile::GEMM_1), + GEMM_2 = static_cast(GemmProfilerBackend::GemmToProfile::GEMM_2), + LAYER = static_cast(3), +}; + namespace { // Abstract class for routing config @@ -358,6 +365,10 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture constexpr static int64_t FP4_VECTOR_SIZE = NVFP4 ? TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaleVectorSize : TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaleVectorSize; + constexpr static int64_t MinNDimAlignment = NVFP4 ? TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentNVFP4 + : TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX; + constexpr static int64_t MinKDimAlignment = NVFP4 ? TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentNVFP4 + : TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentMXFPX; std::vector managed_buffers; int* mSelectedExperts{}; @@ -365,6 +376,7 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture int64_t mHiddenSize{}; int64_t mNumExperts{}; + int64_t mNumExpertsPerNode{}; int64_t mK{}; constexpr static nvinfer1::DataType toDTypeID() @@ -497,6 +509,8 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture } CutlassMoeFCRunner mMoERunner{}; + GemmProfilerBackend mGemmProfilerBackend{}; + char* mGemmProfilerWorkspace{}; char* mWorkspace{}; float* mScaleProbs{}; WeightStorage* mExpertWeight1{}; @@ -544,6 +558,7 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture std::optional mSelectedConfig = std::nullopt; int64_t mBufferIndex = 0; + size_t mGemmProfilerWorkspaceSize = 0; size_t mWorkspaceSize = 0; size_t mExpertWeight1Size = 0; size_t mExpertWeight2Size = 0; @@ -559,10 +574,15 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture size_t mExpertIntScale1Size = 0; size_t mExpertIntScale2Size = 0; + size_t padSize(size_t size) + { + return ceilDiv(size, 128) * 128; + } + template T* allocBuffer(size_t size) { - size_t size_padded = ceilDiv(size * sizeof(T), 128) * 128; + size_t size_padded = padSize(size) * sizeof(T); auto i_buffer = bufferManager->gpu(size_padded); check_cuda_error(cudaGetLastError()); managed_buffers.emplace_back(std::move(i_buffer)); @@ -572,7 +592,7 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture } void initBuffersPermute(int64_t num_tokens, int64_t hidden_size, int64_t inter_size, int64_t num_experts, int64_t k, - int64_t routing_config, MOEParallelismConfig parallelism_config) + int64_t routing_config, MOEParallelismConfig parallelism_config, GemmToProfile gemm_to_profile) { assert(hidden_size % BASE_HIDDEN_SIZE == 0); @@ -582,104 +602,160 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture mHiddenSize = hidden_size; mInterSize = inter_size / parallelism_config.tp_size; mNumExperts = num_experts; + mNumExpertsPerNode = num_experts / parallelism_config.ep_size; mK = k; mIsGated = isGatedActivation(mActType); mGatedMultiplier = mIsGated ? 2 : 1; auto const gated_inter = mInterSize * mGatedMultiplier; + size_t const expert_matrix_size = padSize(mNumExpertsPerNode * mHiddenSize * mInterSize); - mWorkspaceSize = mMoERunner.getWorkspaceSize(mTotalTokens, mHiddenSize, mInterSize, mNumExperts, mK, mActType, - {}, mUseLora, /*use_deepseek_fp8_block_scale=*/false, /*min_latency_mode=*/false, mUsePrequantScale); - - mWorkspace = allocBuffer(mWorkspaceSize * NUM_BUFFERS); - size_t const expert_matrix_size = mNumExperts * mHiddenSize * mInterSize; - - mExpertWeight1Size = expert_matrix_size * mGatedMultiplier / WEIGHT_ELEM_PER_BYTE; - mExpertWeight2Size = expert_matrix_size / WEIGHT_ELEM_PER_BYTE; - mExpertWeight1 = allocBuffer(mExpertWeight1Size * NUM_BUFFERS); - mExpertWeight2 = allocBuffer(mExpertWeight2Size * NUM_BUFFERS); + bool need_weight_1 = gemm_to_profile == GemmToProfile::GEMM_1 || gemm_to_profile == GemmToProfile::LAYER; + bool need_weight_2 = gemm_to_profile == GemmToProfile::GEMM_2 || gemm_to_profile == GemmToProfile::LAYER; + mExpertWeight1Size = need_weight_1 ? expert_matrix_size * mGatedMultiplier / WEIGHT_ELEM_PER_BYTE : 0; + mExpertWeight2Size = need_weight_2 ? expert_matrix_size / WEIGHT_ELEM_PER_BYTE : 0; + mExpertWeight1 = need_weight_1 ? allocBuffer(mExpertWeight1Size * NUM_BUFFERS) : nullptr; + mExpertWeight2 = need_weight_2 ? allocBuffer(mExpertWeight2Size * NUM_BUFFERS) : nullptr; - mExpertBias1 = nullptr; - mExpertBias2 = nullptr; - if (mUseBias) + if (gemm_to_profile == GemmToProfile::LAYER) { - mExpertBias1Size = mNumExperts * gated_inter; - mExpertBias2Size = mNumExperts * mHiddenSize; - mExpertBias1 = allocBuffer(mExpertBias1Size * NUM_BUFFERS); - mExpertBias2 = allocBuffer(mExpertBias2Size * NUM_BUFFERS); - } - if constexpr (INT_QUANT) - { - mExpertIntScale1Size = mNumExperts * gated_inter; - mExpertIntScale2Size = mNumExperts * mHiddenSize; - mExpertIntScale1 = allocBuffer(mExpertIntScale1Size * NUM_BUFFERS); - mExpertIntScale2 = allocBuffer(mExpertIntScale2Size * NUM_BUFFERS); + mWorkspaceSize = mMoERunner.getWorkspaceSize(mTotalTokens, mHiddenSize, mInterSize, mNumExperts, mK, + mActType, parallelism_config, mUseLora, /*use_deepseek_fp8_block_scale=*/false, + /*min_latency_mode=*/false, mUsePrequantScale); - for (int i = 0; i < NUM_BUFFERS; i++) + mWorkspace = allocBuffer(mWorkspaceSize * NUM_BUFFERS); + + mExpertBias1 = nullptr; + mExpertBias2 = nullptr; + if (mUseBias) { - mQuantParams[i] = QuantParams::Int( - mExpertIntScale1 + mExpertIntScale1Size * i, mExpertIntScale2 + mExpertIntScale2Size * i); + mExpertBias1Size = padSize(mNumExpertsPerNode * gated_inter); + mExpertBias2Size = padSize(mNumExpertsPerNode * mHiddenSize); + mExpertBias1 = allocBuffer(mExpertBias1Size * NUM_BUFFERS); + mExpertBias2 = allocBuffer(mExpertBias2Size * NUM_BUFFERS); } - } - else if constexpr (FP8) - { - mExpertFP8Scale1 = allocBuffer(mNumExperts); - mExpertFP8Scale2 = allocBuffer(1); - mExpertFP8Scale3 = allocBuffer(mNumExperts); - for (int i = 0; i < NUM_BUFFERS; i++) + if constexpr (INT_QUANT) { - mQuantParams[i] = QuantParams::FP8(mExpertFP8Scale1, mExpertFP8Scale2, mExpertFP8Scale3); + mExpertIntScale1Size = padSize(mNumExpertsPerNode * gated_inter); + mExpertIntScale2Size = padSize(mNumExpertsPerNode * mHiddenSize); + mExpertIntScale1 = allocBuffer(mExpertIntScale1Size * NUM_BUFFERS); + mExpertIntScale2 = allocBuffer(mExpertIntScale2Size * NUM_BUFFERS); + + for (int i = 0; i < NUM_BUFFERS; i++) + { + mQuantParams[i] = QuantParams::Int( + mExpertIntScale1 + mExpertIntScale1Size * i, mExpertIntScale2 + mExpertIntScale2Size * i); + } } - } - else if constexpr (ANY_FP4) - { - mExpertFP4ActScale1 = allocBuffer(1); - mExpertFP4WeightSf1Size = num_experts * gated_inter * mHiddenSize / FP4_VECTOR_SIZE; - mExpertFP4WeightSf1 = allocBuffer(mExpertFP4WeightSf1Size * NUM_BUFFERS); - mExpertFP4GlobalScale1 = allocBuffer(num_experts); + else if constexpr (FP8) + { + mExpertFP8Scale1 = allocBuffer(mNumExpertsPerNode); + mExpertFP8Scale2 = allocBuffer(1); + mExpertFP8Scale3 = allocBuffer(mNumExpertsPerNode); - mExpertFP4ActScale2 = allocBuffer(1); - mExpertFP4WeightSf2Size = num_experts * mInterSize * mHiddenSize / FP4_VECTOR_SIZE; - mExpertFP4WeightSf2 = allocBuffer(mExpertFP4WeightSf2Size * NUM_BUFFERS); - mExpertFP4GlobalScale2 = allocBuffer(num_experts); + for (int i = 0; i < NUM_BUFFERS; i++) + { + mQuantParams[i] = QuantParams::FP8(mExpertFP8Scale1, mExpertFP8Scale2, mExpertFP8Scale3); + } + } + else if constexpr (ANY_FP4) + { + mExpertFP4ActScale1 = allocBuffer(mNumExpertsPerNode); + mExpertFP4WeightSf1Size = mNumExpertsPerNode + * TmaWarpSpecializedGroupedGemmInput::alignToSfDim(gated_inter, MinNDimAlignment) + * TmaWarpSpecializedGroupedGemmInput::alignToSfDim(mHiddenSize, MinKDimAlignment) / FP4_VECTOR_SIZE; + mExpertFP4WeightSf1 = allocBuffer(mExpertFP4WeightSf1Size * NUM_BUFFERS); + mExpertFP4GlobalScale1 = allocBuffer(mNumExpertsPerNode); + + mExpertFP4ActScale2 = allocBuffer(mNumExpertsPerNode); + mExpertFP4WeightSf2Size = mNumExpertsPerNode + * TmaWarpSpecializedGroupedGemmInput::alignToSfDim(mInterSize, MinNDimAlignment) + * TmaWarpSpecializedGroupedGemmInput::alignToSfDim(mHiddenSize, MinKDimAlignment) / FP4_VECTOR_SIZE; + mExpertFP4WeightSf2 = allocBuffer(mExpertFP4WeightSf2Size * NUM_BUFFERS); + mExpertFP4GlobalScale2 = allocBuffer(mNumExpertsPerNode); + + auto func = NVFP4 ? QuantParams::FP4 : QuantParams::FP8MXFP4; + for (int i = 0; i < NUM_BUFFERS; i++) + { + mQuantParams[i] = func(mExpertFP4ActScale1, mExpertFP4WeightSf1 + mExpertFP4WeightSf1Size * i, + mExpertFP4GlobalScale1, mExpertFP4ActScale2, mExpertFP4WeightSf2 + mExpertFP4WeightSf2Size * i, + mExpertFP4GlobalScale2, false, false); + } + } - auto func = NVFP4 ? QuantParams::FP4 : QuantParams::FP8MXFP4; + mSelectedExpertsSize = padSize(mTotalTokens * mK); + mSelectedExperts = allocBuffer(mSelectedExpertsSize * NUM_BUFFERS); + mScaleProbsSize = padSize(mTotalTokens * mK); + mScaleProbs = allocBuffer(mScaleProbsSize * NUM_BUFFERS); + mInputTensorSize = padSize(mTotalTokens * mHiddenSize); + mInputTensor = allocBuffer(mInputTensorSize * NUM_BUFFERS); + mFinalOutputSize = padSize(mTotalTokens * mHiddenSize); + mFinalOutput = allocBuffer(mFinalOutputSize * NUM_BUFFERS); + + mSourceToExpandedMapSize = padSize(mTotalTokens * mK); + mSourceToExpandedMap = allocBuffer(mSourceToExpandedMapSize * NUM_BUFFERS); + mRoutingConfigIndex = routing_config; + auto tactic = routingConfigCache.at(routing_config); + tactic->start(); for (int i = 0; i < NUM_BUFFERS; i++) { - mQuantParams[i] = func(mExpertFP4ActScale1, mExpertFP4WeightSf1 + mExpertFP4WeightSf1Size * i, - mExpertFP4GlobalScale1, mExpertFP4ActScale2, mExpertFP4WeightSf2 + mExpertFP4WeightSf2Size * i, - mExpertFP4GlobalScale2, false, false); + tactic->setRouting(mSelectedExperts + mSelectedExpertsSize * i, mNumExperts, mK, mTotalTokens); } } - mSelectedExpertsSize = mTotalTokens * mK; - mSelectedExperts = allocBuffer(mSelectedExpertsSize * NUM_BUFFERS); - mScaleProbsSize = mTotalTokens * mK; - mScaleProbs = allocBuffer(mScaleProbsSize * NUM_BUFFERS); - mInputTensorSize = mTotalTokens * mHiddenSize; - mInputTensor = allocBuffer(mInputTensorSize * NUM_BUFFERS); - mFinalOutputSize = mTotalTokens * mHiddenSize; - mFinalOutput = allocBuffer(mFinalOutputSize * NUM_BUFFERS); - - mSourceToExpandedMapSize = mTotalTokens * mK; - mSourceToExpandedMap = allocBuffer(mSourceToExpandedMapSize * NUM_BUFFERS); - - mRoutingConfigIndex = routing_config; - auto tactic = routingConfigCache.at(routing_config); - tactic->start(); - for (int i = 0; i < NUM_BUFFERS; i++) +#ifdef USING_OSS_CUTLASS_MOE_GEMM + mGemmProfilerBackend.init(mMoERunner, GemmProfilerBackend::GemmToProfile::Undefined, typeToDtypeID(), + typeToDtypeID(), typeToDtypeID(), mNumExperts, mK, mHiddenSize, mInterSize, + mGroupSize, mActType, mUseBias, mUseLora, /*min_latency_mode=*/false, + /*need_weights=*/false, parallelism_config, /*enable_alltoall=*/false); +#else + mGemmProfilerBackend.init(mMoERunner, GemmProfilerBackend::GemmToProfile::Undefined, typeToDtypeID(), + typeToDtypeID(), typeToDtypeID(), mNumExperts, mK, mHiddenSize, mInterSize, + mGroupSize, mActType, mUseBias, mUseLora, /*min_latency_mode=*/false, + /*need_weights=*/false, parallelism_config); +#endif + + mGemmProfilerWorkspaceSize = 0; + if (gemm_to_profile == GemmToProfile::GEMM_1 || gemm_to_profile == GemmToProfile::LAYER) + { + mGemmProfilerBackend.mGemmToProfile = GemmProfilerBackend::GemmToProfile::GEMM_1; + mGemmProfilerWorkspaceSize + = std::max(mGemmProfilerWorkspaceSize, mGemmProfilerBackend.getWorkspaceSize(mTotalTokens)); + } + + if (gemm_to_profile == GemmToProfile::GEMM_2 || gemm_to_profile == GemmToProfile::LAYER) { - tactic->setRouting(mSelectedExperts + mSelectedExpertsSize * i, mNumExperts, mK, mTotalTokens); + mGemmProfilerBackend.mGemmToProfile = GemmProfilerBackend::GemmToProfile::GEMM_2; + mGemmProfilerWorkspaceSize + = std::max(mGemmProfilerWorkspaceSize, mGemmProfilerBackend.getWorkspaceSize(mTotalTokens)); } + int64_t num_gemm_buffers = gemm_to_profile == GemmToProfile::LAYER ? 1 : NUM_BUFFERS; + mGemmProfilerWorkspaceSize = padSize(mGemmProfilerWorkspaceSize); + mGemmProfilerWorkspace = mGemmProfilerWorkspaceSize > 0 + ? allocBuffer(mGemmProfilerWorkspaceSize * num_gemm_buffers) + : nullptr; + check_cuda_error(cudaStreamSynchronize(streamPtr->get())); } + void prepareGemmProfiler(GemmToProfile gemm_to_profile) + { + if (gemm_to_profile == GemmToProfile::LAYER) + return; + mGemmProfilerBackend.mGemmToProfile = static_cast(gemm_to_profile); + auto* expert_weights = gemm_to_profile == GemmToProfile::GEMM_1 ? mExpertWeight1 : mExpertWeight2; + auto expert_weights_size = gemm_to_profile == GemmToProfile::GEMM_1 ? mExpertWeight1Size : mExpertWeight2Size; + mGemmProfilerBackend.prepare(mTotalTokens, mGemmProfilerWorkspace + mGemmProfilerWorkspaceSize * mBufferIndex, + /*expert_weights=*/expert_weights + expert_weights_size * mBufferIndex, streamPtr->get()); + } + std::array mGraph{}; + std::array mGraphInstance{}; - void createGraph(MOEParallelismConfig parallelism_config) + void createGraph(MOEParallelismConfig parallelism_config, GemmToProfile gemm_to_profile) { if (!useCudaGraph) return; @@ -689,9 +765,11 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture for (int i = 0; i < NUM_BUFFERS; i++) { mBufferIndex = i; + // Each buffer will have a different routing config for the gemm profiler + prepareGemmProfiler(gemm_to_profile); check_cuda_error(cudaGraphCreate(&mGraph[i], 0)); check_cuda_error(cudaStreamBeginCapture(streamPtr->get(), cudaStreamCaptureModeThreadLocal)); - runMoEPermute(parallelism_config); + runMoEPermute(parallelism_config, gemm_to_profile); check_cuda_error(cudaStreamEndCapture(streamPtr->get(), &mGraph[i])); check_cuda_error(cudaGraphInstantiate(&mGraphInstance[i], mGraph[i], nullptr, nullptr, 0)); } @@ -711,13 +789,23 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture } } - float benchmarkLoop(MOEParallelismConfig parallelism_config) + float benchmarkLoop(MOEParallelismConfig parallelism_config, GemmToProfile gemm_to_profile) { mBufferIndex = (mBufferIndex + 1) % NUM_BUFFERS; - auto tactic = routingConfigCache.at(mRoutingConfigIndex); - if (!tactic->isDeterministic()) + + // Setup the profiler state for this iteration. CUDA Graphs will do this when it captures the graph. + if (gemm_to_profile != GemmToProfile::LAYER && !useCudaGraph) + { + prepareGemmProfiler(gemm_to_profile); + } + else if (gemm_to_profile == GemmToProfile::LAYER) { - tactic->setRouting(mSelectedExperts + mSelectedExpertsSize * mBufferIndex, mNumExperts, mK, mTotalTokens); + auto tactic = routingConfigCache.at(mRoutingConfigIndex); + if (!tactic->isDeterministic()) + { + tactic->setRouting( + mSelectedExperts + mSelectedExpertsSize * mBufferIndex, mNumExperts, mK, mTotalTokens); + } } { @@ -729,7 +817,7 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture } else { - runMoEPermute(parallelism_config); + runMoEPermute(parallelism_config, gemm_to_profile); } check_cuda_error(cudaEventRecord(mEndEvent, streamPtr->get())); check_cuda_error(cudaStreamSynchronize(streamPtr->get())); @@ -742,27 +830,19 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture // An imprecise benchmark pass for picking the best tactic. // Runs for 3 iterations or 1 second and picks the best option - int pickBestTactic(MOEParallelismConfig parallelism_config, GemmProfilerBackend::GemmToProfile gemm_to_profile) + int pickBestTactic(MOEParallelismConfig parallelism_config, GemmToProfile gemm_to_profile) { auto tactics = mMoERunner.getTactics(); ::nvtx3::scoped_range nvtx(tensorrt_llm::common::nvtx::nextColor(), "Tactic Profiling GEMM " + std::to_string(static_cast(gemm_to_profile))); + // We save space by reusing the same workspace buffer for all tactics when doing full layer profiling. So we + // need to hardcode the buffer index to 0. + auto old_buffer_index = mBufferIndex; + mBufferIndex = 0; + prepareGemmProfiler(gemm_to_profile); + mBufferIndex = old_buffer_index; - GemmProfilerBackend profiler; -#ifdef USING_OSS_CUTLASS_MOE_GEMM - profiler.init(mMoERunner, gemm_to_profile, typeToDtypeID(), typeToDtypeID(), - typeToDtypeID(), mNumExperts, mK, mHiddenSize, mInterSize, mGroupSize, mActType, mUseBias, - mUseLora, /*min_latency_mode=*/false, /*need_weights=*/true, parallelism_config, /*enable_alltoall=*/false); -#else - profiler.init(mMoERunner, gemm_to_profile, typeToDtypeID(), typeToDtypeID(), - typeToDtypeID(), mNumExperts, mK, mHiddenSize, mInterSize, mGroupSize, mActType, mUseBias, - mUseLora, /*min_latency_mode=*/false, /*need_weights=*/true, parallelism_config); -#endif - auto workspace_size = profiler.getWorkspaceSize(mTotalTokens); - auto workspace = bufferManager->gpu(workspace_size); - - profiler.prepare( - mTotalTokens, static_cast(workspace->data()), /*expert_weights=*/nullptr, streamPtr->get()); + auto* mGemmProfilerExpertWeights = gemm_to_profile == GemmToProfile::GEMM_1 ? mExpertWeight1 : mExpertWeight2; float best_time = INFINITY; int best_idx = -1; @@ -778,13 +858,13 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture { ::nvtx3::scoped_range nvtx(tensorrt_llm::common::nvtx::nextColor(), "Tactic Profiling Warm-Up"); // Warm-Up run - profiler.runProfiler(mTotalTokens, t, static_cast(workspace->data()), - /*expert_weights=*/nullptr, streamPtr->get()); + mGemmProfilerBackend.runProfiler(mTotalTokens, t, mGemmProfilerWorkspace, + /*expert_weights=*/mGemmProfilerExpertWeights, streamPtr->get()); check_cuda_error(cudaStreamSynchronize(streamPtr->get())); } // Profile all samples or for 1 sec - int const max_iters = profiler.NUM_ROUTING_SAMPLES; + int const max_iters = mGemmProfilerBackend.NUM_ROUTING_SAMPLES; float const max_time_ms = 1000.f; float time = 0.f; @@ -796,8 +876,8 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture "Tactic Profiling Iteration " + std::to_string(iter)); check_cuda_error(cudaEventRecord(mStartEvent, streamPtr->get())); - profiler.runProfiler(mTotalTokens, t, static_cast(workspace->data()), - /*expert_weights=*/nullptr, streamPtr->get()); + mGemmProfilerBackend.runProfiler(mTotalTokens, t, mGemmProfilerWorkspace, + /*expert_weights=*/mGemmProfilerExpertWeights, streamPtr->get()); check_cuda_error(cudaEventRecord(mEndEvent, streamPtr->get())); check_cuda_error(cudaStreamSynchronize(streamPtr->get())); } @@ -838,17 +918,26 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture return best_idx; } - std::pair setTactic(int tactic_idx1, int tactic_idx2, MOEParallelismConfig parallelism_config) + int mBestTacticGemm1 = -1; + int mBestTacticGemm2 = -1; + + std::pair setTactic( + int tactic_idx1, int tactic_idx2, MOEParallelismConfig parallelism_config, GemmToProfile gemm_to_profile) { auto tactics = mMoERunner.getTactics(); - for (auto& t_ptr : {&tactic_idx1, &tactic_idx2}) + std::vector, GemmToProfile>> tactics_to_profile{ + {tactic_idx1, GemmToProfile::GEMM_1}, {tactic_idx2, GemmToProfile::GEMM_2}}; + for (auto& combo : tactics_to_profile) { - auto& t = *t_ptr; + auto& t = combo.first.get(); + if (combo.second != gemm_to_profile && gemm_to_profile != GemmToProfile::LAYER) + { + t = 0; // Unneeded tactic, set to 0 + continue; + } if (t == -1) { - t = pickBestTactic(parallelism_config, - t_ptr == &tactic_idx1 ? GemmProfilerBackend::GemmToProfile::GEMM_1 - : GemmProfilerBackend::GemmToProfile::GEMM_2); + t = pickBestTactic(parallelism_config, combo.second); } if (t < 0 || t >= tactics.size()) @@ -858,38 +947,66 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture } mMoERunner.setTactic(tactics[tactic_idx1], tactics[tactic_idx2]); + mBestTacticGemm1 = tactic_idx1; + mBestTacticGemm2 = tactic_idx2; return {tactic_idx1, tactic_idx2}; } - void runMoEPermute(MOEParallelismConfig parallelism_config) + void runMoEPermute(MOEParallelismConfig parallelism_config, GemmToProfile gemm_to_profile) { - auto stream = streamPtr->get(); - MoeMinLatencyParams min_latency_params; + switch (gemm_to_profile) + { + case GemmToProfile::GEMM_1: + case GemmToProfile::GEMM_2: + { + auto tactic_idx = gemm_to_profile == GemmToProfile::GEMM_1 ? mBestTacticGemm1 : mBestTacticGemm2; + auto* expert_weights = gemm_to_profile == GemmToProfile::GEMM_1 ? mExpertWeight1 : mExpertWeight2; + auto expert_weights_size + = gemm_to_profile == GemmToProfile::GEMM_1 ? mExpertWeight1Size : mExpertWeight2Size; + + auto tactics = mMoERunner.getTactics()[tactic_idx]; + if (static_cast(gemm_to_profile) != static_cast(mGemmProfilerBackend.mGemmToProfile)) + { + throw std::runtime_error("Configuration mismatch between mGemmProfilerBackend and runMoEPermute"); + } + mGemmProfilerBackend.mSampleIndex = mBufferIndex % mGemmProfilerBackend.NUM_ROUTING_SAMPLES; + mGemmProfilerBackend.runProfiler(mTotalTokens, tactics, + mGemmProfilerWorkspace + mGemmProfilerWorkspaceSize * mBufferIndex, + /*expert_weights=*/expert_weights + expert_weights_size * mBufferIndex, streamPtr->get()); + break; + } + case GemmToProfile::LAYER: + { + auto stream = streamPtr->get(); + MoeMinLatencyParams min_latency_params; #ifdef USING_OSS_CUTLASS_MOE_GEMM - mMoERunner.runMoe(mInputTensor + mInputTensorSize * mBufferIndex, nullptr, - mSelectedExperts + mSelectedExpertsSize * mBufferIndex, - mUseFinalScale ? mScaleProbs + mScaleProbsSize * mBufferIndex : nullptr, - mExpertWeight1 + mExpertWeight1Size * mBufferIndex, mExpertBias1 + mExpertBias1Size * mBufferIndex, - mActType, mExpertWeight2 + mExpertWeight2Size * mBufferIndex, - mExpertBias2 + mExpertBias2Size * mBufferIndex, mQuantParams[mBufferIndex], mTotalTokens, mHiddenSize, - mInterSize, mNumExperts, mK, mWorkspace + mWorkspaceSize * mBufferIndex, - mFinalOutput + mFinalOutputSize * mBufferIndex, - mSourceToExpandedMap + mSourceToExpandedMapSize * mBufferIndex, parallelism_config, - /*enable_alltoall=*/false, mUseLora, mLoraParams[mBufferIndex], - /*use_fp8_block_scaling=*/false, /*min_latency_mode=*/false, min_latency_params, stream); + mMoERunner.runMoe(mInputTensor + mInputTensorSize * mBufferIndex, nullptr, + mSelectedExperts + mSelectedExpertsSize * mBufferIndex, + mUseFinalScale ? mScaleProbs + mScaleProbsSize * mBufferIndex : nullptr, + mExpertWeight1 + mExpertWeight1Size * mBufferIndex, mExpertBias1 + mExpertBias1Size * mBufferIndex, + mActType, mExpertWeight2 + mExpertWeight2Size * mBufferIndex, + mExpertBias2 + mExpertBias2Size * mBufferIndex, mQuantParams[mBufferIndex], mTotalTokens, mHiddenSize, + mInterSize, mNumExperts, mK, mWorkspace + mWorkspaceSize * mBufferIndex, + mFinalOutput + mFinalOutputSize * mBufferIndex, + mSourceToExpandedMap + mSourceToExpandedMapSize * mBufferIndex, parallelism_config, + /*enable_alltoall=*/false, mUseLora, mLoraParams[mBufferIndex], + /*use_fp8_block_scaling=*/false, /*min_latency_mode=*/false, min_latency_params, stream); #else - mMoERunner.runMoe(mInputTensor + mInputTensorSize * mBufferIndex, nullptr, - mSelectedExperts + mSelectedExpertsSize * mBufferIndex, - mUseFinalScale ? mScaleProbs + mScaleProbsSize * mBufferIndex : nullptr, - mExpertWeight1 + mExpertWeight1Size * mBufferIndex, mExpertBias1 + mExpertBias1Size * mBufferIndex, - mActType, mExpertWeight2 + mExpertWeight2Size * mBufferIndex, - mExpertBias2 + mExpertBias2Size * mBufferIndex, mQuantParams[mBufferIndex], mTotalTokens, mHiddenSize, - mInterSize, mNumExperts, mK, mWorkspace + mWorkspaceSize * mBufferIndex, - mFinalOutput + mFinalOutputSize * mBufferIndex, - mSourceToExpandedMap + mSourceToExpandedMapSize * mBufferIndex, parallelism_config, mUseLora, - mLoraParams[mBufferIndex], - /*use_fp8_block_scaling=*/false, /*min_latency_mode=*/false, min_latency_params, stream); + mMoERunner.runMoe(mInputTensor + mInputTensorSize * mBufferIndex, nullptr, + mSelectedExperts + mSelectedExpertsSize * mBufferIndex, + mUseFinalScale ? mScaleProbs + mScaleProbsSize * mBufferIndex : nullptr, + mExpertWeight1 + mExpertWeight1Size * mBufferIndex, mExpertBias1 + mExpertBias1Size * mBufferIndex, + mActType, mExpertWeight2 + mExpertWeight2Size * mBufferIndex, + mExpertBias2 + mExpertBias2Size * mBufferIndex, mQuantParams[mBufferIndex], mTotalTokens, mHiddenSize, + mInterSize, mNumExperts, mK, mWorkspace + mWorkspaceSize * mBufferIndex, + mFinalOutput + mFinalOutputSize * mBufferIndex, + mSourceToExpandedMap + mSourceToExpandedMapSize * mBufferIndex, parallelism_config, mUseLora, + mLoraParams[mBufferIndex], + /*use_fp8_block_scaling=*/false, /*min_latency_mode=*/false, min_latency_params, stream); #endif + break; + } + } } void runBenchmark(benchmark::State& state); @@ -913,6 +1030,7 @@ void MixtureOfExpertsBenchmark::runBenchmark(benchmark::State& state int tactic_idx1 = state.range(11); int tactic_idx2 = state.range(12); int const routing_config = state.range(13); + GemmToProfile const gemm_to_profile = static_cast(state.range(14)); state.counters["num_experts"] = num_experts; state.counters["top_k"] = top_k; @@ -928,11 +1046,12 @@ void MixtureOfExpertsBenchmark::runBenchmark(benchmark::State& state state.counters["routing_config"] = (int) routing_config; state.counters["dtype"] = (int) toDTypeID(); state.counters["wtype"] = (int) toWTypeID(); + state.counters["gemm_to_profile"] = (int) gemm_to_profile; std::stringstream ss; - ss << "Experts,K,Hidden,Inter,TP,EP,Rank,Tokens,Bias,Scale,Actfn,Tactic,Routing="; + ss << "Experts,K,Hidden,Inter,TP,EP,Rank,Tokens,Bias,Scale,Actfn,Tactic1,Tactic2,Gemm,Routing="; for (auto v : {num_experts, top_k, hidden_size, inter_size, tp_size, ep_size, world_rank, num_tokens, - (int) mUseBias, (int) mUseFinalScale, (int) mActType, tactic_idx1, tactic_idx2}) + (int) mUseBias, (int) mUseFinalScale, (int) mActType, tactic_idx1, tactic_idx2, (int) gemm_to_profile}) { ss << v << ","; } @@ -942,10 +1061,11 @@ void MixtureOfExpertsBenchmark::runBenchmark(benchmark::State& state // Always use EP size for moe config until we support TP+EP, we just divide the inter size for TP MOEParallelismConfig parallelism_config{tp_size, world_rank / ep_size, ep_size, world_rank % ep_size}; - initBuffersPermute(num_tokens, hidden_size, inter_size, num_experts, top_k, routing_config, parallelism_config); + initBuffersPermute( + num_tokens, hidden_size, inter_size, num_experts, top_k, routing_config, parallelism_config, gemm_to_profile); // Parse the tactic, does checks for "auto" mode and out of range - std::tie(tactic_idx1, tactic_idx2) = setTactic(tactic_idx1, tactic_idx2, parallelism_config); + std::tie(tactic_idx1, tactic_idx2) = setTactic(tactic_idx1, tactic_idx2, parallelism_config, gemm_to_profile); if (tactic_idx1 < 0 || tactic_idx2 < 0) { state.SkipWithMessage("Out of range tactic"); @@ -962,13 +1082,13 @@ void MixtureOfExpertsBenchmark::runBenchmark(benchmark::State& state state.counters["tactic_idx1"] = tactic_idx1; state.counters["tactic_idx2"] = tactic_idx2; - createGraph(parallelism_config); + createGraph(parallelism_config, gemm_to_profile); { - NVTX3_SCOPED_RANGE(BenchmarkRun); + ::nvtx3::scoped_range nvtx(tensorrt_llm::common::nvtx::nextColor(), "BenchmarkRun " + ss.str()); for (auto _ : state) { - float ms = benchmarkLoop(parallelism_config); + float ms = benchmarkLoop(parallelism_config, gemm_to_profile); state.SetIterationTime(ms / 1000.f); } } diff --git a/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkLauncher.cu b/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkLauncher.cu index 663759e3ff77..b784c6d0bc49 100644 --- a/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkLauncher.cu +++ b/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkLauncher.cu @@ -389,11 +389,11 @@ void argGenLoadFile(benchmark::internal::Benchmark* benchmark) { continue; } - else if (std::is_same_v && !hasDtype("float") - && !hasDtype("float32")) - { - continue; - } + // else if (std::is_same_v && !hasDtype("float") + // && !hasDtype("float32")) + // { + // continue; + // } else if (std::is_same_v && !hasDtype("float16") && !hasDtype("half")) { continue; @@ -452,8 +452,38 @@ void argGenLoadFile(benchmark::internal::Benchmark* benchmark) int world_rank = get_or("world_rank", 0); int bias = get_or("bias", 0); int do_final_scale = get_or("do_final_scale", 1); // Default to scales on + int gemm_to_profile = get_or("gemm_to_profile", (int) GemmToProfile::LAYER); TLLM_CHECK_WITH_INFO(world_rank < tp_size * ep_size, "Rank is out of bounds of tp*ep"); + if (gemm_to_profile != (int) GemmToProfile::LAYER && routing_config != UNIFORM_ROUTING_CONFIG) + { + static bool info_printed = false; + if (!info_printed && LOG_LEVEL >= INFO) + { + std::cerr << "Warning: GEMM profiling is experimental, results may be inaccurate" << std::endl; + info_printed = true; + } + + static bool printed = false; + if (LOG_LEVEL >= ERROR && !printed) + { + std::cerr << "Warning: Profiling a specific GEMM will always use uniform random token distribution" + << std::endl; + printed = true; + } + routing_config = UNIFORM_ROUTING_CONFIG; + if (gemm_to_profile == (int) GemmToProfile::GEMM_1) + { + tactic_ids2 = {-1}; + } + else if (gemm_to_profile == (int) GemmToProfile::GEMM_2) + { + if (!has_tactic_ids2) + tactic_ids2 = std::move(tactic_ids1); + tactic_ids1 = {-1}; + } + } + auto get_range = [&](std::string name, int min = 1, int max = INT32_MAX) { auto val = run_config.at(name).get(); @@ -482,7 +512,7 @@ void argGenLoadFile(benchmark::internal::Benchmark* benchmark) get_range("act_fn", 0, (int) ActivationType::Identity), // t1, // t2, // - *routing_config}); + *routing_config, gemm_to_profile}); } } } @@ -518,7 +548,8 @@ void argGenHardcoded(benchmark::internal::Benchmark* benchmark) for (auto tactic2 : cutlass_tactic) for (auto routing : routing_config) benchmark->Args({num_expert, k, size, inter_size, 1, 1, 0, tokens, bias, - 1, (int) act, tactic1, tactic2, routing}); + 1, (int) act, tactic1, tactic2, routing, + (int) GemmToProfile::LAYER}); } } @@ -542,7 +573,7 @@ void argGen(benchmark::internal::Benchmark* benchmark) benchmark->UseManualTime(); benchmark->ArgNames( {"Num Experts", "K", "Hidden Size", "Inter Size", "TP Size", "EP Size", "World Rank", "Num Tokens", "Use Bias", - "Use Final Scale", "Activation Function", "Tactic ID 1", "Tactic ID 2", "Routing ID"}); + "Use Final Scale", "Activation Function", "Tactic ID 1", "Tactic ID 2", "Routing ID", "Gemm To Profile"}); if (workloadFile) argGenLoadFile(benchmark); @@ -550,7 +581,8 @@ void argGen(benchmark::internal::Benchmark* benchmark) argGenHardcoded(benchmark); } -BENCHMARK_BASIC(float, float, float) +// No one cares about float32 +// BENCHMARK_BASIC(float, float, float) BENCHMARK_BASIC(half, half, half) using uint8 = uint8_t; BENCHMARK_BASIC(half, uint8, half) @@ -576,7 +608,7 @@ void delayedRegisterBenchmark() if (workloadFile) { // Extra ones we don't want for hardcoded runs - BENCHMARK_BASIC_DO_REGISTER(float, float, float); + // BENCHMARK_BASIC_DO_REGISTER(float, float, float); BENCHMARK_BASIC_DO_REGISTER(half, uint8, half); BENCHMARK_BASIC_DO_REGISTER(half, uint4b_t, half); #ifdef ENABLE_BF16 @@ -597,6 +629,9 @@ void doCleanup() void help() { + std::cout << "**Disclaimer: This benchmark is intended for developers to help evaluating the impact of new " + "optimisations. This benchmark does not meet the same quality standards as other parts of TRT-LLM. " + "Please use with caution**\n\n"; std::cout << "Usage: mixtureOfExpertsBackendBenchmark [--disable_cuda_graphs] [--input_file ] [benchmark " "options]\n"; std::cout @@ -624,6 +659,7 @@ void help() " \"routing_name\": string, (optional)\n" " \"selected_experts\": [int, ...], or string, (optional, length is a multiple of k)\n" " \"expert_distribution\": [float, ...], or string, (optional, length is num_experts)\n" + " \"gemm_to_profile\": int, (experimental, optional, 1 = gemm1, 2 = gemm2, 3 = layer)\n" " },\n" " ...\n" "]\n" @@ -664,7 +700,7 @@ void help() "Useful for quick perf tests, prefer a full sweep and manually setting the tactic for more accurate " "results" "- dtypes - A list of dtypes to run this config through.\n" - "Allowed values are: fp8, fp4, wfp4afp8, int4, int8, float, half, bfloat16\n" + "Allowed values are: fp8, fp4, wfp4afp8, int4, int8, half, bfloat16\n" "If this argument is omitted all dtypes will be run. Note, not all tactics are supported for all " "dtypes,\n" "unsupported tactics will be skipped with a warning.\n" @@ -681,6 +717,8 @@ void help() "- \"expert_distribution\" - instead of explicitly setting selected_experts, define a random distribution " "that experts will be randomly sampled from." "There is also pre-defined config \"uniform\", which is short-hand for a random uniform distribution\n" + "- \"gemm_to_profile\" - the gemm to profile, 1 = gemm1, 2 = gemm2, 3 = full layer. (default layer). If a " + "specific GEMM is profiled, it will always use uniform random token distribution\n" "\n"; std::cout << "benchmark options:\n"; diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index fa6c3cfc4ed5..848360b23da3 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -34,6 +34,7 @@ #include #include #include +#include namespace tensorrt_llm::batch_manager::kv_cache_manager { @@ -116,41 +117,45 @@ void checkAlternateWindow(BaseKVCacheManager* cacheManager, BaseCacheFormatter:: } } -std::vector CacheFormatter::pickRecvConnections( - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig) const +std::vector CacheFormatter::pickRecvConnections( + size_t numConnections, CacheState const& selfConfig, SizeType32 selfIdx, CacheState const& destConfig) const { auto targetInfo = executor::kv_cache::targetIRanks(destConfig, selfConfig, selfIdx); if (targetInfo.mPeerDupHeadFactor <= 1) { - return connections; + std::vector ret(numConnections); + std::iota(ret.begin(), ret.end(), 0); + return ret; } - TLLM_CHECK(connections.size() == targetInfo.mIRanks.size()); + TLLM_CHECK(numConnections == targetInfo.mIRanks.size()); - std::vector ret; + std::vector ret; for (int i = 0; i < targetInfo.mDomainTPSize; i++) { if (i % targetInfo.mPeerDupHeadFactor == 0) { for (int j = 0; j < targetInfo.mDomainPPSize; j++) { - ret.push_back(connections.at((i * targetInfo.mDomainPPSize) + j)); + ret.push_back((i * targetInfo.mDomainPPSize) + j); } } } return ret; } -void CacheFormatter::formatOutput(LlmRequest const& llmRequest, - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig, runtime::BufferManager const& bufferManager) +void CacheFormatter::format(TransferSession& session) { - NVTX3_SCOPED_RANGE(formatOutput); + NVTX3_SCOPED_RANGE(CacheFormatter_format); + auto const& llmRequest = session.getLlmRequest(); TLLM_LOG_DEBUG( mpi::MpiComm::world().getRank(), "Start sending KV cache for request ID: %ld.", llmRequest.mRequestId); TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently, only beam width 1 is supported."); - TLLM_CHECK(!connections.empty()); + auto const& connections = session.getConnections(); + auto const& selfConfig = session.getSelfState().getCacheState().value(); + auto const& destConfig = session.getOtherState().getCacheState().value(); + auto const selfIdx = session.getSelfState().getCommState().value().getSelfIdx(); + auto& bufferManager = session.getBufferManager(); if (!needSendCache(selfConfig, destConfig, selfIdx)) { return; @@ -186,10 +191,10 @@ void CacheFormatter::formatOutput(LlmRequest const& llmRequest, TLLM_LOG_DEBUG("Block %p of pool %d shape = %s", it->data(), poolIdx, runtime::ITensor::toString(it->getShape()).c_str()); } - for (auto const& connection : connections) + for (size_t i = 0; i < connections.size(); i++) { TLLM_LOG_DEBUG("Send layer %d(%d-%d)", layerIdx, poolIdx, layerIdxInPool); - TransferHelper::sendBuffer(*connection, *layer, llmRequest.mRequestId); + session.send(i, layer->data(), layer->getSizeInBytes()); } } } @@ -247,13 +252,13 @@ void CacheFormatter::formatOutput(LlmRequest const& llmRequest, TLLM_CHECK(connections.size() == 1); TLLM_CUDA_CHECK(cudaSetDevice(deviceId)); - for (auto const& connection : connections) + for (size_t i = 0; i < connections.size(); i++) { for (auto const& [window, blocks] : inputKvCacheBlocks) { for (auto const& block : blocks) { - TransferHelper::sendBuffer(*connection, *block, llmRequest.mRequestId); + session.send(i, block->data(), block->getSizeInBytes()); } } } @@ -308,9 +313,8 @@ void CacheFormatter::formatOutput(LlmRequest const& llmRequest, if (bufferIdx < bufferCoverTargetNum) { - size = (*outputSplitCaches[bufferIdx]).getSizeInBytes(); - TransferHelper::sendBuffer( - *connections[processIdx], *outputSplitCaches[bufferIdx], llmRequest.mRequestId); + size = outputSplitCaches[bufferIdx]->getSizeInBytes(); + session.send(processIdx, outputSplitCaches[bufferIdx]->data(), size); } else if (bufferCoverTargetNum > 0) { @@ -318,9 +322,8 @@ void CacheFormatter::formatOutput(LlmRequest const& llmRequest, auto sendBufferIdx = bufferIdx % bufferCoverTargetNum; bufferManager.copy(*outputSplitCaches[processIdx], *outputSplitCaches.at(sendBufferIdx)); bufferManager.getStream().synchronize(); - size = (*outputSplitCaches.at(sendBufferIdx)).getSizeInBytes(); - TransferHelper::sendBuffer( - *connections[processIdx], *outputSplitCaches.at(sendBufferIdx), llmRequest.mRequestId); + size = outputSplitCaches.at(sendBufferIdx)->getSizeInBytes(); + session.send(processIdx, outputSplitCaches.at(sendBufferIdx)->data(), size); } else { @@ -341,7 +344,7 @@ void CacheFormatter::formatOutput(LlmRequest const& llmRequest, auto copyTargetSlice = runtime::ITensor::slice(preAllocSendBuffer, 0, sendSize); bufferManager.copy(*copySlice, *copyTargetSlice); bufferManager.getStream().synchronize(); - TransferHelper::sendBuffer(*connections[processIdx], *copyTargetSlice, llmRequest.mRequestId); + session.send(processIdx, copyTargetSlice->data(), sendSize); remainSendSize -= sendSize; } } @@ -401,18 +404,21 @@ void CacheFormatter::formatOutput(LlmRequest const& llmRequest, mpi::MpiComm::world().getRank(), "End the sending of KV cache for the request ID:%ld ", llmRequest.mRequestId); } -void CacheFormatter::formatInput(LlmRequest const& llmRequest, - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig, runtime::BufferManager const& bufferManager) +void CacheFormatter::unformat(TransferSession& session) { - NVTX3_SCOPED_RANGE(formatInput); + NVTX3_SCOPED_RANGE(CacheFormatter_unformat); + auto const& llmRequest = session.getLlmRequest(); TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "Start receiving KV cache for request ID: %ld, context request ID: %ld.", llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); - TLLM_CHECK(!connections.empty()); + auto const& connections = session.getConnections(); + auto const& selfConfig = session.getSelfState().getCacheState().value(); + auto const& destConfig = session.getOtherState().getCacheState().value(); + auto const selfIdx = session.getSelfState().getCommState().value().getSelfIdx(); + auto& bufferManager = session.getBufferManager(); auto blockRange = getBlockRangeForReceiving(mCacheManager, llmRequest); - auto pickUpConnections = pickRecvConnections(connections, selfConfig, selfIdx, destConfig); + auto pickUpConnections = pickRecvConnections(connections.size(), selfConfig, selfIdx, destConfig); TLLM_LOG_DEBUG("pickUpConnections size: %d connections size: %d", pickUpConnections.size(), connections.size()); std::vector recvBufferTmps; @@ -453,7 +459,6 @@ void CacheFormatter::formatInput(LlmRequest const& llmRequest, { NVTX3_SCOPED_RANGE(formatInputRecvBuffer); - auto reqId = llmRequest.getContextPhaseParams().value().getReqId(); auto dataType = mCacheManager->getPrimaryPool(0)->getDataType(); bool layerWise = common::getEnvDisaggLayerwise() && numPools == 1; if (layerWise) @@ -494,14 +499,14 @@ void CacheFormatter::formatInput(LlmRequest const& llmRequest, TLLM_LOG_DEBUG("Buffer %d of pool %d shape = %s", idx, poolIdx, runtime::ITensor::toString(recvBufferTmps[idx]->getShape()).c_str()); } - for (auto const& connection : pickUpConnections) + for (size_t i = 0; i < pickUpConnections.size(); i++) { TLLM_LOG_DEBUG("Receive layer %d(%d-%d)", layerIdx, poolIdx, layerIdxInPool); // Buffer dim: [numLayersInPool * layerVolume] auto layer = runtime::ITensor::slice(recvBufferTmps[idx], layerIdxInPool * layerVolume, layerVolume); llmRequest.updateKvCacheSize((*layer).getSizeInBytes()); - TransferHelper::recvBuffer(*connection, *layer, reqId); + session.recv(pickUpConnections[i], layer->data(), layer->getSizeInBytes()); idx++; } } @@ -528,14 +533,14 @@ void CacheFormatter::formatInput(LlmRequest const& llmRequest, TLLM_CHECK(pickUpConnections.size() == 1); TLLM_CUDA_CHECK(cudaSetDevice(deviceId)); - for (auto const& connection : pickUpConnections) + for (size_t i = 0; i < pickUpConnections.size(); i++) { for (auto const& [window, blocks] : outputBuffersPerWindow) { for (auto const& block : blocks) { llmRequest.updateKvCacheSize((*block).getSizeInBytes()); - TransferHelper::recvBuffer(*connection, *block, reqId); + session.recv(pickUpConnections[i], block->data(), block->getSizeInBytes()); } } } @@ -586,7 +591,7 @@ void CacheFormatter::formatInput(LlmRequest const& llmRequest, else { auto* agentConnnecion - = dynamic_cast(pickUpConnections[0]); + = dynamic_cast(connections[pickUpConnections[0]]); if (agentConnnecion != nullptr) { cacheBufferId = agentConnnecion->getCacheBufferId(); @@ -639,8 +644,8 @@ void CacheFormatter::formatInput(LlmRequest const& llmRequest, size_t blockIdx = idx % (blockNum); size_t recvBufferIdx = blockIdx * pickUpConnections.size() + commIdx; llmRequest.updateKvCacheSize((*recvSplitCaches[recvBufferIdx]).getSizeInBytes()); - TransferHelper::recvBuffer( - *pickUpConnections[processIdx], *recvSplitCaches.at(recvBufferIdx), reqId); + auto& buffer = recvSplitCaches.at(recvBufferIdx); + session.recv(pickUpConnections[processIdx], buffer->data(), buffer->getSizeInBytes()); idx++; } } @@ -649,15 +654,16 @@ void CacheFormatter::formatInput(LlmRequest const& llmRequest, if (processIdx >= remainNoCoverTargetNum) { llmRequest.updateKvCacheSize((*recvSplitCaches.at(processIdx)).getSizeInBytes()); - TransferHelper::recvBuffer(*pickUpConnections[processIdx], *recvSplitCaches[processIdx], reqId); + auto& buffer = recvSplitCaches[processIdx]; + session.recv(pickUpConnections[processIdx], buffer->data(), buffer->getSizeInBytes()); } else if (bufferCoverTargetNum > 0) { auto recvBufferIdx = processIdx % bufferCoverTargetNum + remainNoCoverTargetNum; // caches.at(recvBufferIdx) is allocated by cudaMalloc llmRequest.updateKvCacheSize((*recvSplitCaches.at(recvBufferIdx)).getSizeInBytes()); - TransferHelper::recvBuffer( - *pickUpConnections[processIdx], *recvSplitCaches.at(recvBufferIdx), reqId); + auto& buffer = recvSplitCaches.at(recvBufferIdx); + session.recv(pickUpConnections[processIdx], buffer->data(), buffer->getSizeInBytes()); bufferManager.copy(*recvSplitCaches.at(recvBufferIdx), *recvSplitCaches[processIdx]); bufferManager.getStream().synchronize(); } @@ -674,7 +680,7 @@ void CacheFormatter::formatInput(LlmRequest const& llmRequest, auto copySlice = runtime::ITensor::slice( recvSplitCaches[processIdx], targetBufferSize - remainRecvSize, recvSize); llmRequest.updateKvCacheSize((*recvSlice).getSizeInBytes()); - TransferHelper::recvBuffer(*pickUpConnections[processIdx], *recvSlice, reqId); + session.recv(pickUpConnections[processIdx], recvSlice->data(), recvSlice->getSizeInBytes()); bufferManager.copy(*recvSlice, *copySlice); bufferManager.getStream().synchronize(); remainRecvSize -= recvSize; diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.h b/cpp/tensorrt_llm/batch_manager/cacheFormatter.h index a30a3b0a1885..36f6f57d169a 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.h +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.h @@ -26,36 +26,14 @@ #include "tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iTensor.h" +#include "tensorrt_llm/runtime/utils/mpiUtils.h" #include -#include #include #include -#include namespace tensorrt_llm::batch_manager::kv_cache_manager { -class TransferHelper -{ -public: - static void sendBuffer( - executor::kv_cache::Connection const& connection, runtime::IBuffer const& buf, uint64_t requestId) - { - int const tag = ((requestId & 0xFFF) << 8) | (kDATA_TAG & 0xFF); - connection.send(executor::kv_cache::DataContext{tag}, buf.data(), buf.getSizeInBytes()); - } - - static void recvBuffer(executor::kv_cache::Connection const& connection, runtime::IBuffer& buf, uint64_t requestId) - { - int const tag = ((requestId & 0xFFF) << 8) | (kDATA_TAG & 0xFF); - connection.recv(executor::kv_cache::DataContext{tag}, buf.data(), buf.getSizeInBytes()); - } - -private: - static constexpr int32_t kDATA_TAG{43}; -}; - BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, LlmRequest const& llmRequest); BlockRange getBlockRangeForReceiving(BaseKVCacheManager* cacheManager, LlmRequest const& llmRequest); @@ -67,15 +45,13 @@ class BaseCacheFormatter using SizeType32 = tensorrt_llm::runtime::SizeType32; using CacheState = executor::kv_cache::CacheState; - virtual void formatOutput(LlmRequest const& llmRequest, - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig, runtime::BufferManager const& bufferManager) - = 0; + /// @brief Format the cache data into bytes for sending. + /// @param session The transfer session. + virtual void format(TransferSession& session) = 0; - virtual void formatInput(LlmRequest const& llmRequest, - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig, runtime::BufferManager const& bufferManager) - = 0; + /// @brief Unformat the cache data from received bytes. + /// @param session The transfer session. + virtual void unformat(TransferSession& session) = 0; /// @brief Determine whether the sender is applicable to the source and target. /// @param selfConfig Source data arrangement. @@ -94,9 +70,8 @@ class BaseCacheFormatter [[nodiscard]] virtual BaseKVCacheManager* getCacheManager() const noexcept = 0; - [[nodiscard]] virtual std::vector pickRecvConnections( - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig) const + [[nodiscard]] virtual std::vector pickRecvConnections( + size_t numConnections, CacheState const& selfConfig, SizeType32 selfIdx, CacheState const& destConfig) const = 0; /// @brief Destructor. @@ -116,13 +91,9 @@ class CacheFormatter final : public BaseCacheFormatter TLLM_CHECK(mCacheTransBufferManager); } - void formatOutput(LlmRequest const& llmRequest, - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig, runtime::BufferManager const& bufferManager) override; + void format(TransferSession& session) override; - void formatInput(LlmRequest const& llmRequest, - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig, runtime::BufferManager const& bufferManager) override; + void unformat(TransferSession& session) override; [[nodiscard]] bool inquireSupport(CacheState const& selfConfig, CacheState const& destConfig) const override; @@ -138,9 +109,8 @@ class CacheFormatter final : public BaseCacheFormatter } static bool needSendCache(CacheState const& selfConfig, CacheState const& destConfig, runtime::SizeType32 selfIdx); - std::vector pickRecvConnections( - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig) const override; + std::vector pickRecvConnections(size_t numConnections, CacheState const& selfConfig, SizeType32 selfIdx, + CacheState const& destConfig) const override; private: BaseKVCacheManager* mCacheManager; diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp index 51b06feaf71e..1a3aed54f416 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp @@ -210,7 +210,7 @@ CacheTransBufferManager::CacheTransBufferManager( { auto poolIdx = mCacheManager->getBlockManager().getLayerPoolIdx(layerId); auto windowSize = static_cast(mCacheManager->getBlockManager().getPoolWindowSize(poolIdx)); - auto validTokenNum = windowSize < maxNumTokens.value() ? windowSize : maxNumTokens.value(); + auto validTokenNum = (windowSize < maxNumTokens.value() ? windowSize : maxNumTokens.value()); bufferSizeFromMaxNumToken += validTokenNum * kvCacheByteSizePerTokenPerLayer; } } @@ -230,26 +230,37 @@ CacheTransBufferManager::CacheTransBufferManager( TLLM_LOG_INFO( "CacheTransBufferManager: mMaxNumTokens:%ld, mRecvBufferCount:%ld, " "mSendBufferCount:%ld,mTransferBufferSize:%ld, mPreAllocBufferSize:%ld,mOnlyUseDynamicBuffer:%d " - "mUseFabricMemory:%d", + "mUseFabricMemory:%d mDataType:%d", maxNumTokens.has_value() ? maxNumTokens.value() : 0, mRecvBufferCount, mSendBufferCount, mTransferBufferSize, - mPreAllocBufferSize, mOnlyUseDynamicBuffer, mUseFabricMemory); - bool to_allocate = common::getEnvUseMPIKvCache() || common::getEnvUseUCXKvCache() || common::getEnvUseNixlKvCache(); + mPreAllocBufferSize, mOnlyUseDynamicBuffer, mUseFabricMemory, mDataType); - TLLM_CHECK_WITH_INFO(to_allocate, "CacheTransBufferManager: to_allocate is false"); allocateBuffer(); } -size_t CacheTransBufferManager::preAllocBufferSize(std::optional maxNumTokens) +size_t CacheTransBufferManager::preAllocBufferSize( + std::map const& cacheSizeBytesPerTokenPerWindow, + std::optional const& cacheTransceiverConfig) { - bool to_allocate = common::getEnvUseMPIKvCache() || common::getEnvUseUCXKvCache() || common::getEnvUseNixlKvCache(); - if (!to_allocate) + if (!cacheTransceiverConfig.has_value()) { return 0; } + if (!cacheTransceiverConfig->getBackendType().has_value()) + { + return 0; + } + auto maxNumTokens = cacheTransceiverConfig->getMaxTokensInBuffer(); size_t TransferBufferSize = common::getEnvMemSizeForKVCacheTransferBuffer(); if (maxNumTokens.has_value()) { - TransferBufferSize = maxNumTokens.value(); + TransferBufferSize = 0; + for (auto const& [windowSize, cacheSizeBytesPerToken] : cacheSizeBytesPerTokenPerWindow) + { + auto validTokenNum + = (static_cast(windowSize) < maxNumTokens.value() ? static_cast(windowSize) + : maxNumTokens.value()); + TransferBufferSize += validTokenNum * cacheSizeBytesPerToken; + } } bool useFabricMemory = FabricMemory::supportFbaricMemory() && (!(common::getEnvKVCacheTransferUseSyncBuffer() || common::getEnvKVCacheTransferUseAsyncBuffer())); @@ -329,6 +340,14 @@ std::tuple, size_t, bool> CacheTransBuf size_t bufferCoverTargetNum = std::min( static_cast(targetNum), mTransferBufferSize / (targetBufferEleSize * common::getDTypeSize(mDataType))); TLLM_LOG_DEBUG("getOrAllocateBuffers bufferCoverTargetNum:%d", bufferCoverTargetNum); + if (bufferCoverTargetNum < static_cast(targetNum)) + { + TLLM_LOG_WARNING( + "CacheTransceiver getOrAllocateBuffers: bufferCoverTargetNum:%d < targetNum:%d, may use dynamic buffer, " + "it's better to increase MaxTokensInBuffer in cacheTransceiverConfig, otherwise, the performance may " + "be degraded", + bufferCoverTargetNum, targetNum); + } if (bufferId.has_value()) { TLLM_CHECK(static_cast(bufferId.value()) < concurrenceResource.mBuffers.size()); diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h index d534e2b4ac68..e7b050388fe6 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h @@ -18,6 +18,7 @@ #pragma once #include "tensorrt_llm/batch_manager/kvCacheManager.h" +#include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" #include @@ -59,7 +60,8 @@ class CacheTransBufferManager CacheTransBufferManager( KVCacheManager::BaseKVCacheManager* cacheManager, std::optional maxNumTokens = std::nullopt); - static size_t preAllocBufferSize(std::optional maxNumTokens = std::nullopt); + static size_t preAllocBufferSize(std::map const& cacheSizeBytesPerTokenPerWindow, + std::optional const& cacheTransceiverConfig = std::nullopt); std::optional assignBufferIndexForSend(); void freeBufferIndexForSend(std::optional bufferId); diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 3dd85b7dd4f4..599a89cef037 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -62,41 +62,49 @@ std::unique_ptr CacheTransceiverFactory::createCacheTransc runtime::WorldConfig const& worldConfig, executor::kv_cache::CacheState::AttentionType attentionType, std::optional cacheTransceiverConfig) { - - std::optional commType; - if (common::getEnvUseUCXKvCache()) - { - commType = CacheTransceiver::CommType::UCX; - TLLM_LOG_INFO("Enable UCX KV cache transport."); - } - else if (common::getEnvUseNixlKvCache()) + if (!cacheTransceiverConfig.has_value() || !cacheTransceiverConfig.value().getBackendType().has_value()) { - commType = CacheTransceiver::CommType::NIXL; - TLLM_LOG_INFO("Enable NIXL KV cache transport."); + TLLM_LOG_INFO("CacheTransceiver is disabled."); + return nullptr; } - else if (common::getEnvUseMPIKvCache()) + auto backendType = cacheTransceiverConfig.value().getBackendType(); + if (backendType.value() == executor::CacheTransceiverConfig::BackendType::DEFAULT) { - commType = CacheTransceiver::CommType::MPI; - TLLM_LOG_INFO("Enable MPI KV cache transport."); + if (common::getEnvUseUCXKvCache()) + { + backendType = executor::CacheTransceiverConfig::BackendType::UCX; + TLLM_LOG_INFO("Enable UCX KV cache transport."); + } + else if (common::getEnvUseNixlKvCache()) + { + backendType = executor::CacheTransceiverConfig::BackendType::NIXL; + TLLM_LOG_INFO("Enable NIXL KV cache transport."); + } + else if (common::getEnvUseMPIKvCache()) + { + backendType = executor::CacheTransceiverConfig::BackendType::MPI; + TLLM_LOG_INFO("Enable MPI KV cache transport."); + TLLM_LOG_WARNING("MPI KV cache transport is deprecated, please use UCX or NIXL instead."); + } + else + { + backendType = executor::CacheTransceiverConfig::BackendType::UCX; + } } + cacheTransceiverConfig.value().setBackendType(backendType); - if (commType) - { - executor::kv_cache::CacheState::ModelConfig cacheStateCfg{ - modelConfig.getNumKvHeadsPerLayer(), modelConfig.getSizePerHead(), modelConfig.getTokensPerBlock()}; + executor::kv_cache::CacheState::ModelConfig cacheStateCfg{ + modelConfig.getNumKvHeadsPerLayer(), modelConfig.getSizePerHead(), modelConfig.getTokensPerBlock()}; - return std::make_unique(cacheManager, commType.value(), cacheStateCfg, worldConfig, - modelConfig.getKvDataType(), attentionType, cacheTransceiverConfig); - } - return nullptr; + return std::make_unique( + cacheManager, cacheStateCfg, worldConfig, modelConfig.getKvDataType(), attentionType, cacheTransceiverConfig); } -CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, CommType commType, +CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, executor::kv_cache::CacheState::ModelConfig const& cacheStateModelCfg, runtime::WorldConfig const& worldConfig, nvinfer1::DataType dataType, executor::kv_cache::CacheState::AttentionType attentionType, std::optional cacheTransceiverConfig) - : mCommType{commType} - , mMpiGroupComm(std::addressof(tensorrt_llm::mpi::MpiComm::session())) + : mMpiGroupComm(std::addressof(tensorrt_llm::mpi::MpiComm::session())) , mCacheTransceiverConfig{cacheTransceiverConfig} { using tensorrt_llm::batch_manager::kv_cache_manager::CacheFormatter; @@ -138,59 +146,59 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa } } bool isMLA = attentionType == executor::kv_cache::CacheState::AttentionType::kMLA; - if (mCommType == CommType::MPI || mCommType == CommType::UCX || mCommType == CommType::NIXL) - { - std::optional maxNumTokens = std::nullopt; - if (mCacheTransceiverConfig.has_value()) - { - maxNumTokens = mCacheTransceiverConfig.value().getMaxNumTokens(); - } - mCacheTransBufferManager - = std::make_unique(cacheManager, maxNumTokens); - if (mCommType == CommType::UCX) - { - std::lock_guard lock(mDllMutex); - mWrapperLibHandle = dllOpen(UCX_WRAPPER_LIB_NAME); - TLLM_CHECK_WITH_INFO(mWrapperLibHandle != nullptr, "UCX wrapper library is not open correctly."); - auto load_sym = [](void* handle, char const* name) - { - void* ret = dllGetSym(handle, name); - TLLM_CHECK_WITH_INFO(ret != nullptr, - "Unable to load UCX wrapper library symbol, possible cause is that TensorRT-LLM library is not " - "built with UCX support, please rebuild in UCX-enabled environment."); - return ret; - }; - std::unique_ptr (*makeUcxConnectionManager)(); - *(void**) (&makeUcxConnectionManager) = load_sym(mWrapperLibHandle, "makeUcxConnectionManager"); - mManager = makeUcxConnectionManager(); - TLLM_LOG_INFO("UCX Connection Manager created"); - } - else if (mCommType == CommType::NIXL) - { - mManager = std::make_unique( - mCacheTransBufferManager.get()); - TLLM_LOG_INFO("NIXL Connection Manager created"); - } - else - { - mMpiWorldComm = std::addressof(tensorrt_llm::mpi::MpiComm::world()); - mManager = std::make_unique(mMpiWorldComm); - TLLM_LOG_INFO("MPI Connection Manager created"); - } + TLLM_CHECK_WITH_INFO(mCacheTransceiverConfig.has_value(), "CacheTransceiverConfig is not set."); + auto backendType = mCacheTransceiverConfig.value().getBackendType(); + TLLM_CHECK_WITH_INFO( + backendType.has_value() && (backendType.value() != executor::CacheTransceiverConfig::BackendType::DEFAULT), + " CacheTransceiverConfig::BackendType is not set."); - using tensorrt_llm::batch_manager::kv_cache_manager::MLACacheFormatter; - auto makeFormatter = [cacheManager, isMLA, this]() - { return createCacheFormatter(cacheManager, mCacheTransBufferManager.get(), isMLA); }; + std::optional maxNumTokens = mCacheTransceiverConfig.value().getMaxTokensInBuffer(); - mDataResponder = std::make_unique( - std::make_unique(mManager.get(), *mCacheState, worldConfig.getRank(), makeFormatter())); - mDataRequester = std::make_unique( - std::make_unique(mManager.get(), *mCacheState, worldConfig.getRank(), makeFormatter())); + mCacheTransBufferManager = std::make_unique(cacheManager, maxNumTokens); + if (backendType.value() == executor::CacheTransceiverConfig::BackendType::UCX) + { + std::lock_guard lock(mDllMutex); + mWrapperLibHandle = dllOpen(UCX_WRAPPER_LIB_NAME); + TLLM_CHECK_WITH_INFO(mWrapperLibHandle != nullptr, "UCX wrapper library is not open correctly."); + auto load_sym = [](void* handle, char const* name) + { + void* ret = dllGetSym(handle, name); + TLLM_CHECK_WITH_INFO(ret != nullptr, + "Unable to load UCX wrapper library symbol, possible cause is that TensorRT-LLM library is not " + "built with UCX support, please rebuild in UCX-enabled environment."); + return ret; + }; + std::unique_ptr (*makeUcxConnectionManager)(); + *(void**) (&makeUcxConnectionManager) = load_sym(mWrapperLibHandle, "makeUcxConnectionManager"); + mManager = makeUcxConnectionManager(); + TLLM_LOG_INFO("UCX Connection Manager created"); + } + else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL) + { + mManager = std::make_unique( + mCacheTransBufferManager.get()); + TLLM_LOG_INFO("NIXL Connection Manager created"); + } + else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MPI) + { + mMpiWorldComm = std::addressof(tensorrt_llm::mpi::MpiComm::world()); + mManager = std::make_unique(mMpiWorldComm); + TLLM_LOG_INFO("MPI Connection Manager created"); } else { - TLLM_THROW("Unsupported communication type."); + TLLM_THROW("Unsupported cache transceiver backend type "); } + + using tensorrt_llm::batch_manager::kv_cache_manager::MLACacheFormatter; + auto makeFormatter = [cacheManager, isMLA, this]() + { return createCacheFormatter(cacheManager, mCacheTransBufferManager.get(), isMLA); }; + + mDataResponder = std::make_unique( + std::make_unique(mManager.get(), *mCacheState, worldConfig.getRank(), makeFormatter())); + mDataRequester = std::make_unique( + std::make_unique(mManager.get(), *mCacheState, worldConfig.getRank(), makeFormatter())); + initializeCommState(); } diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 1d06ac0e860f..baa51f47e730 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -63,7 +63,10 @@ void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffe SizeType32 batchIdx{0}; for (auto const& llmReq : contextRequests) { - auto const currentSequenceLen = llmReq->mPromptLen + llmReq->getMaxNumGeneratedTokens(); + auto const disaggFirstGenTokenSize + = llmReq->getContextPhaseParams() ? llmReq->getContextPhaseParams().value().getFirstGenTokens().size() : 0; + auto const currentSequenceLen + = llmReq->mPromptLen + llmReq->getMaxNumGeneratedTokens() + disaggFirstGenTokenSize; // Get position of the current sequence in the decoder auto const seqSlot = llmReq->mSeqSlot.value(); batchSlotsRange[batchIdx] = seqSlot; diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 88f1380e0b34..a4617c0d53d9 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -363,8 +363,8 @@ class DataRequester::Impl llmRequest.getContextPhaseParams().value().getReqId()); llmRequest.setKvCacheTransferStart(std::chrono::steady_clock::now()); TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); - mReceiver->sendRequestInfo(llmRequest); - mReceiver->receiveSync(llmRequest); + auto session = mReceiver->sendRequestInfo(llmRequest); + mReceiver->receiveSync(session); llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index 638893f1a976..e92c6112de2f 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -34,6 +34,11 @@ namespace tensorrt_llm::batch_manager { +// TODO: unify the following class into a namespace like tensorrt_llm::transmission +using DataContext = tensorrt_llm::executor::kv_cache::DataContext; +using Connection = tensorrt_llm::executor::kv_cache::Connection; +using ConnectionManager = tensorrt_llm::executor::kv_cache::ConnectionManager; + // Used to store the information that needs to be sent to the context executor to ensure the generation // executor smoothly receives the data. class RequestInfo @@ -89,6 +94,84 @@ class RequestInfo executor::DataTransceiverState mTransState; }; +class TransferSession +{ +public: + TransferSession(std::vector connections, DataContext dataContext, + executor::DataTransceiverState const& selfState, executor::DataTransceiverState otherState, + runtime::BufferManager const& bufferManager, LlmRequest const* llmRequest = nullptr) + : mConnections(std::move(connections)) + , mDataContext(dataContext) + , mSelfState(&selfState) + , mOtherState(std::move(otherState)) + , mBufferManager(&bufferManager) + , mRequest(llmRequest) + { + TLLM_CHECK(!mConnections.empty()); + } + + [[nodiscard]] std::vector const& getConnections() const + { + return mConnections; + } + + // should be called only during the initialization of the TransferSession + void setConnection(size_t idx, Connection const* conn) + { + mConnections.at(idx) = conn; + } + + [[nodiscard]] DataContext const& getDataContext() const + { + return mDataContext; + } + + [[nodiscard]] executor::DataTransceiverState const& getSelfState() const + { + return *mSelfState; + } + + [[nodiscard]] executor::DataTransceiverState const& getOtherState() const + { + return mOtherState; + } + + [[nodiscard]] runtime::BufferManager const& getBufferManager() const + { + return *mBufferManager; + } + + void send(size_t idx, void const* data, size_t size) + { + mConnections.at(idx)->send(mDataContext, data, size); + } + + void recv(size_t idx, void* data, size_t size) + { + mConnections.at(idx)->recv(mDataContext, data, size); + } + + [[nodiscard]] LlmRequest const& getLlmRequest() const + { + TLLM_CHECK(mRequest != nullptr); + return *mRequest; + } + + // in DataSender, the LlmRequest is not available until the sendSync is called + void setLlmRequest(LlmRequest const& llmRequest) + { + mRequest = &llmRequest; + } + +private: + std::vector mConnections; + DataContext mDataContext; + executor::DataTransceiverState const* mSelfState; // stored in DataRequester/DataResponder + executor::DataTransceiverState mOtherState; + runtime::BufferManager const* mBufferManager; + LlmRequest const* mRequest; +}; + // Operators required for data transmission in specific communication protocols. class DataSender { @@ -123,11 +206,11 @@ class DataReceiver public: /// @brief Send the request information. /// @param llmRequest The request object to which the information belongs. - virtual void sendRequestInfo(LlmRequest const& llmRequest) = 0; + virtual TransferSession sendRequestInfo(LlmRequest const& llmRequest) = 0; /// @brief Synchronously receive data. - /// @param llmRequest The request object to which the data belongs. - virtual void receiveSync(LlmRequest const& llmRequest) = 0; + /// @param session The transfer session. + virtual void receiveSync(TransferSession& session) = 0; /// @brief Destructor. virtual ~DataReceiver() = default; diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiverImpl.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiverImpl.cpp index dddd491d504e..e8adabed7f20 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiverImpl.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiverImpl.cpp @@ -24,6 +24,12 @@ namespace tensorrt_llm::batch_manager { +static int32_t tagFromRequestId(LlmRequest::RequestIdType requestId) +{ + constexpr int32_t kDATA_TAG{43}; + return ((requestId & 0xFFF) << 8) | (kDATA_TAG & 0xFF); +} + DataSenderImpl::DataSenderImpl(executor::kv_cache::ConnectionManager* manager, executor::kv_cache::CacheState selfCacheState, SizeType32 selfIndex, std::unique_ptr formatter) : mManager{manager} @@ -75,33 +81,25 @@ DataSenderImpl::DataSenderImpl(executor::kv_cache::ConnectionManager* manager, peerRelativeRanks.begin(), peerRelativeRanks.end(), info.getTransState().getCommState()->getSelfIdx())); { std::unique_lock lk(mMtxForMap); - auto it = mRequestToComms.find(requestId); - if (it == mRequestToComms.end()) + auto it = mRequestToSession.find(requestId); + if (it == mRequestToSession.end()) { - int recvExpectCount = peerRelativeRanks.size(); - { - it = mRequestToComms.emplace(requestId, RequestMapInfo{}).first; - it->second.resize(recvExpectCount); - } + auto session = TransferSession(std::vector(peerRelativeRanks.size(), nullptr), + DataContext{tagFromRequestId(requestId)}, mSelfState, info.getTransState(), mBufferManager); + it = mRequestToSession.emplace(requestId, std::move(session)).first; } - it->second[peerIdx] = {connection, info.getTransState()}; + it->second.setConnection(peerIdx, connection); } return info; } void DataSenderImpl::sendSync(LlmRequest const& llmRequest) { - std::vector connections; - auto it = mRequestToComms.find(llmRequest.mRequestId); - TLLM_CHECK(it != mRequestToComms.end()); - auto const& reqToComm = it->second; - for (auto&& [connection, dataTransceiverState] : reqToComm) - { - connections.emplace_back(connection); - } - auto&& dataTransceiverState = reqToComm.at(0).second; - mFormatter->formatOutput(llmRequest, std::move(connections), mSelfState.getCacheState().value(), - mSelfState.getCommState().value().getSelfIdx(), dataTransceiverState.getCacheState().value(), mBufferManager); + auto it = mRequestToSession.find(llmRequest.mRequestId); + TLLM_CHECK(it != mRequestToSession.end()); + auto& session = it->second; + session.setLlmRequest(llmRequest); + mFormatter->format(session); } [[nodiscard]] executor::kv_cache::CommState const& DataSenderImpl::getCommState() const @@ -116,17 +114,17 @@ void DataSenderImpl::setCommState(executor::kv_cache::CommState commState) [[nodiscard]] size_t DataSenderImpl::getCounterpartsCount(LlmRequest::RequestIdType requestId) const { - auto it = mRequestToComms.find(requestId); - TLLM_CHECK(it != mRequestToComms.end()); - return it->second.size(); + auto it = mRequestToSession.find(requestId); + TLLM_CHECK(it != mRequestToSession.end()); + return it->second.getConnections().size(); } void DataSenderImpl::release(LlmRequest::RequestIdType requestId) { - auto it = mRequestToComms.find(requestId); - TLLM_CHECK(it != mRequestToComms.end()); + auto it = mRequestToSession.find(requestId); + TLLM_CHECK(it != mRequestToSession.end()); std::unique_lock lk(mMtxForMap); - mRequestToComms.erase(it); + mRequestToSession.erase(it); } DataReceiverImpl::DataReceiverImpl(executor::kv_cache::ConnectionManager* manager, @@ -140,7 +138,7 @@ DataReceiverImpl::DataReceiverImpl(executor::kv_cache::ConnectionManager* manage TLLM_CHECK(mFormatter); } -void DataReceiverImpl::sendRequestInfo(LlmRequest const& llmRequest) +TransferSession DataReceiverImpl::sendRequestInfo(LlmRequest const& llmRequest) { uint64_t requestId = llmRequest.getContextPhaseParams().value().getReqId(); auto const& contextState = llmRequest.getDataTransceiverState(); @@ -179,17 +177,17 @@ void DataReceiverImpl::sendRequestInfo(LlmRequest const& llmRequest) auto const* connection = connections.at(index); counterPartConnections.emplace_back(connection); } - auto pickUpConnections = mFormatter->pickRecvConnections(counterPartConnections, mSelfState.getCacheState().value(), + auto pickUpIdx = mFormatter->pickRecvConnections(counterParts.size(), mSelfState.getCacheState().value(), mSelfState.getCommState().value().getSelfIdx(), destCacheState); - for (auto connection : counterPartConnections) + for (size_t i = 0; i < counterPartConnections.size(); i++) { + auto const* connection = counterPartConnections[i]; // if Manager is agentConnectionManager, then send request info to agent auto* agentConnectionManager = dynamic_cast(mManager); if (agentConnectionManager != nullptr) { // TODO: index -> validConnectionIdx conversion - auto valideConnectionIdx - = std::find(pickUpConnections.begin(), pickUpConnections.end(), connection) - pickUpConnections.begin(); + auto valideConnectionIdx = std::find(pickUpIdx.begin(), pickUpIdx.end(), i) - pickUpIdx.begin(); auto* agentConnection = dynamic_cast(connection); TLLM_CHECK(agentConnection != nullptr); TLLM_CHECK(cacheBufferId.has_value()); @@ -201,23 +199,14 @@ void DataReceiverImpl::sendRequestInfo(LlmRequest const& llmRequest) sendRequestInfo(connection, requestInfo); } } + auto const& resource = getReceiveCacheResource(llmRequest); + return TransferSession(std::move(counterPartConnections), DataContext{tagFromRequestId(requestId)}, mSelfState, + contextState, resource->mBufferManager, &llmRequest); } -void DataReceiverImpl::receiveSync(LlmRequest const& llmRequest) +void DataReceiverImpl::receiveSync(TransferSession& session) { - auto const& contextState = llmRequest.getDataTransceiverState(); - auto const& commState = contextState.getCommState().value(); - auto const& destCacheState = contextState.getCacheState().value(); - std::vector connections; - for (auto index : mFormatter->getCounterparts( - mSelfState.getCacheState().value(), mSelfState.getCommState().value().getSelfIdx(), destCacheState)) - { - auto const* connection = mManager->getConnections(commState).at(index); - connections.emplace_back(connection); - } - auto const& resource = getReceiveCacheResource(llmRequest); - mFormatter->formatInput(llmRequest, std::move(connections), mSelfState.getCacheState().value(), - mSelfState.getCommState().value().getSelfIdx(), destCacheState, resource->mBufferManager); + mFormatter->unformat(session); } void DataReceiverImpl::sendRequestInfo(executor::kv_cache::Connection const* connection, RequestInfo const& info) diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiverImpl.h b/cpp/tensorrt_llm/batch_manager/dataTransceiverImpl.h index 8f328bfaede3..fa8d27283292 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiverImpl.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiverImpl.h @@ -44,8 +44,6 @@ class DataSenderImpl : public DataSender, public TransceiverTag { public: using SizeType32 = tensorrt_llm::runtime::SizeType32; - using RequestMapInfo - = std::vector>; DataSenderImpl(executor::kv_cache::ConnectionManager* manager, executor::kv_cache::CacheState selfCacheState, SizeType32 selfIndex, std::unique_ptr formatter); @@ -64,7 +62,7 @@ class DataSenderImpl : public DataSender, public TransceiverTag private: executor::kv_cache::ConnectionManager* mManager; - std::map mRequestToComms; + std::map mRequestToSession; executor::DataTransceiverState mSelfState; std::unique_ptr mFormatter; std::mutex mMtxForMap; @@ -79,9 +77,9 @@ class DataReceiverImpl : public DataReceiver, public TransceiverTag DataReceiverImpl(executor::kv_cache::ConnectionManager* manager, executor::kv_cache::CacheState selfCacheState, SizeType32 selfIndex, std::unique_ptr formatter); - void sendRequestInfo(LlmRequest const& llmRequest) override; + TransferSession sendRequestInfo(LlmRequest const& llmRequest) override; - void receiveSync(LlmRequest const& llmRequest) override; + void receiveSync(TransferSession& session) override; private: struct ReceiveCacheResource diff --git a/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp b/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp index f48e12d6c88f..fd67bb55e89d 100644 --- a/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp +++ b/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp @@ -31,7 +31,7 @@ namespace tensorrt_llm::batch_manager { DecoderInputBuffers::DecoderInputBuffers( - SizeType32 maxNumSequences, SizeType32 maxBatchSize, SizeType32 maxDecoderSteps, BufferManager const& manager) + SizeType32 maxBatchSize, SizeType32 maxDecoderSteps, BufferManager const& manager) { auto const maxBatchSizeShape = ITensor::makeShape({maxBatchSize}); auto const nvSizeType = TRTDataType::value; @@ -49,8 +49,6 @@ DecoderInputBuffers::DecoderInputBuffers( { forwardBatchSlots.emplace_back(BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize}), nvSizeType)); } - - logits.resize(maxNumSequences); } void DecoderInputBuffers::setupMedusaLogits(SizeType32 maxNumSequences, ModelConfig const& modelConfig) diff --git a/cpp/tensorrt_llm/batch_manager/guidedDecoder.cpp b/cpp/tensorrt_llm/batch_manager/guidedDecoder.cpp index 871a33e3ee55..a5a7502c330d 100644 --- a/cpp/tensorrt_llm/batch_manager/guidedDecoder.cpp +++ b/cpp/tensorrt_llm/batch_manager/guidedDecoder.cpp @@ -16,6 +16,7 @@ */ #include "tensorrt_llm/batch_manager/guidedDecoder.h" +#include "tensorrt_llm/batch_manager/decoderBuffers.h" #include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/kernels/logitsBitmask.h" @@ -136,8 +137,7 @@ void GuidedDecoder::build(ScheduledRequests const& scheduledRequests) } } -void GuidedDecoder::execute(ScheduledRequests const& scheduledRequests, BufferManager const& runtimeBufferManager, - std::vector const& decoderBuffersLogits) +void GuidedDecoder::execute(DecoderInputBuffers const& decoderInputBuffers, BufferManager const& runtimeBufferManager) { auto const& stream = runtimeBufferManager.getStream(); @@ -150,32 +150,28 @@ void GuidedDecoder::execute(ScheduledRequests const& scheduledRequests, BufferMa mCopyBufferManager.getStream().record(event); stream.wait(event); - SizeType32 batchIdx{0}; - if (mGuidedDecodingBackend == executor::GuidedDecodingConfig::GuidedDecodingBackend::kXGRAMMAR) + if (mGuidedDecodingBackend == executor::GuidedDecodingConfig::GuidedDecodingBackend::kXGRAMMAR + && !decoderInputBuffers.decoderRequests.empty()) { - for (auto const& requests : {scheduledRequests.contextRequests, scheduledRequests.generationRequests}) + SizeType32 batchIdx{0}; + for (size_t requestIdx = 0; requestIdx < decoderInputBuffers.decoderRequests.size(); ++requestIdx) { - for (auto const& llmReq : requests) + auto const& llmReq = decoderInputBuffers.decoderRequests.at(requestIdx); + + auto const& guidedDecodingParams = llmReq->getGuidedDecodingParams(); + if (guidedDecodingParams.has_value()) { - if (llmReq->isContextInitState() && !llmReq->isLastContextChunk()) - { - continue; - } - auto const& guidedDecodingParams = llmReq->getGuidedDecodingParams(); - if (guidedDecodingParams.has_value()) - { - auto const seqSlot = llmReq->mSeqSlot.value(); + auto const seqSlot = llmReq->mSeqSlot.value(); - auto const& logits = decoderBuffersLogits.at(seqSlot); - auto const logitsBitmask = ITensor::at(mLogitsBitmask, {seqSlot}); + auto const& logits = decoderInputBuffers.logits.at(requestIdx); + auto const logitsBitmask = ITensor::at(mLogitsBitmask, {seqSlot}); - // Use void* to unify the code for different mLogitsDtype - *reinterpret_cast(ITensor::at(mLogitsPtrVecHost, {batchIdx})->data()) = logits->data(); - *reinterpret_cast(ITensor::at(mLogitsBitmaskPtrVecHost, {batchIdx})->data()) - = logitsBitmask->data(); + // Use void* to unify the code for different mLogitsDtype + *reinterpret_cast(ITensor::at(mLogitsPtrVecHost, {batchIdx})->data()) = logits->data(); + *reinterpret_cast(ITensor::at(mLogitsBitmaskPtrVecHost, {batchIdx})->data()) + = logitsBitmask->data(); - ++batchIdx; - } + ++batchIdx; } } if (batchIdx > 0) diff --git a/cpp/tensorrt_llm/batch_manager/handleContextLogits.cpp b/cpp/tensorrt_llm/batch_manager/handleContextLogits.cpp index e7ead88fb349..df3840c14b46 100644 --- a/cpp/tensorrt_llm/batch_manager/handleContextLogits.cpp +++ b/cpp/tensorrt_llm/batch_manager/handleContextLogits.cpp @@ -76,6 +76,13 @@ SizeType32 HandleContextLogits::operator()(DecoderInputBuffers& inputBuffers, Re TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); NVTX3_SCOPED_RANGE(HandleContextLogits); + auto& decoderRequests = inputBuffers.decoderRequests; + decoderRequests.clear(); + decoderRequests.reserve(contextRequests.size()); + auto& allDecoderLogits = inputBuffers.logits; + allDecoderLogits.clear(); + allDecoderLogits.reserve(contextRequests.size()); + SizeType32 batchIndex{0}; SizeType32 logitsIndex{0}; // Copy logits into decoderBuffers.logits @@ -115,7 +122,6 @@ SizeType32 HandleContextLogits::operator()(DecoderInputBuffers& inputBuffers, Re // Get the logits from the last context token and draft tokens auto const numDecoderLogits = 1 + draftLength; auto const seqSlot = llmReq->mSeqSlot.value(); - auto& decoderLogits = inputBuffers.logits.at(seqSlot); TensorPtr logitsView = ITensor::slice(logits, logitsIndex - numDecoderLogits, numDecoderLogits); if (modelConfig.getSpeculativeDecodingMode().hasDraftLogits()) @@ -136,22 +142,28 @@ SizeType32 HandleContextLogits::operator()(DecoderInputBuffers& inputBuffers, Re TLLM_CHECK_DEBUG_WITH_INFO(tru::tensorHasInvalid(*logitsView, manager, "logits") == false, "Found invalid number (NaN or Inf) in logits"); - // Scatter the output logits to the 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 + + if (llmReq->isLastContextChunk()) { - auto const logitsViewShape = logitsView->getShape(); - decoderLogits - = ITensor::view(logitsView, ITensor::makeShape({logitsViewShape.d[0], 1, logitsViewShape.d[1]})); + 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; diff --git a/cpp/tensorrt_llm/batch_manager/handleGenerationLogits.cpp b/cpp/tensorrt_llm/batch_manager/handleGenerationLogits.cpp index a5cecc54751f..5018ae36290d 100644 --- a/cpp/tensorrt_llm/batch_manager/handleGenerationLogits.cpp +++ b/cpp/tensorrt_llm/batch_manager/handleGenerationLogits.cpp @@ -22,6 +22,7 @@ #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" @@ -82,6 +83,11 @@ void HandleGenerationLogits::operator()(DecoderInputBuffers& inputBuffers, Reque TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); NVTX3_SCOPED_RANGE(HandleGenerationLogits); + auto& decoderRequests = inputBuffers.decoderRequests; + decoderRequests.reserve(decoderRequests.size() + generationRequests.size()); + auto& allDecoderLogits = inputBuffers.logits; + allDecoderLogits.reserve(allDecoderLogits.size() + generationRequests.size()); + for (auto const& llmReq : generationRequests) { auto const reqBeamWidth = llmReq->getBeamWidthByIter(); @@ -101,8 +107,9 @@ void HandleGenerationLogits::operator()(DecoderInputBuffers& inputBuffers, Reque 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"); - auto& decoderLogits = inputBuffers.logits.at(seqSlot); - auto const logitsViewShape = logitsView->getShape(); + + TLLM_CHECK(llmReq->isGenerationInProgressState()); + TensorPtr decoderLogits; if (reqBeamWidth > 1) { decoderLogits = logitsView; @@ -110,9 +117,11 @@ void HandleGenerationLogits::operator()(DecoderInputBuffers& inputBuffers, Reque } else { - decoderLogits - = ITensor::view(logitsView, ITensor::makeShape({logitsViewShape.d[0], 1, logitsViewShape.d[1]})); + decoderLogits = logitsView; + decoderLogits->unsqueeze(1); } + decoderRequests.push_back(llmReq); + allDecoderLogits.emplace_back(std::move(decoderLogits)); if (llmReq->getReturnGenerationLogits()) { diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 540dee9148b7..d30ba27be3ab 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -76,14 +76,82 @@ std::list> chopVectorIntoBlocks( return blockedVectors; } +inline uint8_t getNthByte(SizeType32 hashPart, uint8_t byteIdx) noexcept +{ + return static_cast((hashPart >> (24 - byteIdx * 8)) & 0xFF); +} + +std::vector generateBlockHashExtraKeys( + tensorrt_llm::batch_manager::LlmRequest const& llmRequest, SizeType32 startTokenIdx, SizeType32 endTokenIdx) +{ + auto const multimodalHashes = llmRequest.getMultimodalHashes(); + auto const multimodalPositions = llmRequest.getMultimodalPositions(); + auto const multimodalLengths = llmRequest.getMultimodalLengths(); + + if (!multimodalHashes || !multimodalPositions || !multimodalLengths || !(*multimodalHashes) + || (*multimodalHashes)->empty() || !(*multimodalPositions) || (*multimodalPositions)->empty() + || !(*multimodalLengths) || (*multimodalLengths)->empty()) + { + return {}; + } + + if ((*multimodalHashes)->size() != (*multimodalPositions)->size() + || (*multimodalPositions)->size() != (*multimodalLengths)->size()) + { + TLLM_LOG_WARNING("Multimodal data arrays have mismatched sizes"); + return {}; + } + + std::vector extraKeys; // MmKey = std::pair, SizeType32> + extraKeys.reserve((*multimodalPositions)->size()); + std::array mmHashArray; + + for (size_t i = 0; i < (*multimodalPositions)->size(); ++i) + { + auto const& startPos = (*(*multimodalPositions))[i]; + auto const& length = (*(*multimodalLengths))[i]; + auto const& mmHashVector = (*(*multimodalHashes))[i]; + + TLLM_CHECK_WITH_INFO(mmHashVector.size() == 8, "Multimodal hash vector has unexpected size: %zu (expected 8)", + mmHashVector.size()); + + // mmHashVector[j] comes from Python's int(hex_chunk, 16) + // where hex_chunk like "00010203" means 0x00 is MSB and 0x03 is LSB (big endian) + // Convert 8x 32-bit integers into a 32-byte array preserving Blake3 hash byte order + // Example: hashPart = 0x00010203 → mmHashArray[0:3] = [0x00, 0x01, 0x02, 0x03] + for (size_t j = 0; j < 8; ++j) + { + auto const& hashPart = mmHashVector[j]; + for (uint8_t byteIdx = 0; byteIdx < 4; ++byteIdx) + { + mmHashArray[j * 4 + byteIdx] = getNthByte(hashPart, byteIdx); + } + } + + // Check if this multimodal content overlaps with the current block + if (endTokenIdx > startPos && startTokenIdx < startPos + length) + { + SizeType32 mmStartInBlock = (startPos >= startTokenIdx) ? 0 : startTokenIdx - startPos; + extraKeys.emplace_back(mmHashArray, mmStartInBlock); + } + } + + return extraKeys; +} + std::vector buildBlockKeys( std::list& blockedUniqueTokens, tensorrt_llm::batch_manager::LlmRequest const& llmRequest) { std::vector blockKeys; + + SizeType32 currentTokenIdx = 0; for (auto& uniqueTokens : blockedUniqueTokens) { - blockKeys.emplace_back( - llmRequest.getInputTokensExtraIds().has_value(), llmRequest.getLoraTaskId(), std::move(uniqueTokens)); + auto extraKeys = generateBlockHashExtraKeys(llmRequest, currentTokenIdx, currentTokenIdx + uniqueTokens.size()); + currentTokenIdx += uniqueTokens.size(); + + blockKeys.emplace_back(llmRequest.getInputTokensExtraIds().has_value(), llmRequest.getLoraTaskId(), + std::move(uniqueTokens), std::move(extraKeys)); } return blockKeys; } @@ -92,9 +160,11 @@ std::vector buildBlockKeys( namespace tensorrt_llm::batch_manager::kv_cache_manager { - size_t BlockKeyHasher::hash(BlockKey const& blockKey, std::size_t parentHash) noexcept { + // Hashing algorithm adapted from StackOverflow: + // https://stackoverflow.com/questions/664014/what-integer-hash-function-are-good-that-accepts-an-integer-hash-key + // Constants provide very good distribution - each input bit affects each output bit with ~50% probability. size_t seed = blockKey.uniqueTokens.size() ^ parentHash * UINT64_C(0xbf58476d1ce4e5b9); for (auto const& uniqueToken : blockKey.uniqueTokens) @@ -122,7 +192,36 @@ size_t BlockKeyHasher::hash(BlockKey const& blockKey, std::size_t parentHash) no c = c ^ (c >> 31); seed ^= c + 0x9e3779b9 + (seed << 6) + (seed >> 2); } - // TODO: support external hashes for multimodal + + // Add extra keys for multimodal data mixing in external multimodal item hash and token offset within this sequence + // block + if (!blockKey.extraKeys.empty()) + { + for (auto const& [mmHash, startOffset] : blockKey.extraKeys) + { + // Hash the multimodal hash array in 32-bit chunks (more efficient) + for (size_t i = 0; i < 32; i += 4) + { + // Combine 4 bytes into a 32-bit word (construct as little endian order) + uint32_t word = static_cast(mmHash[i]) | (static_cast(mmHash[i + 1]) << 8) + | (static_cast(mmHash[i + 2]) << 16) | (static_cast(mmHash[i + 3]) << 24); + + // Mix the word into the seed + word = ((word >> 16) ^ word) * 0x45d9f3b; + word = ((word >> 16) ^ word) * 0x45d9f3b; + word = (word >> 16) ^ word; + seed ^= word + 0x9e3779b9 + (seed << 6) + (seed >> 2); + } + + // Hash the start offset + uint64_t e = static_cast(startOffset); + e = (e ^ (e >> 30)) * UINT64_C(0xbf58476d1ce4e5b9); + e = (e ^ (e >> 27)) * UINT64_C(0x94d049bb133111eb); + e = e ^ (e >> 31); + seed ^= e + 0x9e3779b9 + (seed << 6) + (seed >> 2); + } + } + return seed; } @@ -2235,13 +2334,8 @@ BlocksPerWindow BaseKVCacheManager::calculateMaxNumBlocks(executor::KvCacheConfi cacheSizeBytesPerTokenPerWindow[windowSize] = cacheSizeBytesPerToken; } - auto const extraCostMemoryBytes = extraCostMemory - * std::accumulate(cacheSizeBytesPerTokenPerWindow.cbegin(), cacheSizeBytesPerTokenPerWindow.cend(), - SizeType32{0}, [](SizeType32 acc, auto const cost) { return acc + cost.second; }); - - TLLM_LOG_DEBUG( - "extraCostMemoryBytes [all windows] [Gib]: %0.2f", extraCostMemoryBytes / static_cast(1 << 30)); - + TLLM_LOG_DEBUG("extraCostMemory [Gib]: %0.2f", extraCostMemory / static_cast(1 << 30)); + allottedPrimaryMemBytes = allottedPrimaryMemBytes - extraCostMemory; auto const tokensPerBlock = modelConfig.getTokensPerBlock(); auto const calculatePrimaryBlocks = [&](SizeType32 windowSize, float windowSizeShare, SizeType32 cacheSizeBytesPerToken) diff --git a/cpp/tensorrt_llm/batch_manager/logitsPostProcessor.cpp b/cpp/tensorrt_llm/batch_manager/logitsPostProcessor.cpp index 10210c3f4eb0..dd34de0ef9a0 100644 --- a/cpp/tensorrt_llm/batch_manager/logitsPostProcessor.cpp +++ b/cpp/tensorrt_llm/batch_manager/logitsPostProcessor.cpp @@ -17,25 +17,24 @@ #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" -#include "tensorrt_llm/runtime/tllmRuntime.h" namespace tr = tensorrt_llm::runtime; 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; -bool LogitsPostProcessor::operator()(RequestVector const& contextRequests, RequestVector const& generationRequests, - bool replicateLogitsPostProcessor, std::vector& seqSlotLogits, tr::WorldConfig const& worldConfig, - tr::TllmRuntime& runtime, std::optional logitsPostProcessorBatched) const +bool LogitsPostProcessor::operator()(DecoderInputBuffers& inputBuffers, bool replicateLogitsPostProcessor, + tr::WorldConfig const& worldConfig, CudaStreamPtr const& stream, + std::optional logitsPostProcessorBatched) const { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); NVTX3_SCOPED_RANGE(LogitsPostProcessor); @@ -47,35 +46,28 @@ bool LogitsPostProcessor::operator()(RequestVector const& contextRequests, Reque std::vector> clientIdsVec; bool logitsPostProcessorIsApplied = false; - for (auto const& requests : {contextRequests, generationRequests}) + for (size_t batchIdx = 0; batchIdx < inputBuffers.decoderRequests.size(); ++batchIdx) { - for (auto const& llmReq : requests) + auto const& llmReq = inputBuffers.decoderRequests.at(batchIdx); + auto& logits = inputBuffers.logits.at(batchIdx); + + // Invoke non-batched processor or collect arguments for batched processor + if (llmReq->mLogitsPostProcessor) { - if (llmReq->isContextInitState() ? llmReq->isLastContextChunk() : llmReq->isGenerationInProgressState()) + logitsPostProcessorIsApplied = true; + if (replicateLogitsPostProcessor || worldConfig.isFirstTensorParallelRank()) { - // Invoke non-batched processor or collect arguments for batched processor - if (llmReq->mLogitsPostProcessor) - { - logitsPostProcessorIsApplied = true; - if (replicateLogitsPostProcessor || worldConfig.isFirstTensorParallelRank()) - { - auto& logits = seqSlotLogits.at(llmReq->mSeqSlot.value()); - (*llmReq->mLogitsPostProcessor)( - llmReq->mRequestId, logits, llmReq->getTokens(), runtime.getStreamPtr(), llmReq->mClientId); - } - } - else if (llmReq->mApplyLogitsPostProcessorBatched) - { - reqIdsVec.push_back(llmReq->mRequestId); - - auto& logits = seqSlotLogits.at(llmReq->mSeqSlot.value()); - logitsVec.push_back(logits); - - beamTokensVec.emplace_back(llmReq->getTokens()); - clientIdsVec.push_back(llmReq->mClientId); - } + (*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 @@ -84,7 +76,7 @@ bool LogitsPostProcessor::operator()(RequestVector const& contextRequests, Reque logitsPostProcessorIsApplied = true; if (replicateLogitsPostProcessor || worldConfig.isFirstTensorParallelRank()) { - (*logitsPostProcessorBatched)(reqIdsVec, logitsVec, beamTokensVec, runtime.getStreamPtr(), clientIdsVec); + (*logitsPostProcessorBatched)(reqIdsVec, logitsVec, beamTokensVec, stream, clientIdsVec); } } diff --git a/cpp/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.cpp b/cpp/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.cpp index 64dedbc44972..c9b2bb0b9371 100644 --- a/cpp/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.cpp +++ b/cpp/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.cpp @@ -33,7 +33,7 @@ using TensorPtr = MakeDecodingBatchInputOutput::TensorPtr; std::unique_ptr MakeDecodingBatchInputOutput::createDecoderBatchInputs( std::vector const& activeSlots, runtime::decoder::DecoderState const& decoderState, - std::vector const& logits, SizeType32 maxNumSequences, std::vector const& batchSlots) + std::vector const& decoderLogits, SizeType32 maxNumSequences, std::vector const& batchSlots) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); @@ -47,40 +47,35 @@ std::unique_ptr MakeDecodingBatchInputOutput::createDe batchSlots.at(step)->resize(maxNumSequences); } - std::vector batchIdx(maxDecoderSteps); + auto constexpr singleRequest = 1; + + std::vector batchSizes(maxDecoderSteps); + std::vector> batchLogits(maxDecoderSteps); auto maxActiveDecoderSteps = 1; - for (auto const slot : activeSlots) + 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[batchIdx[step]] = slot; - batchIdx[step]++; + batchSlotsRange[batchSizes[step]] = slot; + batchSizes[step]++; + TensorPtr 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(batchIdx[step]); - } - - auto constexpr singleRequest = 1; - std::vector> logitsVec(maxActiveDecoderSteps); - for (SizeType32 step = 0; step < maxActiveDecoderSteps; ++step) - { - auto batchSlotsRange = tr::BufferRange(*batchSlots.at(step)); - - for (auto slot : batchSlotsRange) - { - auto const& targetLogits = logits.at(slot); - TensorPtr logitsSlice = tr::ITensor::slice(targetLogits, step, singleRequest); - logitsVec.at(step).push_back(logitsSlice); - } + batchSlots.at(step)->resize(batchSizes[step]); } + batchLogits.resize(maxActiveDecoderSteps); - auto decodingInput = std::make_unique(logitsVec, maxActiveDecoderSteps); + auto decodingInput = std::make_unique(batchLogits, maxActiveDecoderSteps); decodingInput->batchSlots = batchSlots; TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); return decodingInput; @@ -89,21 +84,14 @@ std::unique_ptr MakeDecodingBatchInputOutput::createDe namespace { -std::pair, std::vector> getActiveSlots( - RequestVector const& contextRequests, RequestVector const& generationRequests) +std::pair, std::vector> getActiveSlots(RequestVector const& decoderRequests) { std::vector activeSlots; std::vector generationSteps; - for (auto const& requests : {contextRequests, generationRequests}) + for (auto const& llmReq : decoderRequests) { - for (auto const& llmReq : requests) - { - if (llmReq->isGenerationInProgressState() || llmReq->isLastContextChunk()) - { - activeSlots.push_back(llmReq->mSeqSlot.value()); - generationSteps.push_back(llmReq->getDecodingIter()); - } - } + activeSlots.push_back(llmReq->mSeqSlot.value()); + generationSteps.push_back(llmReq->getDecodingIter()); } return {activeSlots, generationSteps}; @@ -167,14 +155,13 @@ void setEagleInputs(tr::DecodingInput& dInput, RuntimeBuffers const& fusedRuntim } // namespace -std::unique_ptr MakeDecodingBatchInputOutput::operator()(RequestVector const& contextRequests, - RequestVector const& generationRequests, DecoderInputBuffers const& inputBuffers, +std::unique_ptr MakeDecodingBatchInputOutput::operator()(DecoderInputBuffers& inputBuffers, runtime::decoder::DecoderState& decoderState, runtime::ModelConfig const& modelConfig, SizeType32 maxNumSequences, OptionalRef fusedRuntimeBuffers) const { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto [activeSlots, generationSteps] = getActiveSlots(contextRequests, generationRequests); + auto [activeSlots, generationSteps] = getActiveSlots(inputBuffers.decoderRequests); auto decodingInput = createDecoderBatchInputs( activeSlots, decoderState, inputBuffers.logits, maxNumSequences, inputBuffers.forwardBatchSlots); diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index cff94e85544a..21ebabb309c6 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -38,20 +38,17 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager { // some context rank in connection -std::vector MLACacheFormatter::pickRecvConnections( - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig) const +std::vector MLACacheFormatter::pickRecvConnections( + size_t numConnections, CacheState const& selfConfig, SizeType32 selfIdx, CacheState const& destConfig) const { - TLLM_CHECK(!connections.empty()); - auto targetInfo = executor::kv_cache::targetIRanks(destConfig, selfConfig, selfIdx); - TLLM_CHECK(targetInfo.mIRanks.size() == connections.size()); - std::vector ret; + TLLM_CHECK(numConnections == targetInfo.mIRanks.size()); + std::vector ret; // targetInfo , mRanks [tpranks, dpranks] for (int i = 0; i < targetInfo.mDomainPPSize; i++) { - ret.push_back(connections.at(i)); + ret.push_back(i); } return ret; } @@ -87,13 +84,17 @@ bool MLACacheFormatter::needSendCache( return selfTpRank % (selfTPNum / destTPNum) == 0; } -void MLACacheFormatter::formatOutput(LlmRequest const& llmRequest, - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig, runtime::BufferManager const& bufferManager) +void MLACacheFormatter::format(TransferSession& session) { - NVTX3_SCOPED_RANGE(formatOutput); + NVTX3_SCOPED_RANGE(MLACacheFormatter_format); + auto const& llmRequest = session.getLlmRequest(); TLLM_LOG_DEBUG( mpi::MpiComm::world().getRank(), "Start sending KV cache for request ID: %ld.", llmRequest.mRequestId); + auto const& selfConfig = session.getSelfState().getCacheState().value(); + auto const& destConfig = session.getOtherState().getCacheState().value(); + auto const selfIdx = session.getSelfState().getCommState().value().getSelfIdx(); + auto const& connections = session.getConnections(); + auto& bufferManager = session.getBufferManager(); TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently only supports beam width 1."); TLLM_CHECK(!connections.empty()); // diff start @@ -103,7 +104,6 @@ void MLACacheFormatter::formatOutput(LlmRequest const& llmRequest, } // diff end - auto reqId = llmRequest.mRequestId; auto const numPools = mCacheManager->getBlockManager().getNumPools(); auto blockRange = getBlockRangeForSending(mCacheManager, llmRequest); @@ -130,11 +130,11 @@ void MLACacheFormatter::formatOutput(LlmRequest const& llmRequest, NVTX3_SCOPED_RANGE(sendBufferFun); TLLM_CUDA_CHECK(cudaSetDevice(deviceId)); - for (auto const& connection : connections) + for (size_t i = 0; i < connections.size(); i++) { for (auto const& block : inputKvCacheBlocks) { - TransferHelper::sendBuffer(*connection, *block, reqId); + session.send(i, block->data(), block->getSizeInBytes()); } } @@ -193,7 +193,7 @@ void MLACacheFormatter::formatOutput(LlmRequest const& llmRequest, if (cacheIdx < bufferCoverTargetNum) { size = outputSplitCaches.at(cacheIdx)->getSizeInBytes(); - TransferHelper::sendBuffer(*connections.at(processIdx), *outputSplitCaches.at(cacheIdx), reqId); + session.send(processIdx, outputSplitCaches.at(cacheIdx)->data(), size); } else if (bufferCoverTargetNum > 0) { @@ -202,7 +202,7 @@ void MLACacheFormatter::formatOutput(LlmRequest const& llmRequest, size = outputSplitCaches.at(sendBufferIdx)->getSizeInBytes(); bufferManager.copy(*outputSplitCaches.at(cacheIdx), *outputSplitCaches.at(sendBufferIdx)); bufferManager.getStream().synchronize(); - TransferHelper::sendBuffer(*connections.at(processIdx), *outputSplitCaches.at(sendBufferIdx), reqId); + session.send(processIdx, outputSplitCaches.at(sendBufferIdx)->data(), size); } else { @@ -220,7 +220,7 @@ void MLACacheFormatter::formatOutput(LlmRequest const& llmRequest, auto copyTargetSlice = runtime::ITensor::slice(preAllocSendBuffer, 0, sendSize); bufferManager.copy(*copySlice, *copyTargetSlice); bufferManager.getStream().synchronize(); - TransferHelper::sendBuffer(*connections.at(processIdx), *copyTargetSlice, reqId); + session.send(processIdx, copyTargetSlice->data(), sendSize); remainSendSize -= sendSize; } @@ -277,19 +277,21 @@ void MLACacheFormatter::formatOutput(LlmRequest const& llmRequest, mpi::MpiComm::world().getRank(), "End the sending of KV cache for the request ID: %ld.", llmRequest.mRequestId); } -void MLACacheFormatter::formatInput(LlmRequest const& llmRequest, - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig, runtime::BufferManager const& bufferManager) +void MLACacheFormatter::unformat(TransferSession& session) { - NVTX3_SCOPED_RANGE(formatInput); + NVTX3_SCOPED_RANGE(MLACacheFormatter_unformat); + auto const& llmRequest = session.getLlmRequest(); TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently only supports beam width 1."); TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "Start receiving KV cache for request ID: %ld, context request ID: %ld.", llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); - auto reqId = llmRequest.getContextPhaseParams().value().getReqId(); - TLLM_CHECK(!connections.empty()); + auto const& selfConfig = session.getSelfState().getCacheState().value(); + auto const& destConfig = session.getOtherState().getCacheState().value(); + auto const selfIdx = session.getSelfState().getCommState().value().getSelfIdx(); + auto const& connections = session.getConnections(); + auto& bufferManager = session.getBufferManager(); // diff start - auto pickUpConnections = pickRecvConnections(connections, selfConfig, selfIdx, destConfig); + auto pickUpConnections = pickRecvConnections(connections.size(), selfConfig, selfIdx, destConfig); // diff end auto blockRange = getBlockRangeForReceiving(mCacheManager, llmRequest); std::vector recvBufferTmps; @@ -319,11 +321,12 @@ void MLACacheFormatter::formatInput(LlmRequest const& llmRequest, NVTX3_SCOPED_RANGE(recvBufferFun); TLLM_CUDA_CHECK(cudaSetDevice(deviceId)); TLLM_CHECK(pickUpConnections.size() == 1); - for (auto const& connection : pickUpConnections) + for (size_t i = 0; i < pickUpConnections.size(); i++) { for (auto const& block : outputBuffers) { - TransferHelper::recvBuffer(*connection, *block, reqId); + llmRequest.updateKvCacheSize(block->getSizeInBytes()); + session.recv(pickUpConnections[i], block->data(), block->getSizeInBytes()); } } TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), @@ -375,15 +378,17 @@ void MLACacheFormatter::formatInput(LlmRequest const& llmRequest, if (processIdx >= remainNoCoverTargetNum) { - - TransferHelper::recvBuffer(*pickUpConnections.at(processIdx), *recvSplitCaches.at(processIdx), reqId); + auto& buffer = recvSplitCaches.at(processIdx); + llmRequest.updateKvCacheSize(buffer->getSizeInBytes()); + session.recv(pickUpConnections.at(processIdx), buffer->data(), buffer->getSizeInBytes()); } else if (bufferCoverTargetNum > 0) { auto recvBufferIdx = processIdx % bufferCoverTargetNum + remainNoCoverTargetNum; // caches.at(recvBufferIdx) is allocated by cudaMalloc - TransferHelper::recvBuffer( - *pickUpConnections.at(processIdx), *recvSplitCaches.at(recvBufferIdx), reqId); + auto& buffer = recvSplitCaches.at(recvBufferIdx); + llmRequest.updateKvCacheSize(buffer->getSizeInBytes()); + session.recv(pickUpConnections.at(processIdx), buffer->data(), buffer->getSizeInBytes()); bufferManager.copy(*recvSplitCaches.at(recvBufferIdx), *recvSplitCaches.at(processIdx)); bufferManager.getStream().synchronize(); } @@ -399,7 +404,8 @@ void MLACacheFormatter::formatInput(LlmRequest const& llmRequest, auto recvSlice = runtime::ITensor::slice(preAllocRecvBuffer, 0, recvSize); auto copySlice = runtime::ITensor::slice( recvSplitCaches.at(processIdx), targetBufferSize - remainRecvSize, recvSize); - TransferHelper::recvBuffer(*pickUpConnections.at(processIdx), *recvSlice, reqId); + llmRequest.updateKvCacheSize(recvSlice->getSizeInBytes()); + session.recv(pickUpConnections.at(processIdx), recvSlice->data(), recvSlice->getSizeInBytes()); bufferManager.copy(*recvSlice, *copySlice); bufferManager.getStream().synchronize(); remainRecvSize -= recvSize; diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.h b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.h index 9d977d733876..c96e000e612b 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.h +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.h @@ -35,13 +35,9 @@ class MLACacheFormatter final : public BaseCacheFormatter TLLM_CHECK(mCacheTransBufferManager); } - void formatOutput(LlmRequest const& llmRequest, - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig, runtime::BufferManager const& bufferManager) override; + void format(TransferSession& session) override; - void formatInput(LlmRequest const& llmRequest, - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig, runtime::BufferManager const& bufferManager) override; + void unformat(TransferSession& session) override; [[nodiscard]] bool inquireSupport(CacheState const& selfConfig, CacheState const& destConfig) const override; @@ -57,9 +53,8 @@ class MLACacheFormatter final : public BaseCacheFormatter } static bool needSendCache(CacheState const& selfConfig, CacheState const& destConfig, runtime::SizeType32 selfIdx); - std::vector pickRecvConnections( - std::vector const& connections, CacheState const& selfConfig, - SizeType32 selfIdx, CacheState const& destConfig) const override; + std::vector pickRecvConnections(size_t numConnections, CacheState const& selfConfig, SizeType32 selfIdx, + CacheState const& destConfig) const override; private: BaseKVCacheManager* mCacheManager; diff --git a/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp index 8eeca23df35f..f513f2a3a102 100644 --- a/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp @@ -591,9 +591,10 @@ SizeType32 PeftCacheManager::determineNumPages(std::shared_ptr llmRe TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); if (llmRequest->getLoraTaskId().has_value()) { + auto taskId = llmRequest->getLoraTaskId().value(); try { - return mHostLoraCache->determineNumPages(llmRequest->getLoraTaskId().value()); + return mHostLoraCache->determineNumPages(taskId); } catch (std::runtime_error& e) { @@ -601,10 +602,17 @@ SizeType32 PeftCacheManager::determineNumPages(std::shared_ptr llmRe { return mHostLoraCache->determineNumPages(llmRequest->getLoraConfig().value()); } - else + if (!llmRequest->getLoraWeights().has_value()) { - throw; + auto const reqId = llmRequest->mRequestId; + std::string errMsg + = "Request ID " + std::to_string(reqId) + " has no LoRA adapter weights while configured with LoRA task " + + std::to_string(taskId) + " that's not found in LoRA CPU cache." + " Note that currently a request with LoRA task that was already loaded is sent without its LoRA weights to save its serialization, copy and deserialization," + " so if this LoRA task was evicted from LoRA CPU cache, then its reuse is currently not supported."; + throw PeftTaskNotCachedException(errMsg); } + throw; } } TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 1bc80ac21564..80418b2bc730 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -264,10 +264,35 @@ TrtGptModelInflightBatching::TrtGptModelInflightBatching(std::shared_ptr const& maxAttentionWindowVec, bool isCrossAttention, SizeType32 kvFactor) + { + auto [numKvHeadsPerLayerBegin, numKvHeadsPerLayerEnd] = modelConfig.getNumKvHeadsPerLayerLocalRange( + worldConfig.getPipelineParallelism(), worldConfig.getPipelineParallelRank(), isCrossAttention); + auto numKvHeadsPerLayer = std::vector(numKvHeadsPerLayerBegin, numKvHeadsPerLayerEnd); + auto windowSizeLayers + = BaseKVCacheManager::groupLayersByWindowSize(maxAttentionWindowVec, modelConfig.getNbLayers()); + std::map cacheSizeBytesPerTokenPerWindow; + for (auto const& [windowSize, managedLayers] : windowSizeLayers) + { + auto const cacheSizePerToken = BaseKVCacheManager::calculateCacheSizePerTokenForSingleWindowSize( + modelConfig, managedLayers, isCrossAttention, kvFactor); + auto const cacheSizeBytesPerToken + = cacheSizePerToken * BufferDataType(modelConfig.getKvDataType()).getSize(); + cacheSizeBytesPerTokenPerWindow[windowSize] = cacheSizeBytesPerToken; + } + + return cacheSizeBytesPerTokenPerWindow; + }; auto cacheTransceiverConfig = executorConfig.getCacheTransceiverConfig().value_or(executor::CacheTransceiverConfig()); - auto cacheTransPreAllocaSize - = kv_cache_manager::CacheTransBufferManager::preAllocBufferSize(cacheTransceiverConfig.getMaxNumTokens()); + + auto const cacheSizeBytesPerTokenPerWindow = calculateCacheSizePerToken( + mModelConfig, mWorldConfig, getMaxAttentionWindowVec(), mModelConfig.useCrossAttention(), 2); + auto cacheTransPreAllocaSize = kv_cache_manager::CacheTransBufferManager::preAllocBufferSize( + cacheSizeBytesPerTokenPerWindow, cacheTransceiverConfig); auto const [freePrimaryMemBytes, freeSecondaryMemBytes] = BaseKVCacheManager::calculateFreeMemBytes(mRuntime->getBufferManager(), kvCacheConfig); @@ -879,8 +904,9 @@ void TrtGptModelInflightBatching::forwardSync() { // TODO: skip if sending layer-wise { - TLLM_CHECK_WITH_INFO( - mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration."); + TLLM_CHECK_WITH_INFO(mCacheTransceiver, + "Disaggregated serving is not enabled, please check the configuration of " + "cacheTransceiverConfig."); mCacheTransceiver->respondAndSendAsync(llmReq.get()); } mSeqSlotManager->freeSequenceSlot(llmReq->mRequestId); @@ -1504,7 +1530,7 @@ void TrtGptModelInflightBatching::createBuffers(executor::DecodingConfig const& for (SizeType32 i = 0; i < mNumMicroBatches; ++i) { mDecoderInputBuffers.emplace_back( - getMaxNumSequences(), getMaxBatchSize(), mModelConfig.getMaxDecodingTokens(), mRuntime->getBufferManager()); + getMaxBatchSize(), mModelConfig.getMaxDecodingTokens(), mRuntime->getBufferManager()); mDecoderInputBuffers.back().setupMedusaLogits(getMaxNumSequences(), mModelConfig); mDecoderOutputBuffers.emplace_back(getMaxNumSequences(), mOperatingBeamWidth, getMaxSequenceLen(), mModelConfig.getMaxDecodingTokens(), mRuntime->getBufferManager()); @@ -1780,8 +1806,8 @@ void TrtGptModelInflightBatching::executeStep( bufferCast(*mBuffers[bufferId]->transformerBuffers->contextProgressHost)[0] = progress.get(); if (progress) { - TLLM_CHECK_WITH_INFO( - mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration."); + TLLM_CHECK_WITH_INFO(mCacheTransceiver, + "Disaggregated serving is not enabled, please check the configuration of cacheTransceiverConfig."); mCacheTransceiver->respondAndSendLayerWise(layerWiseRequests, progress); } } @@ -2003,7 +2029,6 @@ runtime::CudaEvent TrtGptModelInflightBatching::decoderStepAsync(ScheduledReques NVTX3_SCOPED_RANGE(decoderStepAsync); auto& decoderInputBuffers = mDecoderInputBuffers.at(getFusedBufferId()); - auto& seqSlotLogits = decoderInputBuffers.logits; auto const contextBufferId = mCtxGenFusion ? getFusedBufferId() : getContextBufferId(); auto& contextRuntimeBuffers = mBuffers.at(contextBufferId); @@ -2023,22 +2048,20 @@ runtime::CudaEvent TrtGptModelInflightBatching::decoderStepAsync(ScheduledReques copyCacheIndirectionFromOutputsToInputs(scheduledRequests, genBufferId); } - mLogitsPostProcessorIsApplied - = (*mLogitsPostProcessor)(scheduledRequests.contextRequests, scheduledRequests.generationRequests, - mReplicateLogitsPostProcessor, seqSlotLogits, mWorldConfig, *mRuntime, mLogitsPostProcessorBatched); + mLogitsPostProcessorIsApplied = (*mLogitsPostProcessor)(decoderInputBuffers, mReplicateLogitsPostProcessor, + mWorldConfig, mRuntime->getStreamPtr(), mLogitsPostProcessorBatched); if (mGuidedDecoder) { - mGuidedDecoder->execute(scheduledRequests, mRuntime->getBufferManager(), seqSlotLogits); + mGuidedDecoder->execute(decoderInputBuffers, mRuntime->getBufferManager()); } auto const fusedBufferId = getFusedBufferId(); auto& fusedRuntimeBuffers = mBuffers.at(fusedBufferId); auto& decodingInput = mDecodingInputs.at(mMicroBatchId); - decodingInput = (*mMakeDecodingBatchInputOutput)(scheduledRequests.contextRequests, - scheduledRequests.generationRequests, mDecoderInputBuffers.at(fusedBufferId), *mDecoderState, mModelConfig, - getMaxNumSequences(), *fusedRuntimeBuffers); + decodingInput = (*mMakeDecodingBatchInputOutput)(mDecoderInputBuffers.at(fusedBufferId), *mDecoderState, + mModelConfig, getMaxNumSequences(), *fusedRuntimeBuffers); auto decoderFinishEvent = mDecoder->forwardAsync(*mDecoderState, *decodingInput); diff --git a/cpp/tensorrt_llm/common/attentionOp.h b/cpp/tensorrt_llm/common/attentionOp.h index d19a9cbcc4e2..b738fdaf2fdb 100644 --- a/cpp/tensorrt_llm/common/attentionOp.h +++ b/cpp/tensorrt_llm/common/attentionOp.h @@ -341,6 +341,11 @@ class AttentionOp void debugCheckSemaphores(cudaStream_t stream); + [[nodiscard]] int getMultiProcessorCount() const + { + return mMultiProcessorCount; + } + [[nodiscard]] std::string toString() const; int mLayerIdx = -1; diff --git a/cpp/tensorrt_llm/deep_ep/CMakeLists.txt b/cpp/tensorrt_llm/deep_ep/CMakeLists.txt index 603f26796e62..f4c3f48bbb23 100644 --- a/cpp/tensorrt_llm/deep_ep/CMakeLists.txt +++ b/cpp/tensorrt_llm/deep_ep/CMakeLists.txt @@ -1,4 +1,4 @@ -set(DEEP_EP_COMMIT c381dadf43a85062f6a8947592017ee513abc70b) +set(DEEP_EP_COMMIT 7b15af835942675df041eca2dcb9930b880287e1) set(NVSHMEM_URL_HASH SHA256=eb2c8fb3b7084c2db86bd9fd905387909f1dfd483e7b45f7b3c3d5fcf5374b5a) @@ -36,6 +36,9 @@ if(NOT DEEP_EP_CUDA_ARCHITECTURES) return() endif() +# Ensure that dependent libraries are installed +find_library(MLX5_lib NAMES mlx5 REQUIRED) + # Prepare files # ============= diff --git a/cpp/tensorrt_llm/executor/cacheTransceiverConfig.cpp b/cpp/tensorrt_llm/executor/cacheTransceiverConfig.cpp index 1f392ef0583e..6919d213642e 100644 --- a/cpp/tensorrt_llm/executor/cacheTransceiverConfig.cpp +++ b/cpp/tensorrt_llm/executor/cacheTransceiverConfig.cpp @@ -21,24 +21,36 @@ namespace tensorrt_llm::executor { -CacheTransceiverConfig::CacheTransceiverConfig(std::optional maxNumTokens) - : mMaxNumTokens(maxNumTokens) +CacheTransceiverConfig::CacheTransceiverConfig( + std::optional backendType, std::optional maxNumTokens) + : mBackendType(backendType) + , mMaxTokensInBuffer(maxNumTokens) { } bool CacheTransceiverConfig::operator==(CacheTransceiverConfig const& other) const { - return mMaxNumTokens == other.mMaxNumTokens; + return mMaxTokensInBuffer == other.mMaxTokensInBuffer && mBackendType == other.mBackendType; } -std::optional CacheTransceiverConfig::getMaxNumTokens() const +void CacheTransceiverConfig::setBackendType(std::optional backendType) { - return mMaxNumTokens; + mBackendType = backendType; } -void CacheTransceiverConfig::setMaxNumTokens(size_t maxNumTokens) +void CacheTransceiverConfig::setMaxTokensInBuffer(std::optional maxTokensInBuffer) { - mMaxNumTokens = maxNumTokens; + mMaxTokensInBuffer = maxTokensInBuffer; +} + +std::optional CacheTransceiverConfig::getBackendType() const +{ + return mBackendType; +} + +std::optional CacheTransceiverConfig::getMaxTokensInBuffer() const +{ + return mMaxTokensInBuffer; } } // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp index 2ea6c26dc733..65718f0405d6 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -1258,19 +1258,22 @@ size_t Serialization::serializedSize(SchedulerConfig const& schedulerConfig) // CacheTransceiverConfig CacheTransceiverConfig Serialization::deserializeCacheTransceiverConfig(std::istream& is) { - auto maxNumTokens = su::deserialize>(is); - return CacheTransceiverConfig{maxNumTokens}; + auto backendType = su::deserialize>(is); + auto maxTokensInBuffer = su::deserialize>(is); + return CacheTransceiverConfig{backendType, maxTokensInBuffer}; } void Serialization::serialize(CacheTransceiverConfig const& cacheTransceiverConfig, std::ostream& os) { - su::serialize(cacheTransceiverConfig.getMaxNumTokens(), os); + su::serialize(cacheTransceiverConfig.getBackendType(), os); + su::serialize(cacheTransceiverConfig.getMaxTokensInBuffer(), os); } size_t Serialization::serializedSize(CacheTransceiverConfig const& cacheTransceiverConfig) { size_t totalSize = 0; - totalSize += su::serializedSize(cacheTransceiverConfig.getMaxNumTokens()); + totalSize += su::serializedSize(cacheTransceiverConfig.getBackendType()); + totalSize += su::serializedSize(cacheTransceiverConfig.getMaxTokensInBuffer()); return totalSize; } diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlTwoShotAllreduceKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlTwoShotAllreduceKernels.cu index f2e87e39ddac..2176ba759f47 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlTwoShotAllreduceKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlTwoShotAllreduceKernels.cu @@ -27,6 +27,10 @@ namespace tensorrt_llm::kernels::mnnvl { + +// Guard for internal helper functions +namespace +{ __device__ bool isNegZero(float v) { return v == 0.f && signbit(v); @@ -49,6 +53,12 @@ inline __device__ float toFloat<__nv_bfloat16>(__nv_bfloat16 val) return __bfloat162float(val); } +template <> +inline __device__ float toFloat<__nv_half>(__nv_half val) +{ + return __half2float(val); +} + template inline __device__ T fromFloat(float val) { @@ -61,6 +71,77 @@ inline __device__ __nv_bfloat16 fromFloat<__nv_bfloat16>(float val) return __float2bfloat16(val); } +template <> +inline __device__ __nv_half fromFloat<__nv_half>(float val) +{ + return __float2half(val); +} + +inline __device__ float2 loadfloat2(void const* ptr) +{ + float2 return_value; + asm volatile("ld.volatile.global.v2.f32 {%0, %1}, [%2];\n" : "=f"(return_value.x), "=f"(return_value.y) : "l"(ptr)); + return return_value; +} + +template +inline __device__ T divUp(T val, T divisor) +{ + return (val + divisor - 1) / divisor; +} + +__device__ struct __attribute__((aligned(32))) LamportFlags +{ + uint32_t buffer_size; + uint32_t input_offset; + uint32_t clear_offset; + uint32_t num_tokens_prev; + uint32_t* offset_access_ptr; + uint32_t* buffer_flags; + + __device__ explicit LamportFlags(uint32_t* buffer_flags) + : offset_access_ptr(&buffer_flags[4]) + , buffer_flags(buffer_flags) + { + uint4 flag = reinterpret_cast(buffer_flags)[0]; + buffer_size = flag.z; + input_offset = flag.x * (buffer_size << 1U); + clear_offset = flag.y * (buffer_size << 1U); + num_tokens_prev = flag.w; + } + + __device__ void cta_arrive() + { + __syncthreads(); + if (threadIdx.x == 0) + { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) + asm volatile("red.async.release.global.gpu.add.u32 [%0], %1;" ::"l"(offset_access_ptr), "r"(1) : "memory"); +#elif (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("red.global.gpu.add.u32 [%0], %1;" ::"l"(offset_access_ptr), "r"(1) : "memory"); +#else + atomicAdd(offset_access_ptr, 1); +#endif + } + } + + __device__ void wait_and_update(uint32_t num_tokens) + { + if (threadIdx.x == 0 && blockIdx.x == gridDim.x - 1 && blockIdx.y == 0) + { + while (*reinterpret_cast(offset_access_ptr) < gridDim.x * gridDim.y) + { + } + uint4 flag = reinterpret_cast(buffer_flags)[0]; + buffer_flags[0] = (flag.x + 1) % 3; + buffer_flags[1] = (flag.y + 1) % 3; + buffer_flags[3] = num_tokens; + *(offset_access_ptr) = 0; + } + } +}; +} // namespace + template __global__ void twoshot_allreduce_kernel(T* output_ptr, T* shard_ptr, T** input_ptrs, T* mcast_ptr, int num_tokens, int buffer_M, int token_dim, int rank, uint32_t* buffer_flags, bool wait_for_results) @@ -74,20 +155,14 @@ __global__ void twoshot_allreduce_kernel(T* output_ptr, T* shard_ptr, T** input_ cudaGridDependencySynchronize(); #endif - uint32_t* offset_access_ptr = &buffer_flags[3]; - // Buffer size is M * N, and we need two buffers for reduce-scatter and allgather - uint32_t buffer_size = (buffer_flags[2] << 1); - uint32_t input_offset = buffer_flags[0] * buffer_size; - uint32_t clear_offset = buffer_flags[1] * buffer_size; + LamportFlags flags(buffer_flags); - if (wait_for_results) - { - __syncthreads(); - if (threadIdx.x == 0) - { - atomicAdd(offset_access_ptr, 1); - } - } + // Capture the number of tokens in previous iteration so that we can properly clear the buffer + // The scatter stage will use the buffer in WORLD_SIZE granularity, thus we need to round up + uint32_t clr_toks_cta + = divUp(flags.num_tokens_prev > num_tokens ? flags.num_tokens_prev : num_tokens, WORLD_SIZE) + * WORLD_SIZE; + clr_toks_cta = divUp(clr_toks_cta, gridDim.x); if (elt < token_dim) { @@ -97,31 +172,34 @@ __global__ void twoshot_allreduce_kernel(T* output_ptr, T* shard_ptr, T** input_ T val = shard_ptr[token * token_dim + elt]; if (isNegZero(val)) val = fromFloat(0.f); - input_ptrs[dest_rank][input_offset + dest_token_offset * token_dim * WORLD_SIZE + rank * token_dim + elt] = val; - - // Reduce and broadcast + input_ptrs[dest_rank][flags.input_offset + dest_token_offset * token_dim * WORLD_SIZE + rank * token_dim + elt] + = val; - int global_token = token * WORLD_SIZE + rank; - if (global_token < num_tokens) + // Clear the buffer used by the previous call. Note the number of tokens to clear could be larger than the + // number of tokens in the current call. + for (int clr_tok = 0; clr_tok < clr_toks_cta; clr_tok++) { + uint32_t clr_token_idx = token + clr_tok * gridDim.x; + if (clr_token_idx < buffer_M) + { + input_ptrs[rank][flags.clear_offset + clr_token_idx * token_dim + elt] = fromFloat(-0.f); + } + } + // Reduce and broadcast + if ((token % WORLD_SIZE) == rank) + { + int local_token = token / WORLD_SIZE; float accum = 0.f; T values[WORLD_SIZE]; - - for (int r = 0; r < WORLD_SIZE; r++) - { - input_ptrs[rank][clear_offset + token * token_dim * WORLD_SIZE + r * token_dim + elt] - = fromFloat(-0.f); - } - while (1) { bool valid = true; for (int r = 0; r < WORLD_SIZE; r++) { - T volatile* lamport_ptr = (T volatile*) &input_ptrs[rank][input_offset - + token * token_dim * WORLD_SIZE + r * token_dim + elt]; + T volatile* lamport_ptr = (T volatile*) &input_ptrs[rank][flags.input_offset + + local_token * token_dim * WORLD_SIZE + r * token_dim + elt]; values[r] = *lamport_ptr; valid &= !isNegZero(values[r]); } @@ -132,7 +210,7 @@ __global__ void twoshot_allreduce_kernel(T* output_ptr, T* shard_ptr, T** input_ { accum += toFloat(values[r]); } - mcast_ptr[input_offset + buffer_M * token_dim + global_token * token_dim + elt] = fromFloat(accum); + mcast_ptr[flags.input_offset + buffer_M * token_dim + token * token_dim + elt] = fromFloat(accum); } } @@ -140,30 +218,47 @@ __global__ void twoshot_allreduce_kernel(T* output_ptr, T* shard_ptr, T** input_ cudaTriggerProgrammaticLaunchCompletion(); #endif - input_ptrs[rank][clear_offset + buffer_M * token_dim + token * token_dim + elt] = fromFloat(-0.f); + // Similarly clear broadcast buffer here + for (int clr_tok = 0; clr_tok < clr_toks_cta; clr_tok++) + { + uint32_t clr_token_idx = token + clr_tok * gridDim.x; + if (clr_token_idx < buffer_M) + { + input_ptrs[rank][flags.clear_offset + buffer_M * token_dim + clr_token_idx * token_dim + elt] + = fromFloat(-0.f); + } + } // Optionally wait for results if the next layer isn't doing the Lamport check if (wait_for_results) { - T volatile* lamport_ptr - = (T volatile*) &input_ptrs[rank][input_offset + buffer_M * token_dim + token * token_dim + elt]; - T val = *lamport_ptr; - while (isNegZero(val)) - val = *lamport_ptr; - - // Copy if requested - if (output_ptr) - output_ptr[token * token_dim + elt] = val; - if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0) + // Update the atomic counter to indicate the block has read the offsets + flags.cta_arrive(); + + // Only use a set of CTAs for lamport sync, reargange the grid + constexpr int ELTS_PER_LOAD = sizeof(float2) / sizeof(T); + // blockDim.x / ELTS_PER_LOAD should be at least the size of a warp (32) + if (threadIdx.x < (blockDim.x / ELTS_PER_LOAD)) { - // Make sure all blocks have finished reading the offsets, 2-D grid - while (*reinterpret_cast(offset_access_ptr) < gridDim.x * gridDim.y) + uint64_t current_pos = blockIdx.x * token_dim + blockIdx.y * blockDim.x + threadIdx.x * ELTS_PER_LOAD; + + void* lamport_ptr = (void*) &input_ptrs[rank][flags.input_offset + buffer_M * token_dim + current_pos]; + // We have 2 assumptions here: + // 1. The write is atomic in 8B granularity -> Each buffer in the buffer group should be aligned to 8B + // 2. The num_token * token_dim is divisible by ELTS_PER_LOAD (4 for BF16 and 2 for FP32) + float2 val = loadfloat2(lamport_ptr); + while (isNegZero(*(T*) &val)) { + val = loadfloat2(lamport_ptr); + } + if (output_ptr) + { + *((float2*) &output_ptr[current_pos]) = val; } - buffer_flags[0] = (buffer_flags[0] + 1) % 3; - buffer_flags[1] = (buffer_flags[1] + 1) % 3; - *(offset_access_ptr) = 0; } + + // Update the buffer flags + flags.wait_and_update(num_tokens); } } @@ -229,12 +324,28 @@ void twoshot_allreduce_op(AllReduceParams const& params) default: TLLM_CHECK_WITH_INFO(false, "TwoShot AllReduce]: unsupported world_size."); } } + else if (dtype == nvinfer1::DataType::kHALF) + { + switch (world_size) + { + case 2: LAUNCH_ALL_REDUCE_KERNEL(2, __nv_half); break; + case 4: LAUNCH_ALL_REDUCE_KERNEL(4, __nv_half); break; + case 8: LAUNCH_ALL_REDUCE_KERNEL(8, __nv_half); break; + case 16: LAUNCH_ALL_REDUCE_KERNEL(16, __nv_half); break; + case 32: LAUNCH_ALL_REDUCE_KERNEL(32, __nv_half); break; + case 64: LAUNCH_ALL_REDUCE_KERNEL(64, __nv_half); break; + default: TLLM_CHECK_WITH_INFO(false, "TwoShot AllReduce]: unsupported world_size."); + } + } else { TLLM_CHECK_WITH_INFO(false, "TwoShot AllReduce]: unsupported dtype."); } } +// Guard for internal helper functions +namespace +{ template __device__ void copy_f4(T_IN* dst, T_IN const* src) { @@ -251,18 +362,6 @@ __device__ void copy_f4_ldg(T_IN* dst, T_IN const* src) *dst4 = *src4; } -__device__ float4 loadfloat4(void const* ptr) -{ - - float return_value[4]; - - asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];\n" - : "=f"(return_value[0]), "=f"(return_value[1]), "=f"(return_value[2]), "=f"(return_value[3]) - : "l"(ptr)); - - return *(float4*) return_value; -} - template inline __device__ T add(T a, T b) { @@ -295,6 +394,19 @@ inline __device__ float block_reduce_sum(float val) return val; } +__device__ float4 loadfloat4(void const* ptr) +{ + + float4 return_value; + + asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];\n" + : "=f"(return_value.x), "=f"(return_value.y), "=f"(return_value.z), "=f"(return_value.w) + : "l"(ptr)); + + return return_value; +} +} // namespace + template __global__ void __launch_bounds__(128, 1) RMSNorm(T_IN* input_plus_residual, T_OUT* output_norm, T_IN const* buffer_input, T_IN const* gamma, float epsilon, @@ -321,20 +433,11 @@ __global__ void __launch_bounds__(128, 1) int offsets[NUM_INPUTS][DIM / (1 * ELTS_PER_THREAD * NUM_THREADS)]; - uint32_t* offset_access_ptr = &buffer_flags[3]; - // Buffer size is M * N, and we need two buffers for reduce-scatter and allgather - uint32_t buffer_size = buffer_flags[2]; - uint32_t buffer_offset = buffer_flags[0] * (buffer_size << 1); - T_IN const* input = &buffer_input[buffer_offset + buffer_size]; + LamportFlags flags(buffer_flags); + T_IN const* input = &buffer_input[flags.input_offset + flags.buffer_size]; cudaTriggerProgrammaticLaunchCompletion(); - __syncthreads(); - if (threadIdx.x == 0) - { - atomicAdd(offset_access_ptr, 1); - } - for (int i = 0; i < NUM_INPUTS; i++) { for (int j = 0; j < DIM / (1 * ELTS_PER_THREAD * NUM_THREADS); j++) @@ -361,7 +464,7 @@ __global__ void __launch_bounds__(128, 1) } __pipeline_commit(); - + flags.cta_arrive(); // Load all inputs bool valid = false; @@ -491,17 +594,7 @@ __global__ void __launch_bounds__(128, 1) = out4; } // Update the buffer pointers - if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0) - { - // Make sure all blocks have finished accessing the buffer - while (*reinterpret_cast(offset_access_ptr) != gridDim.x * gridDim.y) - { - } - buffer_flags[0] = (buffer_flags[0] + 1) % 3; - buffer_flags[1] = (buffer_flags[1] + 1) % 3; - *(offset_access_ptr) = 0; - } - __syncthreads(); + flags.wait_and_update(batch_size); #endif } @@ -512,8 +605,6 @@ void twoshot_rmsnorm(T* prenorm_output, T* normed_output, T const* input, T cons // input to rmsnorm is the buffer in the twoshot ar // We should use prenorm output to determine the actual used size - // int batch = normed_output.sizes()[0]; - // int dim = normed_output.sizes()[1]; float _epsilon{static_cast(epsilon)}; static constexpr int NUM_THREADS = 128; @@ -576,6 +667,20 @@ void twoshot_rmsnorm_op(RMSNormParams const& params) default: TLLM_CHECK_WITH_INFO(false, "[MNNVL TwoShot RMSNorm]: unsupported hidden_dim."); } } + else if (dtype == nvinfer1::DataType::kHALF) + { + switch (params.hidden_dim) + { + case 2048: LAUNCH_RMSNORM_KERNEL(__nv_half, 2048); break; + case 4096: LAUNCH_RMSNORM_KERNEL(__nv_half, 4096); break; + // Llama-4 Hidden Dimension + case 5120: LAUNCH_RMSNORM_KERNEL(__nv_half, 5120); break; + // DeepSeek Hidden Dimension + case 7168: LAUNCH_RMSNORM_KERNEL(__nv_half, 7168); break; + case 8192: LAUNCH_RMSNORM_KERNEL(__nv_half, 8192); break; + default: TLLM_CHECK_WITH_INFO(false, "[MNNVL TwoShot RMSNorm]: unsupported hidden_dim."); + } + } else { TLLM_CHECK_WITH_INFO(false, "[MNNVL TwoShot RMSNorm]: unsupported dtype."); diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_cubin.h b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_cubin.h index 612d1af7c522..66dc990d184d 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_cubin.h +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_cubin.h @@ -26,8 +26,6 @@ namespace kernels #ifndef EXCLUDE_SM_90 -extern unsigned char cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_tma_ws_sm90_cu_cubin[]; -extern unsigned char cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_tma_ws_sm90_cu_cubin[]; extern unsigned char cubin_fmha_v2_flash_attention_fp16_64_128_S_qkv_128_tma_ws_sm90_cu_cubin[]; extern unsigned char cubin_fmha_v2_flash_attention_fp16_64_128_S_qkv_128_softcapping_tma_ws_sm90_cu_cubin[]; extern unsigned char cubin_fmha_v2_flash_attention_fp16_64_128_S_q_kv_128_tma_ws_sm90_cu_cubin[]; @@ -195,10 +193,12 @@ extern void run_fmha_v2_flash_attention_bf16_64_128_S_qkv_104_tma_ws_sm90(Fused_ extern void run_fmha_v2_flash_attention_bf16_64_64_S_qkv_160_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_64_S_qkv_192_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_64_S_qkv_256_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); +extern void run_fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_64_S_qkv_256_softcapping_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_256_S_q_kv_32_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_256_S_q_kv_64_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_128_S_q_kv_72_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); +extern void run_fmha_v2_flash_attention_bf16_64_128_S_q_kv_192x128_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_256_S_q_paged_kv_32_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_256_S_q_paged_kv_40_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_256_S_q_paged_kv_48_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); @@ -210,10 +210,13 @@ extern void run_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_104_tma_ws_sm90 extern void run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_256_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); +extern void run_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_256_softcapping_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); +extern void run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_256_S_q_kv_32_softmax_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_256_S_q_kv_64_softmax_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_128_S_q_kv_72_softmax_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); +extern void run_fmha_v2_flash_attention_bf16_64_128_S_q_kv_192x128_softmax_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_256_S_qkv_32_alibi_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_256_S_qkv_40_alibi_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); extern void run_fmha_v2_flash_attention_bf16_64_256_S_qkv_48_alibi_tma_ws_sm90(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); @@ -1348,8 +1351,6 @@ extern void run_fmha_v2_flash_attention_fp16_fp32_64_16_S_q_paged_kv_256_softcap #ifndef EXCLUDE_SM_90 -extern uint32_t cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_tma_ws_sm90_cu_cubin_len; -extern uint32_t cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_tma_ws_sm90_cu_cubin_len; extern uint32_t cubin_fmha_v2_flash_attention_fp16_64_128_S_qkv_128_tma_ws_sm90_cu_cubin_len; extern uint32_t cubin_fmha_v2_flash_attention_fp16_64_128_S_qkv_128_softcapping_tma_ws_sm90_cu_cubin_len; extern uint32_t cubin_fmha_v2_flash_attention_fp16_64_128_S_q_kv_128_tma_ws_sm90_cu_cubin_len; @@ -1472,8 +1473,6 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 void (*launcher)(Fused_multihead_attention_params_v2& params, const Launch_params& launch_params, cudaStream_t stream); } sMhaKernelMetaInfosV2[] = { #ifndef EXCLUDE_SM_90 -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_causal_tma_ws_sm90_kernel", 213248, 384, 64, 1, 0, false, true, true, true, false, false, false, false, nullptr}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_causal_tma_ws_sm90_kernel", 213248, 384, 64, 1, 2, false, true, true, true, false, false, false, false, nullptr}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 64, 64, 64, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_fp16_64_32_ldgsts_sm90_kernel", 17408, 128, 0, 0, 0, false, false, false, false, true, false, false, false, run_fmha_v2_fp16_64_32_ldgsts_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 64, 64, 64, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_fp16_64_32_sliding_or_chunked_causal_ldgsts_sm90_kernel", 17408, 128, 0, 2, 0, false, false, false, false, true, false, false, false, run_fmha_v2_fp16_64_32_ldgsts_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 64, 64, 64, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_fp16_64_32_causal_ldgsts_sm90_kernel", 17408, 128, 0, 1, 0, false, false, false, false, true, false, false, false, run_fmha_v2_fp16_64_32_ldgsts_sm90}, @@ -1685,12 +1684,12 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_64_128_S_qkv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_64_128_S_qkv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_64_128_S_qkv_128_causal_tma_ws_sm90_kernel", 164096, 384, 64, 1, 0, false, true, true, false, false, false, false, false, nullptr}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_64_128_S_qkv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_64_128_S_qkv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_64_128_S_qkv_128_sliding_or_chunked_causal_tma_ws_sm90_kernel", 164096, 384, 64, 2, 0, false, true, true, false, false, false, false, false, nullptr}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_64_128_S_qkv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_64_128_S_qkv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_64_128_S_qkv_128_custom_mask_tma_ws_sm90_kernel", 164096, 384, 64, 3, 0, false, true, true, false, false, false, false, false, nullptr}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_160_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_160_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_160_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_160_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_160_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_160_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_192_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_192_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_192_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_192_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_192_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_192_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_160_causal_tma_ws_sm90_kernel", 147712, 384, 64, 1, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_160_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_160_sliding_or_chunked_causal_tma_ws_sm90_kernel", 147712, 384, 64, 2, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_160_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_160_custom_mask_tma_ws_sm90_kernel", 147712, 384, 64, 3, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_160_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_192_causal_tma_ws_sm90_kernel", 147712, 384, 64, 1, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_192_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_192_sliding_or_chunked_causal_tma_ws_sm90_kernel", 147712, 384, 64, 2, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_192_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_192_custom_mask_tma_ws_sm90_kernel", 147712, 384, 64, 3, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_192_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_256_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_256_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_256_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_256_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_256_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 0, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_256_tma_ws_sm90}, @@ -1736,12 +1735,12 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_causal_tma_ws_sm90_kernel", 164096, 384, 64, 1, 2, false, true, true, false, false, false, false, false, nullptr}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_sliding_or_chunked_causal_tma_ws_sm90_kernel", 164096, 384, 64, 2, 2, false, true, true, false, false, false, false, false, nullptr}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_custom_mask_tma_ws_sm90_kernel", 164096, 384, 64, 3, 2, false, true, true, false, false, false, false, false, nullptr}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_causal_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_sliding_or_chunked_causal_tma_ws_sm90_kernel", 147712, 384, 64, 2, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_custom_mask_tma_ws_sm90_kernel", 147712, 384, 64, 3, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_causal_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_sliding_or_chunked_causal_tma_ws_sm90_kernel", 147712, 384, 64, 2, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_custom_mask_tma_ws_sm90_kernel", 147712, 384, 64, 3, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_256_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_256_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_256_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_256_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_256_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 2, false, true, true, false, false, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_256_tma_ws_sm90}, @@ -1766,8 +1765,8 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 96, 96, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_128_S_qkv_96_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 0, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_128_S_qkv_96_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 104, 104, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_128_S_qkv_104_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 0, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_128_S_qkv_104_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_64_128_S_qkv_128_alibi_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_64_128_S_qkv_128_alibi_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_64_128_S_qkv_128_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 0, false, true, true, false, true, false, false, false, nullptr}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_160_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_160_alibi_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_192_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_192_alibi_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_160_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 0, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_160_alibi_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_192_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 0, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_192_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_qkv_256_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_qkv_256_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 256, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_256_S_q_paged_kv_32_causal_alibi_tma_ws_sm90_kernel", 73984, 384, 64, 1, 2, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_256_S_q_paged_kv_32_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 256, 40, 40, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_256_S_q_paged_kv_40_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_256_S_q_paged_kv_40_alibi_tma_ws_sm90}, @@ -1778,8 +1777,8 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 96, 96, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_96_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 2, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_96_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 104, 104, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_104_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 2, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_104_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 2, false, true, true, false, true, false, false, false, nullptr}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_alibi_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_alibi_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_160_alibi_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_192_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_256_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, false, true, false, false, false, run_fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_256_alibi_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 256, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_256_S_qkv_32_tma_ws_sm90_kernel", 73984, 384, 64, 0, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_256_S_qkv_32_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 256, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_256_S_qkv_32_causal_tma_ws_sm90_kernel", 73984, 384, 64, 1, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_256_S_qkv_32_tma_ws_sm90}, @@ -1812,15 +1811,16 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_qkv_128_causal_tma_ws_sm90_kernel", 164096, 384, 64, 1, 0, false, true, true, true, false, false, false, false, nullptr}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_qkv_128_sliding_or_chunked_causal_tma_ws_sm90_kernel", 164096, 384, 64, 2, 0, false, true, true, true, false, false, false, false, nullptr}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_qkv_128_custom_mask_tma_ws_sm90_kernel", 164096, 384, 64, 3, 0, false, true, true, true, false, false, false, false, nullptr}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_160_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_160_tma_ws_sm90}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_160_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_160_tma_ws_sm90}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_160_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_160_tma_ws_sm90}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_192_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_192_tma_ws_sm90}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_192_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_192_tma_ws_sm90}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_192_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_192_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_160_causal_tma_ws_sm90_kernel", 147712, 384, 64, 1, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_160_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_160_sliding_or_chunked_causal_tma_ws_sm90_kernel", 147712, 384, 64, 2, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_160_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_160_custom_mask_tma_ws_sm90_kernel", 147712, 384, 64, 3, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_160_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_192_causal_tma_ws_sm90_kernel", 147712, 384, 64, 1, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_192_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_192_sliding_or_chunked_causal_tma_ws_sm90_kernel", 147712, 384, 64, 2, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_192_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_192_custom_mask_tma_ws_sm90_kernel", 147712, 384, 64, 3, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_192_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_256_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_256_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_256_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_256_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_256_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_256_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_causal_tma_ws_sm90_kernel", 213248, 384, 64, 1, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_128_softcapping_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_128_softcapping_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_qkv_128_causal_softcapping_tma_ws_sm90_kernel", 164096, 384, 64, 1, 0, false, true, true, true, false, false, true, false, nullptr}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_128_softcapping_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_128_softcapping_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_qkv_128_sliding_or_chunked_causal_softcapping_tma_ws_sm90_kernel", 164096, 384, 64, 2, 0, false, true, true, true, false, false, true, false, nullptr}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_256_causal_softcapping_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, true, false, false, true, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_256_softcapping_tma_ws_sm90}, @@ -1833,6 +1833,7 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 72, 72, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_kv_72_custom_mask_tma_ws_sm90_kernel", 164096, 384, 64, 3, 1, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_128_S_q_kv_72_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_tma_ws_sm90_kernel", 164096, 384, 64, 0, 1, false, true, true, true, false, false, false, false, nullptr}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_custom_mask_tma_ws_sm90_kernel", 164096, 384, 64, 3, 1, false, true, true, true, false, false, false, false, nullptr}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_kv_192x128_causal_tma_ws_sm90_kernel", 213248, 384, 64, 1, 1, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_128_S_q_kv_192x128_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 256, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_256_S_q_paged_kv_32_tma_ws_sm90_kernel", 73984, 384, 64, 0, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_256_S_q_paged_kv_32_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 256, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_256_S_q_paged_kv_32_causal_tma_ws_sm90_kernel", 73984, 384, 64, 1, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_256_S_q_paged_kv_32_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 256, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_256_S_q_paged_kv_32_sliding_or_chunked_causal_tma_ws_sm90_kernel", 73984, 384, 64, 2, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_256_S_q_paged_kv_32_tma_ws_sm90}, @@ -1863,19 +1864,21 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_causal_tma_ws_sm90_kernel", 164096, 384, 64, 1, 2, false, true, true, true, false, false, false, false, nullptr}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_sliding_or_chunked_causal_tma_ws_sm90_kernel", 164096, 384, 64, 2, 2, false, true, true, true, false, false, false, false, nullptr}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_custom_mask_tma_ws_sm90_kernel", 164096, 384, 64, 3, 2, false, true, true, true, false, false, false, false, nullptr}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_tma_ws_sm90}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_tma_ws_sm90}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_tma_ws_sm90}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_tma_ws_sm90}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_tma_ws_sm90}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_causal_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_sliding_or_chunked_causal_tma_ws_sm90_kernel", 147712, 384, 64, 2, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_custom_mask_tma_ws_sm90_kernel", 147712, 384, 64, 3, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_causal_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_sliding_or_chunked_causal_tma_ws_sm90_kernel", 147712, 384, 64, 2, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_custom_mask_tma_ws_sm90_kernel", 147712, 384, 64, 3, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_256_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_256_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_256_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_256_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_256_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_256_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_causal_tma_ws_sm90_kernel", 213248, 384, 64, 1, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_causal_softcapping_tma_ws_sm90_kernel", 164096, 384, 64, 1, 2, false, true, true, true, false, false, true, false, nullptr}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_sliding_or_chunked_causal_softcapping_tma_ws_sm90_kernel", 164096, 384, 64, 2, 2, false, true, true, true, false, false, true, false, nullptr}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_256_causal_softcapping_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, true, false, false, true, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_256_softcapping_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_256_sliding_or_chunked_causal_softcapping_tma_ws_sm90_kernel", 196864, 384, 64, 2, 2, false, true, true, true, false, false, true, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_256_softcapping_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_causal_tma_ws_sm90_kernel", 213248, 384, 64, 1, 3, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_bf16_64_128_S_q_k_v_192x128_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 256, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_256_S_q_kv_32_softmax_tma_ws_sm90_kernel", 73984, 384, 64, 0, 1, false, true, true, true, false, false, false, true, run_fmha_v2_flash_attention_bf16_64_256_S_q_kv_32_softmax_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 256, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_256_S_q_kv_32_custom_mask_softmax_tma_ws_sm90_kernel", 73984, 384, 64, 3, 1, false, true, true, true, false, false, false, true, run_fmha_v2_flash_attention_bf16_64_256_S_q_kv_32_softmax_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 256, 64, 64, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_256_S_q_kv_64_softmax_tma_ws_sm90_kernel", 147712, 384, 64, 0, 1, false, true, true, true, false, false, false, true, run_fmha_v2_flash_attention_bf16_64_256_S_q_kv_64_softmax_tma_ws_sm90}, @@ -1884,6 +1887,7 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 72, 72, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_kv_72_custom_mask_softmax_tma_ws_sm90_kernel", 164096, 384, 64, 3, 1, false, true, true, true, false, false, false, true, run_fmha_v2_flash_attention_bf16_64_128_S_q_kv_72_softmax_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_softmax_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_softmax_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_softmax_tma_ws_sm90_kernel", 164096, 384, 64, 0, 1, false, true, true, true, false, false, false, true, nullptr}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_softmax_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_softmax_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_custom_mask_softmax_tma_ws_sm90_kernel", 164096, 384, 64, 3, 1, false, true, true, true, false, false, false, true, nullptr}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 192, 128, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_kv_192x128_causal_softmax_tma_ws_sm90_kernel", 213248, 384, 64, 1, 1, false, true, true, true, false, false, false, true, run_fmha_v2_flash_attention_bf16_64_128_S_q_kv_192x128_softmax_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 256, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_256_S_qkv_32_causal_alibi_tma_ws_sm90_kernel", 73984, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_256_S_qkv_32_alibi_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 256, 40, 40, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_256_S_qkv_40_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_256_S_qkv_40_alibi_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 256, 48, 48, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_256_S_qkv_48_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_256_S_qkv_48_alibi_tma_ws_sm90}, @@ -1893,8 +1897,8 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 96, 96, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_qkv_96_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_128_S_qkv_96_alibi_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 104, 104, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_qkv_104_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_128_S_qkv_104_alibi_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_128_alibi_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_qkv_128_alibi_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_qkv_128_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 0, false, true, true, true, true, false, false, false, nullptr}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_160_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_160_alibi_tma_ws_sm90}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_192_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_192_alibi_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_160_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_160_alibi_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_192_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_192_alibi_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_qkv_256_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_qkv_256_alibi_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 256, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_256_S_q_paged_kv_32_causal_alibi_tma_ws_sm90_kernel", 73984, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_256_S_q_paged_kv_32_alibi_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 256, 40, 40, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_256_S_q_paged_kv_40_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_256_S_q_paged_kv_40_alibi_tma_ws_sm90}, @@ -1905,8 +1909,8 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 96, 96, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_96_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_96_alibi_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 104, 104, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_104_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_104_alibi_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 2, false, true, true, true, true, false, false, false, nullptr}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_alibi_tma_ws_sm90}, -{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_alibi_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_160_alibi_tma_ws_sm90}, +{ DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_192_alibi_tma_ws_sm90}, { DATA_TYPE_BF16, DATA_TYPE_BF16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_256_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_bf16_64_64_S_q_paged_kv_256_alibi_tma_ws_sm90}, { DATA_TYPE_E4M3, DATA_TYPE_E4M3, 0, 64, 256, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_e4m3_64_256_S_qkv_32_tma_ws_sm90_kernel", 82304, 384, 64, 0, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_e4m3_64_256_S_qkv_32_tma_ws_sm90}, { DATA_TYPE_E4M3, DATA_TYPE_E4M3, 0, 64, 256, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_e4m3_64_256_S_qkv_32_causal_tma_ws_sm90_kernel", 82304, 384, 64, 1, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_e4m3_64_256_S_qkv_32_tma_ws_sm90}, @@ -2049,12 +2053,12 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_causal_tma_ws_sm90_kernel", 164096, 384, 64, 1, 0, false, true, true, true, false, false, false, false, nullptr}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_sliding_or_chunked_causal_tma_ws_sm90_kernel", 164096, 384, 64, 2, 0, false, true, true, true, false, false, false, false, nullptr}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_custom_mask_tma_ws_sm90_kernel", 164096, 384, 64, 3, 0, false, true, true, true, false, false, false, false, nullptr}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_causal_tma_ws_sm90_kernel", 147712, 384, 64, 1, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_sliding_or_chunked_causal_tma_ws_sm90_kernel", 147712, 384, 64, 2, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_custom_mask_tma_ws_sm90_kernel", 147712, 384, 64, 3, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_causal_tma_ws_sm90_kernel", 147712, 384, 64, 1, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_sliding_or_chunked_causal_tma_ws_sm90_kernel", 147712, 384, 64, 2, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_custom_mask_tma_ws_sm90_kernel", 147712, 384, 64, 3, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_256_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_256_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_256_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_256_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_256_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 0, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_256_tma_ws_sm90}, @@ -2100,12 +2104,12 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_causal_tma_ws_sm90_kernel", 164096, 384, 64, 1, 2, false, true, true, true, false, false, false, false, nullptr}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_sliding_or_chunked_causal_tma_ws_sm90_kernel", 164096, 384, 64, 2, 2, false, true, true, true, false, false, false, false, nullptr}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_custom_mask_tma_ws_sm90_kernel", 164096, 384, 64, 3, 2, false, true, true, true, false, false, false, false, nullptr}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_causal_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_sliding_or_chunked_causal_tma_ws_sm90_kernel", 147712, 384, 64, 2, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_custom_mask_tma_ws_sm90_kernel", 147712, 384, 64, 3, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_causal_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_sliding_or_chunked_causal_tma_ws_sm90_kernel", 147712, 384, 64, 2, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_custom_mask_tma_ws_sm90_kernel", 147712, 384, 64, 3, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_256_causal_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_256_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_256_sliding_or_chunked_causal_tma_ws_sm90_kernel", 196864, 384, 64, 2, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_256_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_256_custom_mask_tma_ws_sm90_kernel", 196864, 384, 64, 3, 2, false, true, true, true, false, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_256_tma_ws_sm90}, @@ -2130,8 +2134,8 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 96, 96, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_96_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_96_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 104, 104, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_104_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_104_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_alibi_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_alibi_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 0, false, true, true, true, true, false, false, false, nullptr}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_alibi_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_alibi_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_160_alibi_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_192_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_256_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 0, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_qkv_256_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 256, 32, 32, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_256_S_q_paged_kv_32_causal_alibi_tma_ws_sm90_kernel", 73984, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_256_S_q_paged_kv_32_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 256, 40, 40, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_256_S_q_paged_kv_40_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_256_S_q_paged_kv_40_alibi_tma_ws_sm90}, @@ -2142,8 +2146,8 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 96, 96, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_96_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_96_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 104, 104, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_104_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_104_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 128, 128, 128, 0, 0, 0, kSM_90, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90_cu_cubin, cubin_fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90_cu_cubin_len, "fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_causal_alibi_tma_ws_sm90_kernel", 164096, 384, 64, 1, 2, false, true, true, true, true, false, false, false, nullptr}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_alibi_tma_ws_sm90}, -{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_alibi_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 160, 160, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_160_alibi_tma_ws_sm90}, +{ DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 192, 192, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_causal_alibi_tma_ws_sm90_kernel", 147712, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_192_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 64, 64, 256, 256, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_256_causal_alibi_tma_ws_sm90_kernel", 196864, 384, 64, 1, 2, false, true, true, true, true, false, false, false, run_fmha_v2_flash_attention_fp16_fp32_64_64_S_q_paged_kv_256_alibi_tma_ws_sm90}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 128, 128, 16, 16, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_128_128_S_qkv_16_causal_sm90_kernel_nl_tiled", 16384, 128, 128, 1, 0, false, true, false, false, true, true, false, false, run_fmha_v2_flash_attention_fp16_128_128_S_qkv_16_sm90_nl_tiled}, { DATA_TYPE_FP16, DATA_TYPE_FP16, 0, 128, 128, 16, 16, 0, 0, 0, kSM_90, nullptr, 0, "fmha_v2_flash_attention_fp16_128_128_S_qkv_16_sliding_or_chunked_causal_sm90_kernel_nl_tiled", 16384, 128, 128, 2, 0, false, true, false, false, true, true, false, false, run_fmha_v2_flash_attention_fp16_128_128_S_qkv_16_sm90_nl_tiled}, diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_bf16_128_32_ldgsts_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_bf16_128_32_ldgsts_sm90.cubin.cpp index 6a5bc281d0fb..81208594d0f3 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_bf16_128_32_ldgsts_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_bf16_128_32_ldgsts_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:e31701e0a1f29ac57f2e4c48b52366fa6574d470921089ec9fc471d37b5bcc08 -size 1003178 +oid sha256:d5bb139b12206a563daec9fa473dda422319bde5ae5f965d37cf5ca67d325c49 +size 1005546 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_bf16_128_64_ldgsts_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_bf16_128_64_ldgsts_sm90.cubin.cpp index 0ca1b1c20821..7086ad9f4852 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_bf16_128_64_ldgsts_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_bf16_128_64_ldgsts_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:f5cc3e3ce3d000dc88cec8266e85d4f9fc875d8b4ceccb17796cfc40a1ff226c -size 1063956 +oid sha256:c4357a935656d47414a459939720b66311c67213f450168715e1cb0238653768 +size 1066324 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_16_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_16_sm90.cubin.cpp deleted file mode 100644 index cf69a50762ad..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_16_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:0b3bb19010319e0444524e2dcf739027a24c91b88c641113d20105cc2405c76c -size 926650 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_32_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_32_sm90.cubin.cpp deleted file mode 100644 index 431537bb68c3..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_32_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:0cd4e9a8eaa25e922318e3eb4b1ece0682d2c9c2e2202a35fc7cb7b408aea912 -size 1285796 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_40_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_40_sm90.cubin.cpp deleted file mode 100644 index 3adb44e66bcf..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_40_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:dce9c86932a9a89ded198c51acce01a317719d52fa406dc2b66f4e983d1b02bd -size 1101092 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_48_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_48_sm90.cubin.cpp deleted file mode 100644 index f58eb90158d1..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_48_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d2bbd5ce15707920bdcf093eb57fb5f70462658b3d5f559b0fde43ee90796300 -size 1101092 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_64_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_64_sm90.cubin.cpp deleted file mode 100644 index 0bb93648ee82..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_128_128_S_qkv_64_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4b7a146f40a62e6f98d5343a3d1a654a0df4055f19bf4834fef24a8d8794ff0e -size 1534436 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_softmax_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_softmax_tma_ws_sm90.cubin.cpp index 5b497dde23ee..8331dbce4df7 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_softmax_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_softmax_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:088af0f9eac5d140147835562bdce53304ab1c5da28e1e43689bc857611afb50 -size 700094 +oid sha256:3fff0dfc8b05bdfd41b9f00d65567ff8a96f36e56a75b31e5c48835b7d9c90f6 +size 693780 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_tma_ws_sm90.cubin.cpp index 610a3e03060c..652139d10515 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:a6f5cc3a37a17dedcd18c7ca7dc5ac23fc650c7ad78cd4ba619f62a5b72d79d7 -size 649560 +oid sha256:9fa28c23d82290a782267b18eaa36a545213045d493a72513e3a65305c0fb080 +size 672452 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_72_softmax_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_72_softmax_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 14144f6dc012..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_72_softmax_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:731c1cc24ed554d673ed275219ebf7f4ce8b3bcca0d6680223bbd3d1902c44a4 -size 687462 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_72_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_72_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 10bcabb864fc..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_72_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:68620df2dd0071a06f55a6a8ca0b4004ec544386044f753e0cbd5f8594234199 -size 636140 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_104_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_104_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 5a6e4ba2c52c..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_104_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:112f1f9578a95e2a410350dc1fed1fae6afb9974c4ec1d2b28c04c228ba778bb -size 414363 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_104_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_104_tma_ws_sm90.cubin.cpp deleted file mode 100644 index efe0feb330aa..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_104_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ec74c163ee2573ae8d08a37613b03a495c08ef431a7735c8a2f3870eb11c1a15 -size 1253412 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp index b944cc2450e2..a3c98f01b299 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:d0dcf2a57c63f7673f8e4e880c5e32cc7eedaab4b5bd1cc91a1dd8871b3b1665 -size 417519 +oid sha256:70b101d8936e175391d8051967ff5733a144118ff8793b29b612eac92abc581e +size 423439 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp index afaf3f7091c7..ee0ce3074404 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:8c0f738936d51ad7ace6a754fc15e4073d6003ac33cd8fa56840268cecba5bdb -size 1199762 +oid sha256:26ae7817cbed824212d92c0eb8b25d0f6b9d6281e4d4b6e95e9b6d6d2f5f0faf +size 1236860 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp index 72917f9739d6..e65389452d9b 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:c077ab36aa5f5f4eef96b5cfc451ff4ebda2424fc5d878b8b56919f62578dcb8 -size 1663076 +oid sha256:97dcf2a904ca8ce22f2282644a53986b03f7c0d7948803d2b2b401d6a6dfb5a9 +size 1719120 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 81c3d1eb34b7..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_192x128_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a9ed436452ad0453900569fd6d28c0abe034167107b91a56de8a9d223f485be5 -size 473953 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_72_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_72_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index fc62666be2c9..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_72_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e0fcefd3d955edff214c0b7f166d2dcddb38b18eb1b35c42b023a33b0b0bc72b -size 410413 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_72_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_72_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 614070eafac0..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_72_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9d9f105879646cbd61062987d18f456ff0f07b84947c5ad685c57ca619828652 -size 1243150 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_80_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_80_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index e5fe8735bd03..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_80_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:81cbc5b3140634630e90fb36ce7c95e0ec248ca62f4c4e5725d7f46172ad4394 -size 411203 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_80_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_80_tma_ws_sm90.cubin.cpp deleted file mode 100644 index dc3121d7209c..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_80_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2a05c1c1ef932b5d9b1826f0b27161c930d454ba0e732cece75c39feaa1291a1 -size 1245518 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_96_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_96_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index dcdc8a116a7f..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_96_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a7a3d00acced6a644cf2b1b628b0148f1c7298cde59bc398e7425f4ff9459dcc -size 412781 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_96_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_96_tma_ws_sm90.cubin.cpp deleted file mode 100644 index d7de3ee4cf7f..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_96_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:fe8146d83aee45d6459e39262670429227476a297b889c617f75fb1ee94c6efe -size 1250254 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_104_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_104_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index ee8a28e450ca..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_104_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9b2deeda61234dba168895b7fee211723f27d6523942d498cbe10a7dba39d1dc -size 385933 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_104_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_104_sm90.cubin.cpp deleted file mode 100644 index da0441b8c102..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_104_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d77aa92f15587650a4aaabe619b7cac968dfe2047969179361de209620682d62 -size 857188 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_104_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_104_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 608e5e11e708..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_104_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:0ad52cc57226e4530fe202df9aba3dc36daa7a606c80185cddeb735660776c7f -size 1169730 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp index 70bd1df61403..23274d5f7274 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:4f2f8ccf8cc34cddbc2b13022dbdcb1bff71a4280ecb2008bc47d6a3e46a99c8 -size 389089 +oid sha256:d8a9578f22279c7f83f0126eada9fb14a959e3e841efd641b780be06d5e7ebde +size 375277 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_sm90.cubin.cpp index a4ba144fb212..f8d1e75b2f05 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:dfafc2f1fef681c37f474d7ab0dd90625640ccc2b2a75924ca40a39cfebc5e07 -size 1135824 +oid sha256:e8f883e1814759b4e4e643edb51465f132f27dd77392e9403908cd954eccb19e +size 1137402 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_softcapping_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_softcapping_sm90.cubin.cpp index e0791fa93eed..8cf6386b362d 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_softcapping_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_softcapping_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:bc057942f3706196dce52bd61191e219cae2d7accdb1a84ff7ec92b8972b3eb6 -size 651986 +oid sha256:eb96a6fdcae7f8e19516c4bc4064ccd759906a8b0052e5148fd01e59c37e2f4f +size 652776 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp index c9fbca55b7e6..6f8890117cca 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:97ef8fbe175b0246c3051dd9377800540bc7973728343101da2b1a456d56b320 -size 1140548 +oid sha256:93fb97424b5abb3f807b300bc67bc37f14355831d0ff1ffa2d5d9c0fd872731d +size 1137390 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_tma_ws_sm90.cubin.cpp index b18724e50ffe..7e031d3bf852 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_128_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:515d1f6e5c4eb2c31f0b2b1d3ca1014ffc71626ed114630641022b4f57a6ec37 -size 1554924 +oid sha256:a6803c454338b0a0c548204701ba4411ab55602b42cd2122140b5db09cd19660 +size 1537558 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_160_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_160_sm90.cubin.cpp deleted file mode 100644 index 24b64e480bfc..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_160_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5a93df4d0438a2f30da0c502602c1ad19bf0aac7ff4447f38369dbc9cadbbb5d -size 1004004 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_192_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_192_sm90.cubin.cpp deleted file mode 100644 index 409a84a9f459..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_192_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:11bc483d7ebef0b8a46b2cc2df5f9c8a8fda57a432d5a1932fb5254a85f74df0 -size 1067940 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 3ffa164c38a0..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_192x128_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:46622d087774ebb646bd3fbc168a4eee23d4521fdb3ad207b546847d465fbf38 -size 445523 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_256_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_256_sm90.cubin.cpp deleted file mode 100644 index df6b1982e4a3..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_256_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7ad0f33e0a55b590ca1ca77decdd0407be4b0bbf3d41c1bc50749cc0f88c2bf7 -size 1186340 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_256_softcapping_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_256_softcapping_sm90.cubin.cpp deleted file mode 100644 index 1311db50dbca..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_256_softcapping_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:8cf7bc32edaa83ee0dd2a290b1f1bae15f877b4324e49707a1717f4f476ff52c -size 856424 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_72_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_72_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index acf31b8efb3d..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_72_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d341c3b80c5621797ab29a1a38b79bf5f89f9eb71ce69d37adba5ae5a606a893 -size 381983 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_72_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_72_sm90.cubin.cpp deleted file mode 100644 index abb87e806fbf..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_72_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:cdebd5bdf24c4a8c52f8f2af1ede1c2f7f717412c6cda3b8b3644f72136dc8a4 -size 1037944 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_72_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_72_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 0070fe7008a8..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_72_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4e68d820c0ecf286088ad066b9290e394b099b571bb0d777bbbf83e154aa14b2 -size 1529664 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_80_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_80_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index bbef592ae411..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_80_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:23f9e766eb410f41c76ecded10f19fc43fc6b02bd0ac086fc4c3e4bf813d6d29 -size 382773 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_80_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_80_sm90.cubin.cpp deleted file mode 100644 index d663a008fb6f..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_80_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:450c778eb8fc3ae062bf5346ee22bf840451f38a9b2b6fe540f2cb08a1b6af98 -size 807458 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_80_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_80_tma_ws_sm90.cubin.cpp deleted file mode 100644 index a6af3b1ba17f..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_80_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2850264c0fd83ca5b4d91ed81592f77a3f08424d827a9af7a4821fb4e8512327 -size 1162624 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_96_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_96_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 9938691f1622..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_96_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7a8830df3a06a2dbafd642fd408e40d2f3f1d722f1dfe2a5d5b740c1830b1b76 -size 384351 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_96_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_96_sm90.cubin.cpp deleted file mode 100644 index c871942aacd1..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_96_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2dc471dd95c97bb9d2a90480f5523bcd69e99a4284673fac6e06661a88a0452d -size 830350 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_96_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_96_tma_ws_sm90.cubin.cpp deleted file mode 100644 index cc61db72cc71..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_qkv_96_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:23629c3957daeb633bb0a4eab813bb46b3704619690781acbdd7378671aa8e9a -size 1167360 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_q_paged_kv_64_sm86.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_q_paged_kv_64_sm86.cubin.cpp index 08f4a6c8e365..397d8f56d237 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_q_paged_kv_64_sm86.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_q_paged_kv_64_sm86.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:2b7e8c3474bcc4b0bff206b941e102a0c7514424395ee65b4cd315a69b527cab -size 500863 +oid sha256:8396a30929e67e906ac438e011acdd1eac5e2bd2fa887c2f6ae8aa0f5b6ccda8 +size 514281 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_qkv_128_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_qkv_128_sm89.cubin.cpp index cceb3a68d7e4..18ba9e944906 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_qkv_128_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_qkv_128_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:0e2734b87644eb200d2070ab4ee79bbc0ba95998b0fcfc474c3d471d2a4ecce2 -size 665034 +oid sha256:2c51433d1240dc1d8ab205f89b8cb7f83d93e0224850433610fd95555ecf6222 +size 665822 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_qkv_128_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_qkv_128_sm90.cubin.cpp index 02a1ff8706aa..7ad270f3862f 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_qkv_128_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_qkv_128_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:64686f2a0d54fb592493fc8e6ab7c1e1027f9e5ecf6b0cb88b8d8eb5236113fc -size 683534 +oid sha256:60f4a4656af5bbeb2c8552bf9f9c7cd779586a4cb5cc9f6cbb1e38d8b279226d +size 684322 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_qkv_128_softcapping_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_qkv_128_softcapping_sm90.cubin.cpp index ef0d04327104..2f1dde1db827 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_qkv_128_softcapping_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_32_S_qkv_128_softcapping_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:9a364abc18338e88fc655839a4fc9687b1b60845bfae255ad2676dcc399058ac +oid sha256:61dcb9e691d97658eb41885a1801dc84a2818b7b9939163864c60b2f2f698d01 size 370981 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_160_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_160_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index f76f09226f59..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_160_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:3c5a04c0ac00758408ab1b8cb8f6f949f6a522ed39b47bed6f5678bdbaf11ad1 -size 500399 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_160_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_160_tma_ws_sm90.cubin.cpp deleted file mode 100644 index bd0035fda193..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_160_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d70ee4dce214defe4ce9efe773bac36eddbd171660c497dbfff077e5f7fd4c32 -size 1550992 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_192_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_192_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 25698be3b61c..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_192_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7f0edaad3a70a75ade67c325324d4c0ac55f309156e205fcef08a4c7611f8ab2 -size 500399 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_192_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_192_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 264872229f7f..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_192_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:34486462cb4acca6af183b653b4b9201331fabb6891857bb3b984166cd69a9c6 -size 1559674 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_256_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_256_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index bad6672ed502..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_256_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c8fdce6913e287f1d51657216a504d0f070941806d06386ad0dec166cbde3433 -size 500399 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_256_softcapping_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_256_softcapping_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 73d37e803052..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_256_softcapping_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:766d5759c22eee6b5b9ed4ea0afc90c6ebb1ef663706271214adf1a067202b05 -size 1377362 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_256_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_256_tma_ws_sm90.cubin.cpp deleted file mode 100644 index ee2ce8a9e3b8..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_q_paged_kv_256_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:8cf49186adafa2a5a1e441eff2339eb4d829aaf57d06fcd6203add71b45aaa6a -size 1577040 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_160_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_160_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 3358a83b63d4..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_160_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:923a3091ce8024bb30e2e707e056397aac9f9b24e2d0c8818cc40a3f65895bc4 -size 472759 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_160_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_160_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 99c8093f6cc7..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_160_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:86e26ced3524a0de02487867cfed075c202d8fb08a2e590e1ffdb226ce494457 -size 1422316 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_192_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_192_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index d1dfe9660402..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_192_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:079ccc83022fbaa92f4f7823a190f0805420ddbda63ac8e1d22afddcb1d41806 -size 472759 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_192_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_192_tma_ws_sm90.cubin.cpp deleted file mode 100644 index c9ad41e55d6d..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_192_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5cdbfa248a1ddef45fbebcb93848f369462d4ea43fce7f8d12f725b9a84212bb -size 1431788 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_256_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_256_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 45588bc5e86e..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_256_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2c1b08c4dab9a3165db27d880056ddda08ca6e592082ce76a03f8014a3d2d2c1 -size 473549 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_256_softcapping_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_256_softcapping_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 04ca0edb471e..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_256_softcapping_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d955d554942accb0ceefcbe3ea9e29a1924e258a510d48118a411be4e1c8a108 -size 1311044 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_256_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_256_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 1415d53048f9..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_128_S_qkv_256_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ee9653048ea31c603be31c6daa3b1a45c91994133f8511b055c014e8b8cdfebb -size 1449154 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_kv_128_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_kv_128_tma_ws_sm90.cubin.cpp index b67d89874983..2b9e46c7a071 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_kv_128_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_kv_128_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:e72b520e0778628ed37b71c8f456ee449edd82aa83bfef5ffa4a26c19e3d9229 -size 955032 +oid sha256:d188489645839f22b23f7ab60024a38784246dd3cdebb2860afba4b17e555987 +size 981870 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp index ba25b15cf945..536b3a60f9e9 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:529ff642c151809e38653a82e60a289a8255646da874445d5cec353350b62675 -size 589595 +oid sha256:5bc5c98f5bb68ce8457192a8deb66fd33bd4e18181f6543a80ffee90f9fa889c +size 610511 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp index 39e5fb80584d..9ba28ff3ecff 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:e6caf59252e5158018fc675761bc665a5dd3511284ac01fe3cbe07e42fd76089 -size 1817020 +oid sha256:38facf3787477a775cb81819dd32adc2b14302a6e245ea1bd39a7c79a27f6be1 +size 1922792 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp index 18a10673e3d8..079d5342e286 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:b4823adaab9907bddc44e17da39a8f3ec4388b568172557cbfb3d745275ace3c -size 2409786 +oid sha256:49d610072be65cb35753c025a6e34d297cb8b00763e31f032f8068fd49e82746 +size 2606330 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp index 0acfae14aabc..ece0d7125edb 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:30549e4e351877d091b39480e48d9078e7d6335ea806e34e93b9e0ca51f47ad7 -size 564321 +oid sha256:78b4569d41bffce532654f3b0641599049004acba634be1965685863f4485949 +size 570241 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_sage_64_64_256_output_bf16_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_sage_64_64_256_output_bf16_tma_ws_sm90.cubin.cpp index df4b28eceb51..779c84435700 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_sage_64_64_256_output_bf16_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_sage_64_64_256_output_bf16_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:f5daacedea4e507cdbcd62d25937b413d3c7a2e2fd03dd4781423d8fd44b0b0d -size 674872 +oid sha256:12660d6342b533a1023650fe1c40ed8df1e303878035422e4995697de1abce6b +size 692632 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp index e991c1d980d3..f32216bae9c7 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:99d1c5306300720848580b7c349dc13a71740f7ac757794db1c64b20f45928a0 -size 1761754 +oid sha256:ff17dcd50d76036338dc9f3d009b6b10f5d2b8a338342fef9018dd73a79f1b7a +size 1804378 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_tma_ws_sm90.cubin.cpp index 0ab400146a04..a65367f70722 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_64_256_S_qkv_128_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:6089956bc2085ed1c89d78ece97e879216860ec499125f73f04e74b1fc70a144 -size 2287426 +oid sha256:760cc23fd160128f4be3fd1dd6f6ef4bf18551106404b146b7f374af3fb81c4d +size 2338732 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_kv_32_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_kv_32_sm89.cubin.cpp index acd72c65de0a..e4141dd2d30d 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_kv_32_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_kv_32_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:6173ab315983d8844078fbddd8410ea6b99d30092e5c6dc467fda10300620b74 -size 601111 +oid sha256:de60062494c933226d989901d7fc15d886fd5a84c124f1c01fe583cb45281801 +size 601899 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_kv_64_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_kv_64_sm89.cubin.cpp index 13ae87685feb..8906ad11fe30 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_kv_64_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_kv_64_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:f32d82ae86c521360042b14f1b6a6d79b2bcfe23f6d129af99df591787007dee +oid sha256:367458885389381731b08889460600b9a4e9542cc979a38ad05d6ca3992744b3 size 912898 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_32_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_32_sm89.cubin.cpp index d212a4e8a82b..292e1a9232b8 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_32_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_32_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:f7bf690286a3f532c5375cd76db7383ba552a59f60eba114584e5cde0043834a -size 1385720 +oid sha256:87b40dfd9d1ab2258d7de80a89820e686e87243ab43f7dd20990c871d4202841 +size 1408612 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_40_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_40_sm89.cubin.cpp index 0faf145688b6..c9db86ef9ba0 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_40_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_40_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:f73d1f5e15a69c4455a57a351f856f544b097543991c17c0620917d1e1fd3fad -size 1456760 +oid sha256:ea80c0c776d59d68b5a47ed7ba0fc8e37ea38ab189419519795ca57dd7589304 +size 1475704 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_48_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_48_sm89.cubin.cpp index 490b9a06bd2f..398204974d0f 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_48_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_48_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:e56cb50ecd9aac19bd3af9b65ec3f0e04aef868596dc625939a0e4ad0693ff13 -size 1456760 +oid sha256:b3c7887870f3defa8c2595868c2c8b40afb2ca0b090dc241ad8a34c754857ab4 +size 1475704 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_64_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_64_sm89.cubin.cpp index 6a4052e1b32a..ead5c967592c 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_64_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_q_paged_kv_64_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:1aa3a4f9101c656e57a9053f6f669f36d897e97d29d5c0889b0fa74478a315da -size 1979300 +oid sha256:b797da09627dbf7661ccad3e8b7fd741330f008b3f8e033b7a3c7787a7233e1d +size 2003768 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_32_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_32_sm89.cubin.cpp index a0e6270eccce..4faeb657b982 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_32_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_32_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:1ae2f8df40a25cb8b09f6ce2fb838953e8bbab1ad6fb71a372739d9a8a6636ff -size 1389654 +oid sha256:c55e36802f8679e988ed6fac295314367dd9914c5ff457b7c4c5437ab8b53a41 +size 1391232 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_40_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_40_sm89.cubin.cpp index 6ffcc0b3e14f..85f6542b689d 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_40_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_40_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:c93bb4f2f953d9f0d46139642a87a9955c338cf00d757d95c91d02cf0671e329 +oid sha256:7d9a65aa870c5057349809ae2cc7e03837e37ac3ef2e5633d19e69c444358c96 size 1409386 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_48_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_48_sm89.cubin.cpp index 7816afe19de9..15b05089cf6a 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_48_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_48_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:087062c343a9d04afda590db19761e37a7ad53740f4a1919e86dc439d86e9d37 +oid sha256:76cbfb5a29797bbeb2adad93c0c1e0fd4c1c544a6c12faa2a825cdb4eff1dff2 size 1409386 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_64_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_64_sm89.cubin.cpp index b0727995ba2f..ea60da2843bb 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_64_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_128_128_S_qkv_64_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:9d0e082555cbda07638de0d1d838269437f7100e6f12afd98c3a3dc378d2aa7c -size 1948502 +oid sha256:61c16947041287198b160091a89f1677ebe7babed9c9da6f6625436f7b526a6f +size 1946134 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_kv_128_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_kv_128_sm89.cubin.cpp index b3a1253af760..bccbb4b8d850 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_kv_128_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_kv_128_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:5c46353a6c00c154ed5d7bbb52c56b42f8dccf5a700f928243029ccfafee3013 -size 308265 +oid sha256:f1114bbd784a3ea000d86f00e35086435d50c430ed695448a306cfc4bd54f60c +size 309055 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_kv_72_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_kv_72_sm89.cubin.cpp index 969696cebbe7..4d09371f99ef 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_kv_72_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_kv_72_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:f4f0d5736d6801f3614c72f31581c1e227cf51eafb60e009b47f267982f36136 -size 292477 +oid sha256:3c8905ae4aafc41cce6557456bdf08d7ae6eb5a93286ccbf5d0b745fb33cd298 +size 293267 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_104_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_104_sm89.cubin.cpp index 93ce38445bef..41214fa51ddc 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_104_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_104_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:9d1c4f9a5c53d3f226dda0c2f1dd53afac4f3719731130af6a9ce704e9b55d0e -size 515083 +oid sha256:e373ec7eb583a0803821145ec16f2ecf1a173c70f0796207750e51b97c72d604 +size 528501 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_128_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_128_sm89.cubin.cpp index 132492c05c43..a946012b6b52 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_128_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_128_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:8662ebc259db8989f193c69e1aea9bc2de7da97d8f0564ca023d77123cfc05d8 -size 679266 +oid sha256:2805c97b33142d036c8fc510d603e5c0d6d74174ae1f15b04feeedf44f0b5ab6 +size 702156 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_160_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_160_sm89.cubin.cpp index 7d509ef97a23..ce6524aa572f 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_160_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_160_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:33c76fd50a8a68c154e3c5016767f1deef66b9b369885fce6fe5da1ecabe83b5 -size 742412 +oid sha256:111f7cebf93583b831e5714ab597ef6cf9afe9a215a5a9bb1cedf04176f4129b +size 761356 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_192_output_bf16_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_192_output_bf16_sm89.cubin.cpp index 2dcf6621af63..7e03d88b7e6b 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_192_output_bf16_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_192_output_bf16_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:69eef116cc9ceeb142af8d83bf9463fd1678539ac11915712be7b7123f71aed8 -size 782692 +oid sha256:9b44d7f8e5db9b0fd8ccdd905124faf5a703c89c6de326367ba200697fb518fa +size 806372 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_192_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_192_sm89.cubin.cpp index cd3846383cdc..053f856fb3e8 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_192_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_192_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:80da78fcf36253cfa63bc5cd7891cf4f79ed32ade50c3bf4c6ab209abb77cf46 -size 780300 +oid sha256:664ed6e91ccd091fb4733b55a2799d4562df876ef4e3be8ca79e6d0b55bace4a +size 803980 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_256_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_256_sm89.cubin.cpp index 8dfa8144b480..ec8103b8a16b 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_256_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_256_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:798951dbc53219e7402642bd6b49a5eeb01010ff76a0ab8ae99f519effc86080 -size 980002 +oid sha256:98431cb031d4d41035fd7a5a253fbf4b23214ba9e8689749ad23de925d97b0eb +size 999734 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_72_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_72_sm89.cubin.cpp index 33172350e7ba..ebaa17c5c62a 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_72_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_72_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:69aef72f514c7338449e301205aca1a411ed466f90801410547d241f2147f339 -size 507977 +oid sha256:48ab14dd4c3e988db85530381833b1753fc8579a8716df1a81799d122ecc19cd +size 520607 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_80_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_80_sm89.cubin.cpp index be3e06ee6bc2..fe3765594ae8 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_80_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_80_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:737387664ae52b4874af7971c93f70942f17a559dd68dac553b59be682183d60 -size 507977 +oid sha256:a4aa5c1c533f5ce60a50110a6bbfa2af6cd7a0488776cb1fd491ce594b0f94f4 +size 520607 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_96_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_96_sm89.cubin.cpp index 73a65400cdc5..69da730357cd 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_96_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_q_paged_kv_96_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:23785e6d85c7a93d7a0f8691d79a6de1c953fbb4ee057cb8ac13a10c0b1ed6d6 -size 517449 +oid sha256:b0dae8957de096f310cfe6bb977babbe745e7542072920a454a60b9ad05c4318 +size 530867 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_104_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_104_sm89.cubin.cpp index 09e8012c4e32..29a11c7b0bea 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_104_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_104_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:ffefd85f6395becfe5b80d863761617fea35167138b738d924718efcb1736f49 -size 499283 +oid sha256:849c37d9f772de883d6fa358161f977216d48932ef8a27cec2cfe931c9880e06 +size 500861 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_128_sage_64_32_32_output_bf16_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_128_sage_64_32_32_output_bf16_sm89.cubin.cpp index 7bcf78afdc03..b1e2e33414a1 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_128_sage_64_32_32_output_bf16_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_128_sage_64_32_32_output_bf16_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:346b1557eee6957ed0cf3b793c86b78dbcaa799bc806798f15c28eaf6581e110 +oid sha256:189df2e89d79e1969521dcb124bcd71f274493e369b2809fc5ed552e8be1977b size 184391 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_128_sage_64_32_32_output_fp16_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_128_sage_64_32_32_output_fp16_sm89.cubin.cpp index b054bd5be480..76ed2ade986d 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_128_sage_64_32_32_output_fp16_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_128_sage_64_32_32_output_fp16_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:fec694c26cdda7b808b836a7b18918b56eca406c0d42108cec6c60c31d882209 +oid sha256:43ae547cc799f0c688c19daee4bf357d6d2fe2c06d894bcded7ac40e699caced size 184391 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_128_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_128_sm89.cubin.cpp index f150e37b946d..344fd446267f 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_128_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_128_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:039256731f20528aab02a8df3729680d8cc9c9bb03b89047724b58c185d65f74 -size 665832 +oid sha256:39c941a13e14d0cbfcd19e1d11f75047227aaf992d60b56e45f063f92ff80cc8 +size 667412 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_160_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_160_sm89.cubin.cpp index 04fa0c92a53b..50293ac4e5a7 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_160_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_160_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:a4bad8fa30b04f0f3a13edc310a6b9eb6e99ca31cad75a15410e233327babdbd -size 674516 +oid sha256:868ce05564bbf9e23a3f6562bd75d537d1c5e901eeb0bbecb24261bcc7d23370 +size 676094 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_192_output_bf16_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_192_output_bf16_sm89.cubin.cpp index 275115d4f86c..7f2a34961d2f 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_192_output_bf16_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_192_output_bf16_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:001374158c745bc46dec1996a7d1ba0a3b537c8c354ecd6938e5ef9d93339bcc -size 725056 +oid sha256:66d791187f871dc70a6b90cd9d60dc3db06d60c2beaefb3d75c2ff1f949d5458 +size 726636 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_192_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_192_sm89.cubin.cpp index 33eabb64f7c7..13085d8c6674 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_192_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_192_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:4bd5818a16a40b85edb46f08b23b78adcaf3dac0defcc86000fcf0589a6874f1 -size 722664 +oid sha256:6a065d8c65f022875bb49bdc9aa853061149ff2cdfcaf1f8cdf8a3efe456e8a5 +size 723454 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_256_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_256_sm89.cubin.cpp index ec22b91087cb..b5ec7f76b485 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_256_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_256_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:ed8dbc734d33ec27051eac487109d50ef8c63edb6471b4f8b0fd403d807bc173 +oid sha256:212ffad34a9b3002c1ab7e590bbadf1c94cb9847acbb479c311e9057c4e4c44b size 932628 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_72_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_72_sm89.cubin.cpp index d721dfe53b5d..2099dc866529 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_72_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_72_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:b22e753cfbcf3314884fc4557c973d6cf2486cef891f0ed74a680a3e34ffac20 -size 638204 +oid sha256:e70aa7f7c6f8e41c5f142fd268a88fd0390f59ac9aad56b8be062a05f8f49ff8 +size 638994 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_80_sage_64_32_32_output_bf16_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_80_sage_64_32_32_output_bf16_sm89.cubin.cpp index 7d20f6338647..b43312dbda29 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_80_sage_64_32_32_output_bf16_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_80_sage_64_32_32_output_bf16_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:8797953ca8e515e35a955de5e7a173dd2f83be3c807844fb4c4f04128c4840b8 -size 161497 +oid sha256:d0cc18b1e3835a7cc42648d1bd0b63507020427299027667f9dd4faef37450ab +size 169391 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_80_sage_64_32_32_output_fp16_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_80_sage_64_32_32_output_fp16_sm89.cubin.cpp index 6b020e27aab0..bb9d123faddc 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_80_sage_64_32_32_output_fp16_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_80_sage_64_32_32_output_fp16_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:65cf71ff8b657165ff727d1bd90266042fcf1c31e0882953415d9f66e14b8eb3 -size 161497 +oid sha256:90e97d06799b33f0f4ed6c68aa43616f4f2e013680909ca56d2e514a4481f0cf +size 169391 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_80_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_80_sm89.cubin.cpp index 1664e4edd238..8e7857f9ec20 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_80_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_80_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:bc72689d27d04bbff63953c8772069ffde934aac9017fb22be9b27f056fa826d -size 488229 +oid sha256:c48f3c39368e774c4f3c281b7422e0b90e08321fa29591882c7071a635e1c3c6 +size 489019 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_96_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_96_sm89.cubin.cpp index 79fef537b3ca..686a996434f1 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_96_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_32_S_qkv_96_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:960e14c1154c414028e1eec2b88258cd5d6d4db05ad0905836eb59527f0bc7dc -size 500859 +oid sha256:b5edbd9d472583367857e998d65097561a9b36bc68ba1ae94f3b79940c7cb6f3 +size 501649 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_192x128_output_bf16_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_192x128_output_bf16_sm89.cubin.cpp index a70af8524466..dc1b346d2316 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_192x128_output_bf16_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_192x128_output_bf16_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:30f39bd5e745d016a62d93b5bff3b86eba92b91a8391579dac8e9ff3f43b4c89 -size 232533 +oid sha256:9eeb56a178049dbe0869030e20eeb608423fd5e34e3720230e5ed4373717b91a +size 238849 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_192x128_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_192x128_sm89.cubin.cpp index 53245fb936fd..c0b56e6cf06f 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_192x128_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_192x128_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:0a7c5b8d27d0e3470bf7a5600722e8c9cb977802746ce529b9224b2aaf197c40 -size 231721 +oid sha256:00c69c0bfcb04dcd381677913781984ffafa3980922807faa94f125c01d7b901 +size 238035 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_576x512_output_bf16_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_576x512_output_bf16_sm89.cubin.cpp index ed02d1dae9bf..d8dde7184afb 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_576x512_output_bf16_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_576x512_output_bf16_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:1b67c08eebf9ac037c3c0ca6f8cd86c2c66760db4ab48e714e44276e10d4f0cd -size 288577 +oid sha256:cade6eee7a6be594da0a65e270954a11af436082b02bdd036aeddf9486812996 +size 298837 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_576x512_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_576x512_sm89.cubin.cpp index 61eccf02eba0..394e497b7591 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_576x512_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_q_paged_kv_576x512_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:beb4939e0f07e964f53db3bc7f051e124a89d684caacbf53b4d882049c979541 -size 287763 +oid sha256:470b274928968dc99c7cc1299cb906a9c38c2e5ddb556591047677e8b968b2c9 +size 298025 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_qkv_192x128_output_bf16_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_qkv_192x128_output_bf16_sm89.cubin.cpp index aead6698731b..c4a5aff2bd72 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_qkv_192x128_output_bf16_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_qkv_192x128_output_bf16_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:66dcf4cefafc80111d5c517466d3be1b96fdef31975a7fbd0afbe903b90e8694 +oid sha256:6d9c45c07e5f4513fa4666178709a7051042e1fa791d0ddfe9540802ddf36194 size 231731 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_qkv_192x128_sm89.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_qkv_192x128_sm89.cubin.cpp index fc9ed96b2b91..6ba4c09f1efc 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_qkv_192x128_sm89.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_e4m3_fp32_64_64_S_qkv_192x128_sm89.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:341f1667912db3b3cb2f5b98e41c9f41d5458e47c3d0cfd056a4191a81f550ae +oid sha256:682a0bc5821e74d56736641ecd8a7ccb1a7d7352183eda62a56edaa280d99004 size 230917 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_q_paged_kv_64_sm80.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_q_paged_kv_64_sm80.cubin.cpp index fc73ed78374b..8fd17c8d5bbd 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_q_paged_kv_64_sm80.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_q_paged_kv_64_sm80.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:506ac0837ad02e0f474df7005ecd6007834bcbd95d51b8f367ff4982eaa1f6d3 -size 1583834 +oid sha256:2dbba9a30ed262e3096c4e7d7c3e4fdadd3e073e41894e8258de9274e08979d7 +size 1615406 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_16_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_16_sm90.cubin.cpp deleted file mode 100644 index ce86916034f1..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_16_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:6c223dc94354ca23a35b7b4b5a3b6db3148f6bfedc3c2ebbba64116afd80c893 -size 957434 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_32_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_32_sm90.cubin.cpp deleted file mode 100644 index f6f5ccd922cd..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_32_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b6c0be8d476acc18c75a5ded0ed86488606343e37c0819946151f1a0a2cabb72 -size 1300004 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_40_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_40_sm90.cubin.cpp deleted file mode 100644 index 13de4bdfb40e..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_40_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9820b68a7a52187391827e6050cb3aa7d00789523e15a1d6aa67213dcebd8141 -size 1102672 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_48_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_48_sm90.cubin.cpp deleted file mode 100644 index a4c26c46d2e9..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_48_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:bed56fc61e8d6137c68843fc8cc81619eecbb9f18a15608121ea40357a9d07d2 -size 1102672 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_64_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_64_sm90.cubin.cpp deleted file mode 100644 index 90224750ef17..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_128_128_S_qkv_64_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4de4517b69e8db6f9fd570eebc612d93c37156c9c03ca75ac0fbf76b723af5e1 -size 1454714 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_kv_128_softmax_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_kv_128_softmax_tma_ws_sm90.cubin.cpp index ea8efec4677d..b9e28a17c540 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_kv_128_softmax_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_kv_128_softmax_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:68093a7692e95151323982878c48703677b3fbd1f46490d95e00718f79f41c8c -size 731668 +oid sha256:dbd51135c48812f21f53811b57057cabbef6c7a8a7833c411d8f8c47a2285c65 +size 724564 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_kv_128_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_kv_128_tma_ws_sm90.cubin.cpp index 3dac1049d58e..7a93dfaa65c2 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_kv_128_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_kv_128_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:85a12532f106fdd7ba32a5f5e4f82ac7cde4fd4e4634a3f4c26ed2015d0feca3 -size 678766 +oid sha256:c9ca2010bc714808c4e62ad7a66ae070e18bd40f678f46663b5f46d964283e6c +size 704814 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_kv_72_softmax_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_kv_72_softmax_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 9d819d50c7f4..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_kv_72_softmax_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5136cfd28704b70803682f0f2136f9142b4ef232abe0811a736d47a6104d2ff9 -size 725350 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_kv_72_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_kv_72_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 7d5011d919a2..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_kv_72_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:dffe7d4f5738972b3324ab2accc3fbc60629ccce5af7539e027f7bcb3b6eb379 -size 671660 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_104_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_104_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index d021de623396..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_104_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:166c465c2a33088be987261fbbdea6c9bed80e167d2599c800ee5fbe9288623f -size 445147 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_104_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_104_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 7b91ddb310d5..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_104_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:1baa67f5338401a3deb91c06932ef2a6c14c57dd0bf13a01a547655dae36a46f -size 1308666 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp index b6cb9d74bc7c..a16884caed3a 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:b5d1bccb654c37a5912c92af0fcee51d0c48d0e7a79ecb23694b033c819a034c -size 446725 +oid sha256:aff65d92093547c644da83b9800c8d8393f1a9d530f809b6bb35138afbe669c8 +size 454223 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_sm80.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_sm80.cubin.cpp index c0fb3f904c4c..91712bb82ca4 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_sm80.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_sm80.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:746a02b69b59a23700401b3269da63a7c39e1d4f551eb0440a2d0de155c9430f -size 1339930 +oid sha256:3242c721b07ab2f56698b11c16f2766b61f1a27c8c30e9458e5179a71340cf76 +size 1377818 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp index 43c704676ca3..5d684d6316e3 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:88cb677a4f6f1e0dbdd67a53e66438f66ef94c1069c03189e132ca18b00235ad -size 1218706 +oid sha256:cd323cec032400ab6c820d02d9e1c6da22ad0b627a0bf6bf51de0c0ab4aad99c +size 1260540 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp index fbf197218b48..138e82ec0c48 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:00bb96ced6c3120c6012c0a3148f6acb19e7c9902c95340ddbc19df26502a45a -size 1728592 +oid sha256:3adf59ee5801afeed6c1a51c6ca6bf504e534c3c277dd58c91d1818e13c726be +size 1790160 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_72_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_72_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 6b0625c2df7c..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_72_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:59a139b34f9cd01be2adfaea903224755ec32f9a6c220afe553e96f107d53905 -size 443565 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_72_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_72_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 3166df93c69e..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_72_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9052273b78c6e1683cc27ab2a38366c2e430ba2f39ba9915359c3551d0c20b4a -size 1303928 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_80_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_80_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 005a6460cf3e..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_80_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7ae4fbf01b3b00e9e5c69515200048c4b263a877ac3f015b802c363c61b11452 -size 444355 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_80_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_80_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 06e37faff0c8..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_80_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:aa9d3415300b1940f6d78cfd10d45e2f041f215fd22d9cf9732167bdfa24cd96 -size 1305506 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_96_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_96_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index bbef6fb47e4e..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_96_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:607bda6ef568706aee7d7d2d74d02755cd388189f6b01b6223296adbe6964cb0 -size 445145 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_96_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_96_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 93ae415f316b..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_q_paged_kv_96_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:0f1d5702c25c2b4efde52ab1a786425c80b722876d6a50814467475a9811c6bf -size 1307874 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_104_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_104_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 1d076d17157c..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_104_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:04d0d28c881b763046b8c545561b0181c2223b41f145937febfd02a383335b45 -size 429345 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_104_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_104_sm90.cubin.cpp deleted file mode 100644 index ed67845d837e..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_104_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c829eaa0218016c75e572dec7c747b9edfd3649c169ea999d925565ec8f28352 -size 836666 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_104_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_104_tma_ws_sm90.cubin.cpp deleted file mode 100644 index cd71fa129055..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_104_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:0bb919a31dd552d8d07cdc9be071c05302fa570f4680832112f3ba802a52e588 -size 1232876 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp index 22a173a7b7eb..481792268b5c 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:a3696f0ecc3413faea1c7017f9f0c793a048c5b19d342a9f8e22f147f5a27a34 -size 430925 +oid sha256:e17333a518382c1d0980c8c8c4500df358846c602db5f7f2c413f135f3ff263e +size 416321 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_sm90.cubin.cpp index 0191d44e8b95..62e54f7ecc4e 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:3e80dad3e93753dd6bdc463d7f5f490dfde9c864db3f2dbcef26bcd4aeef7440 -size 1107408 +oid sha256:5654ec576d9e76bec93bbc11dfc7142bf4e57d1bc718e8c76e1b8a9c9dced0dc +size 1108986 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_softcapping_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_softcapping_sm90.cubin.cpp index 2c9f708cce19..b485cdcf2eee 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_softcapping_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_softcapping_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:44c08cde104f5fbb7b6afc1f31ea124b60ce248286eb172f1abe278bc1206823 -size 632252 +oid sha256:09f3e9c7de20a1fd78f68d32b4be0301a8426ea8b61c90a361968e143a409dee +size 633042 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp index a76b694dd8d6..84b753442af6 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:b9d77cac38f219b69b29f9c2050a98298ee9c1b436ab1c2c77179a52fb6b4ae6 -size 1161070 +oid sha256:22a85bd4725e2ca09a3f45519b9abd3d353f5de8cb5994f40213f5dca233e0ad +size 1162650 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_tma_ws_sm90.cubin.cpp index 57587463a857..0445af1cfa4a 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_128_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:d2243fd0f40e2b69906ac81f5f07986109c48d9b193c8a4b25af1013e235b140 -size 1633068 +oid sha256:c373d9294f2adc0601433f57e1369eef8ec03a6fc0c0a514b5338ed313e6a6e2 +size 1620438 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_160_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_160_sm90.cubin.cpp deleted file mode 100644 index e1f73fa4f095..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_160_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7c3b69bf7b3375b0bc7d02a44a7c819df352bf79a54ed043ccbd63aaf39045f0 -size 964538 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_192_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_192_sm90.cubin.cpp deleted file mode 100644 index 41d039a1f2a3..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_192_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:898f60eef263a833f82713f3cbfc35de7cb7c4a379f860672089d7f22cbb5aee -size 1011108 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_256_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_256_sm90.cubin.cpp deleted file mode 100644 index 6a36d042529c..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_256_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e17eaebf1bf5aed3844436a7fb66e621398cf29086e0827a267cd995d92ebd01 -size 1061626 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_256_softcapping_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_256_softcapping_sm90.cubin.cpp deleted file mode 100644 index ca1c147945d1..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_256_softcapping_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5cd7f4691e5630e8ece756982dee21d822e2b12298141e41a258c2af3e64119e -size 774332 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_72_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_72_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 69e8a2563887..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_72_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:eb6285b8d8105f3622f48cb86c033b35bfa1ff5ea1c90a84a58f779212b0d5cd -size 426975 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_72_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_72_sm90.cubin.cpp deleted file mode 100644 index f21688f121b9..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_72_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:948056dc6b22f82ecf30c2884dd37c44b779c28a6e73292f614a8710446c2458 -size 1028472 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_72_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_72_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 6396b083006a..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_72_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:75c6414ae1e6e1d8f93e9ec0d0287070a4129752ff0c26649bbee24f372a0375 -size 1620436 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_80_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_80_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 5436b237a2f7..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_80_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e2c564cb41d28cb43f60e239f53e58042e958648c86f511f038ffaf1e6cdca10 -size 427765 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_80_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_80_sm90.cubin.cpp deleted file mode 100644 index c9949c867703..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_80_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a920a2e35442a9d1b8542ebb79224d155eba14801c249013c97c533424be549f -size 797986 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_80_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_80_tma_ws_sm90.cubin.cpp deleted file mode 100644 index e241bcaf72a9..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_80_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:27d74856fb9a4c77a6cb4d3049d5a008edce9f16bb1f9feaa17ed69dea0618f3 -size 1228928 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_96_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_96_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 9e28fd65eb34..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_96_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c034724a933f1a5c9a6e4a8b5036666145fbfd05b8e92f59c58d7d8b145d21e8 -size 428555 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_96_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_96_sm90.cubin.cpp deleted file mode 100644 index fd3666f80465..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_96_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:59bca9ea361c94ec5515bcf4430e260374fdeb5eb8092893b4af57d832b57e77 -size 817720 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_96_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_96_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 7b988cd4030b..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_128_S_qkv_96_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:6d1cfb99fc175ab75e1fa312988b1f32a941cae7efcf88b9eeff0a5b3a0ea6c2 -size 1231296 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_32_S_qkv_128_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_32_S_qkv_128_sm90.cubin.cpp index b91767d0f768..81125e7086ef 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_32_S_qkv_128_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_32_S_qkv_128_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:daab8ced44f0d93a883bb02992718e70f9ccd0ce2a449caf7f9993d1f8d31aba -size 608545 +oid sha256:c70a136dfd55771b4218b60536d034f6dbcf285353ce8ea75c8fc93d33d09450 +size 609335 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_32_S_qkv_128_softcapping_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_32_S_qkv_128_softcapping_sm90.cubin.cpp index 4c466d2d8b3e..8e7059ad2bd8 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_32_S_qkv_128_softcapping_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_64_32_S_qkv_128_softcapping_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:fd36f3da8fbdefa334ef098dcd66b4448ab3fecbe245d94dcaa0a28e435abbe7 -size 332303 +oid sha256:0af8defec56bebfe634eafe3825626e91301937a1beafd5e2cb61d28e18e86dd +size 333093 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_16_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_16_sm90.cubin.cpp deleted file mode 100644 index c0a612b201e3..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_16_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:864b93f2f5b39c858a747390bd11230ba988a4cd22694ca545584760f067a0b2 -size 928238 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_32_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_32_sm90.cubin.cpp deleted file mode 100644 index 9496b7405544..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_32_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:322a29b9b01f4707bdb85d4aea462f6ccd5e986d597eda2d1d686f239585dabe -size 1288174 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_40_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_40_sm90.cubin.cpp deleted file mode 100644 index 1994a04d107a..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_40_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:0b0edb593c51d3123623c83a434d572c864f36bba488a92f0cf580cb02ef4f9c -size 1101892 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_48_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_48_sm90.cubin.cpp deleted file mode 100644 index a993550a3b74..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_48_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:fd5b26698724cf28a93c0e599b7d94c4edd5dfce135148ac04f4a72da7bcb75b -size 1101892 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_64_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_64_sm90.cubin.cpp deleted file mode 100644 index 6ffff18c1966..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_128_128_S_qkv_64_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:f9ac7a8026dfbbb20916d4a3833969e537abb017bf01f74437c7b7cec7ef43d7 -size 1536814 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_kv_128_softmax_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_kv_128_softmax_tma_ws_sm90.cubin.cpp index 3e19ec15864c..813ec5559ea7 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_kv_128_softmax_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_kv_128_softmax_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:cc37c82a5da895cdea5cf64cdf53e7c2111e9baa5520faa6a0862452cb725bdd -size 701682 +oid sha256:9e05e42418d14593b3d990875c8d813441176118804a2b6d79bc19c420ad176d +size 695368 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_kv_128_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_kv_128_tma_ws_sm90.cubin.cpp index ecfd32234db3..131f4659278c 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_kv_128_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_kv_128_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:d4407bbdc5e828d0fdee274d220835fedd95a1df0de5f03eb25c565d77475a11 -size 651150 +oid sha256:3eee694dc657713c85cd5daefb80742ec9789cf01846683d490ecc237863aeda +size 674040 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_kv_72_softmax_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_kv_72_softmax_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 908a6703979b..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_kv_72_softmax_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:56f64ad3e1e105681ff0bcb36ecb975e0c2272c5498e2e4e28a2c974f50e1bbe -size 689840 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_kv_72_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_kv_72_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 1550dde50a0a..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_kv_72_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c6b9c5df24126dc6379d494d4f3c0c111745b4991807d7832b7e07c6fabb6f30 -size 637728 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_104_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_104_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 6226838bd2c7..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_104_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a9404432f9369126cb46f895f58583ec513353401a862f4c839e1cd32a455263 -size 415161 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_104_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_104_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 4775e85371d2..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_104_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4ffc6ccc2a3aa754a835062567c29b6c65030513e089d8e73f52a2d6f13093ca -size 1255002 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp index f75f8face10c..61f3af8c375c 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:9ecb41327499b0afec6ed95c51ea525ff24faecbfb6dbb1bb9306963c63c1024 -size 418319 +oid sha256:8baad0ecf9c9f2afcff799f063c24c3d1475f45f4097977bacdfea37fd9fc6db +size 424239 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp index e38d2fce5bcf..ef55d9b350f4 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:ef96357c675bae747ea535ea9db16f091e5244e11da565ff37153b57639d170c -size 1201350 +oid sha256:693859c24beb3519f369aa92d5b3097fa7323b5f9e911dd508c029f0289bef17 +size 1238450 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp index 9b1c99cf4775..5644a54c5b58 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:4e86ed4f5441192399a918e0c935a8026b87074f9ec85e0851d7131477e96ebe -size 1666244 +oid sha256:5e4ae887df4aaa7f402cc3fc9e44bff89b4211d6b9ad8875a99e44362e188557 +size 1722286 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_72_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_72_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index e9f876edc915..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_72_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2692bb1d337ec37478f1e03d202df0708fd1caef562a6b3a6ce47983bb76e2b6 -size 412003 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_72_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_72_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 8730787928c6..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_72_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:663daceb926f75f3ea35fd3b59e4bcc55ec607cd010655cd93262a4f989548fe -size 1245528 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_80_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_80_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index f79fb129a327..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_80_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e48da9c2af467db9a313f0bb181d7c89e194d8bd7019cccb3cf99d69872f528f -size 412791 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_80_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_80_tma_ws_sm90.cubin.cpp deleted file mode 100644 index e135e15beccd..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_80_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a074b86c12b02ecd7965354a257de8bf04582c26d1a33a46751c0da8d421f057 -size 1247896 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_96_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_96_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 31f3e2fdbd1e..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_96_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:92b6659dacee2367a4667b24922c32f79803fdc6330eff8b1620484261fa9b95 -size 414371 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_96_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_96_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 561b767b54ed..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_q_paged_kv_96_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d70c8d0fc4cb758a6b3bfd4a6d52dc130926cd9b86e6040ada69d65eaa9dd08f -size 1252632 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_104_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_104_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 662adb4773c3..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_104_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b9f21028fd1d004f6ac939e26260629969a44ef54a26e6b66835fc058262402e -size 386731 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_104_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_104_sm90.cubin.cpp deleted file mode 100644 index 9394650f1b0c..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_104_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:da5d152d9ff0b395026ac63e410b97a5dc21bdbe9903fed79c239b4069e32c9b -size 858778 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_104_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_104_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 65c19702664d..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_104_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:0a8b5dfed70618d873005a39a1a8decdbee84c3cc1e3a1a7bf5868d3b758091c -size 1172108 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp index a84c5b9ef5cd..755f0195b6cd 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_alibi_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:283a5d2aa8c629cf11339cf9bf5590c9c1bbe90d31f7a36f333d85759881b4ad -size 389889 +oid sha256:97d53942b6dd1ad8bd7596ffba97f79b5f9c932beb5553a22d7aeaa1f16299f9 +size 376865 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_sm90.cubin.cpp index 4e697362cbe6..f03bac6ad1dc 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:a4bb3112c04d162f34d2f4aeb48d42d90dd6140b03f3440a734c1ca8de95e1ef -size 1138202 +oid sha256:eaf758af72cf17bca3eca50fa0062fe64a354297bc02a4948226e33bbdcb5bb2 +size 1139780 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_softcapping_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_softcapping_sm90.cubin.cpp index 8eb54ceb8e40..17236357122e 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_softcapping_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_softcapping_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:f1a314c7873595f44f8abb24d131b734e22588123d094ff75d58bc500a55b8f7 -size 652786 +oid sha256:13ac9af1a09a4c5ff6eddd9565840aaac11e6072dac3c7a1bb5377705b5d120b +size 653574 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp index 508ea21ce318..55070baa1fb2 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_softcapping_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:4ff78b87d21a504895d0408aad3e10cbb0c2a6006e171bee327ec9a7330b49d6 -size 1142136 +oid sha256:c35488ad990365bc5f50b7b2bfad2572f48ee9060345435e817384d41b4f3b13 +size 1138980 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_tma_ws_sm90.cubin.cpp index c1be56992e5d..1ca06ff0c635 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_tma_ws_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_128_tma_ws_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:d19a450c92c0fa54efc60d5216009a4e0ded9aa67002da37c4f8cd6a33d3e527 -size 1558092 +oid sha256:f0be66ba8c48682577dee9a7a75a5fdd9e363332881a6400c643a38d7dea16ca +size 1539936 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_160_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_160_sm90.cubin.cpp deleted file mode 100644 index b68db813ea06..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_160_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:6487499516c326de9764184082cc310734ab21c1e7f6575636b87eb47c7948fb -size 1004804 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_192_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_192_sm90.cubin.cpp deleted file mode 100644 index 1c5f58b5c376..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_192_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:6f341d23a7b31258e2c9cc5ee8ec1efee8f8ce3ec692d0bc85ba75b0f0e18255 -size 1069530 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_256_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_256_sm90.cubin.cpp deleted file mode 100644 index 8978e7308056..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_256_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:39402fc4921b25f7cc686503b99e548320d5261c152a4da53f2bbe9ff822a7e8 -size 1187930 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_256_softcapping_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_256_softcapping_sm90.cubin.cpp deleted file mode 100644 index 7fbd1d530944..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_256_softcapping_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2488adef67c304e1683f9fca3764ca9349ec30a5f40aa271beb9f3ef906aafb4 -size 857222 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_72_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_72_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 48227580b73f..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_72_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:bc6995f32954da3b8ec44f6b0dfbbd6e628f8f2a53e4637c67c1154b9ec0141f -size 383573 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_72_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_72_sm90.cubin.cpp deleted file mode 100644 index 3fd7d0074b86..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_72_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5188b96258ba3f64eff9e76c6ba123db82f51364a41c69ea18be86b97d4ca58c -size 1039532 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_72_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_72_tma_ws_sm90.cubin.cpp deleted file mode 100644 index ab8b03996b93..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_72_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4863884f64d4dd3d58605afe174ed735e99d69623d5a6556d67d3601e469815b -size 1532042 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_80_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_80_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index b4efd858c894..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_80_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:fa4972a0f2d79a52a0ca9f3433746d1d45aa978cab2e2ecccb6a9d804186ab4c -size 384361 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_80_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_80_sm90.cubin.cpp deleted file mode 100644 index 3d86a698f216..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_80_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:6bc8d4e72f22014a3b43fcae4819b1a77913acd18a6837554ed291906db4c0a1 -size 809048 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_80_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_80_tma_ws_sm90.cubin.cpp deleted file mode 100644 index bc53dc7278e9..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_80_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:8de71b3a330e32573d7644aef5e32dabf9bddd955e5a377b28754655a52078af -size 1164212 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_96_alibi_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_96_alibi_tma_ws_sm90.cubin.cpp deleted file mode 100644 index 7c272c77d038..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_96_alibi_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:979f4cec391f415c87333ad950ff4ae5e90b464c20b91902688d22956c98216b -size 385941 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_96_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_96_sm90.cubin.cpp deleted file mode 100644 index 555bf7292dff..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_96_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:fde1dec4746bca09ef1fcf986ac069de2bce86079fbefa7caee845887d788c98 -size 831938 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_96_tma_ws_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_96_tma_ws_sm90.cubin.cpp deleted file mode 100644 index d8cb87b2eac6..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_128_S_qkv_96_tma_ws_sm90.cubin.cpp +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:34229a727983a774ac1acddeecb051760d7431b02857deda6ff52eaf8e75787a -size 1168948 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_32_S_qkv_128_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_32_S_qkv_128_sm90.cubin.cpp index 40dffe304b84..f76871460c4a 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_32_S_qkv_128_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_32_S_qkv_128_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:1afae26383dce7307d9b12c1e8b6559dc65b7762e8108975a46ec5e7df8dff84 +oid sha256:ce5bcf4c0194abce62b39cd408d5a449e3725badf28d51510e7775df30d0ccd9 size 685912 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_32_S_qkv_128_softcapping_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_32_S_qkv_128_softcapping_sm90.cubin.cpp index b903a8d92713..daf415f99a86 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_32_S_qkv_128_softcapping_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_fp16_fp32_64_32_S_qkv_128_softcapping_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:754a6b3bf9c764fa535c2e73dda1f58d29f37013e421405229d2a0d43d854b09 +oid sha256:fe521017d6cb30dc5f434b809068533a31db662dfa8d19af927ff79761230c62 size 371779 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_128_32_ldgsts_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_128_32_ldgsts_sm90.cubin.cpp index 1ca46e799df6..e2ee736b49d0 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_128_32_ldgsts_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_128_32_ldgsts_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:7c920a8fccb239403c050d00d23e5784c1f3c67598cfa7b26f2e57514964ed4f -size 1018174 +oid sha256:dd930ed415b0303a973a37550ee33fa4975ad6be0cc58d461370b127f9a90f8e +size 1020542 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_128_64_ldgsts_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_128_64_ldgsts_sm90.cubin.cpp index 393bd489fe20..95d9b2bf6473 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_128_64_ldgsts_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_128_64_ldgsts_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:fbd0c0ca6cb0657009e82fd343f1115901db6ab10961e9ec313dcbfb0d168c33 -size 1053694 +oid sha256:4f2b243127e1ce00a850a10cca104ffc42512711f434fbdf8683eeeb49b8ce42 +size 1056062 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_fp32_128_32_ldgsts_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_fp32_128_32_ldgsts_sm90.cubin.cpp index 6f2beba416cb..0c093db643c3 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_fp32_128_32_ldgsts_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_fp32_128_32_ldgsts_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:2f59cf8d14c75513d555ce75a2d93e552ec0a82279c40bbea287c7f4beea5fa0 -size 1005556 +oid sha256:2ce9cc89b1db7f7e4b76b94cf1c3b04db49a2d86b529b1fc85b19057a99bc9fa +size 1007924 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_fp32_128_64_ldgsts_sm90.cubin.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_fp32_128_64_ldgsts_sm90.cubin.cpp index 9365bad44616..c24e239dd0c0 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_fp32_128_64_ldgsts_sm90.cubin.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_fp16_fp32_128_64_ldgsts_sm90.cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:0322cb4741792dbaeba2d75a05330fee7995b6f15749f39c220252a526770d8a -size 1066334 +oid sha256:e176513fa0074d688620299dfca53adc3902491e97ea9b6938a4ceb2fcf17ef5 +size 1068702 diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp index 68c5492bef16..a0f68d8080a2 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp @@ -140,28 +140,47 @@ void FusedMHARunnerV2::setupKernelParams(MHARunnerParams runnerParams) mKernelParams.softmax_stats_ptr = runnerParams.softmaxStatsPtr; mKernelParams.softmax_stats_stride_in_bytes = sizeof(float) * mFixedParams.numQHeads; - // Packed QKV input layout. - mKernelParams.qkv_stride_in_bytes = get_size_in_bytes(mFixedParams.numQHeads * mFixedParams.headSize - + mFixedParams.numKvHeads * mFixedParams.headSize + mFixedParams.numKvHeads * mFixedParams.headSizeV, - mFixedParams.dataType); - // Contiguous Q input layout. - mKernelParams.q_stride_in_bytes - = get_size_in_bytes(mFixedParams.numQHeads * mFixedParams.headSize, mFixedParams.dataType); - // Set the kv_stride_in_bytes when separate kv buffer is used. - if (mFixedParams.attentionInputLayout == AttentionInputLayout::Q_PAGED_KV) - { - // Paged kv cache layout. - mKernelParams.kv_stride_in_bytes = get_size_in_bytes( - runnerParams.pagedKvCache.mTokensPerBlock * mFixedParams.headSize, mFixedParams.dataType); - // only for deepseek - mKernelParams.v_stride_in_bytes = mKernelParams.kv_stride_in_bytes; - } - else if (mFixedParams.attentionInputLayout == AttentionInputLayout::Q_CONTIGUOUS_KV) - { - // Contiguous kv input layout. - mKernelParams.kv_stride_in_bytes - = get_size_in_bytes(2 * mFixedParams.numKvHeads * mFixedParams.headSize, mFixedParams.dataType); + if (mFixedParams.attentionInputLayout == AttentionInputLayout::PACKED_QKV) + { + // Packed QKV input layout, [B, S, H * D + H_kv * D + H_kv * Dv]. + mKernelParams.qkv_ptr = runnerParams.qkvPtr; + mKernelParams.q_stride_in_bytes = mKernelParams.k_stride_in_bytes = mKernelParams.v_stride_in_bytes + = get_size_in_bytes(mFixedParams.numQHeads * mFixedParams.headSize + + mFixedParams.numKvHeads * mFixedParams.headSize + + mFixedParams.numKvHeads * mFixedParams.headSizeV, + mFixedParams.dataType); } + else + { + // Contiguous Q input layout, [B, S, H, D]. + mKernelParams.q_ptr = runnerParams.qPtr; + mKernelParams.q_stride_in_bytes + = get_size_in_bytes(mFixedParams.numQHeads * mFixedParams.headSize, mFixedParams.dataType); + + // Separate q and kv buffers may have different q and kv sequence lengths. + mKernelParams.cu_kv_seqlens = reinterpret_cast(runnerParams.cuKvSeqLenPtr); + + if (mFixedParams.attentionInputLayout == AttentionInputLayout::Q_CONTIGUOUS_KV) + { + // Contiguous kv input layout, [B, S, H_kv * D + H_kv * Dv]. + mKernelParams.kv_ptr = runnerParams.kvPtr; + mKernelParams.k_stride_in_bytes = mKernelParams.v_stride_in_bytes = get_size_in_bytes( + mFixedParams.numKvHeads * (mFixedParams.headSize + mFixedParams.headSizeV), mFixedParams.dataType); + } + else if (mFixedParams.attentionInputLayout == AttentionInputLayout::Q_PAGED_KV) + { + // Paged kv cache layout. + mKernelParams.paged_kv_cache = runnerParams.pagedKvCache.copyKVBlockArrayForContextFMHA(); + mKernelParams.k_stride_in_bytes = get_size_in_bytes( + runnerParams.pagedKvCache.mTokensPerBlock * mFixedParams.headSize, mFixedParams.dataType); + // If d == dv, then v_stride_in_bytes == k_stride_in_bytes. + // For DeepSeek MLA, which is the only case where d != dv, V is padded to the sizeof K. + // Thus, v_stride_in_bytes always equals to k_stride_in_bytes so far. + mKernelParams.v_stride_in_bytes = mKernelParams.k_stride_in_bytes; + } + } + + mKernelParams.o_ptr = runnerParams.outputPtr; // Set the output buffer stride in bytes. mKernelParams.o_stride_in_bytes = get_size_in_bytes(mFixedParams.numQHeads * mFixedParams.headSizeV, mFixedParams.dataTypeOut); @@ -214,11 +233,6 @@ void FusedMHARunnerV2::setupKernelParams(MHARunnerParams runnerParams) mFixedParams.numQHeads, runnerParams.kvSeqLen, mFixedParams.tpSize, mFixedParams.tpRank, scale_after_alibi); } - // Set device pointers. - mKernelParams.qkv_ptr = runnerParams.qkvPtr; - mKernelParams.q_ptr = runnerParams.qPtr; - mKernelParams.kv_ptr = runnerParams.kvPtr; - mKernelParams.o_ptr = runnerParams.outputPtr; if (mFixedParams.attentionMaskType == ContextAttentionMaskType::CUSTOM_MASK) { mKernelParams.packed_mask_ptr = runnerParams.packedMaskPtr; @@ -237,18 +251,6 @@ void FusedMHARunnerV2::setupKernelParams(MHARunnerParams runnerParams) mKernelParams.scale_bmm2_d = reinterpret_cast(runnerParams.scaleBmm2Ptr); } - // Separate q and kv buffers may have different q and kv sequence lengths. - if (mFixedParams.attentionInputLayout != AttentionInputLayout::PACKED_QKV) - { - mKernelParams.cu_kv_seqlens = reinterpret_cast(runnerParams.cuKvSeqLenPtr); - } - - // Paged kv fmha. - if (mFixedParams.attentionInputLayout == AttentionInputLayout::Q_PAGED_KV) - { - mKernelParams.paged_kv_cache = runnerParams.pagedKvCache.copyKVBlockArrayForContextFMHA(); - } - // for sage attention mKernelParams.sage.q.scales = runnerParams.qScalePtr; mKernelParams.sage.k.scales = runnerParams.kScalePtr; @@ -293,11 +295,18 @@ void FusedMHARunnerV2::setupLaunchParams(MHARunnerParams runnerParams) mLaunchParams.total_kv_seqlen = mFixedParams.isSPadded ? runnerParams.b * runnerParams.kvSeqLen : runnerParams.totalKvSeqLen; - // Next power of 2 head size. TLLM_CHECK_WITH_INFO(mFixedParams.headSize > 0, "Head size should be greater than 0."); - mLaunchParams.padded_d = (mFixedParams.headSize & (mFixedParams.headSize - 1)) == 0 + // Pad head size to next power of 2. + int padded_d_next_power_of_2 = (mFixedParams.headSize & (mFixedParams.headSize - 1)) == 0 ? mFixedParams.headSize : pow(2, int(log2(mFixedParams.headSize)) + 1); + // In fact, due to 128B swizzle mode of TMA, only 128 bytes alignment is required, + // so we pad head size to next multiply of 128B. + int d_per_group = 128 / get_size_in_bytes(mFixedParams.dataType); + int d_groups = (mFixedParams.headSize + d_per_group - 1) / d_per_group; + int padded_d_next_multiply_of_128byte = d_groups * d_per_group; + // Choose the smaller one to save SMEM. + mLaunchParams.padded_d = std::min(padded_d_next_power_of_2, padded_d_next_multiply_of_128byte); bool const isSm70 = (mSM == kSM_70); bool const isSm90 = (mSM == kSM_90); @@ -453,273 +462,162 @@ void FusedMHARunnerV2::setupLaunchParams(MHARunnerParams runnerParams) //////////////////////////////////////////////////////////////////////////////////////////////////// // TMA descriptors are used as grid_constant parameters (remove MemCpyH2D operations) -void FusedMHARunnerV2::setPackedQkvTmaDescriptors(MHARunnerParams runnerParams) +void FusedMHARunnerV2::setTmaDescriptors(MHARunnerParams runnerParams) { + const uint32_t d = mKernelParams.d; + const uint32_t dv = mKernelParams.dv; + const uint32_t h = mKernelParams.h; + const uint32_t h_kv = mKernelParams.h_kv; + const uint32_t total_q_seqlen = mLaunchParams.total_q_seqlen; + const uint32_t total_kv_seqlen = mLaunchParams.total_kv_seqlen; + + uint64_t const d_in_bytes = get_size_in_bytes(d, mFixedParams.dataType); + uint64_t const dv_in_bytes = get_size_in_bytes(dv, mFixedParams.dataType); + // split D into multiple groups in order to match the TMA swizzle mode (128B) - uint32_t const d_in_bytes = get_size_in_bytes(mLaunchParams.padded_d, mFixedParams.dataType); - uint32_t const d_groups = d_in_bytes > 128 ? d_in_bytes / 128 : 1; + uint32_t const padded_d_in_bytes = get_size_in_bytes(mLaunchParams.padded_d, mFixedParams.dataType); + uint32_t const d_groups = padded_d_in_bytes > 128 ? padded_d_in_bytes / 128 : 1; + uint32_t const d_bytes_per_group = padded_d_in_bytes / d_groups; + uint32_t const d_per_group = mLaunchParams.padded_d / d_groups; - // separate q, k, v and o tma descriptors - Multiple_tma_descriptor<4> qkv_tma_descriptor; + uint32_t q_step = 0, kv_step = 0; + xmmaKernel->getStepSize(q_step, kv_step, mKernelParams, mLaunchParams); - // tensor size - uint32_t tensor_size_qkv[4]; - if (mKernelParams.h_kv < mKernelParams.h) - { - // if multi-query or grouped-query - tensor_size_qkv[2] = 1; - tensor_size_qkv[1] = (mKernelParams.h + 2 * mKernelParams.h_kv); - tensor_size_qkv[0] = mKernelParams.d; // mKernelParams.d; - } - else - { - tensor_size_qkv[2] = 3; - tensor_size_qkv[1] = mKernelParams.h; - tensor_size_qkv[0] = mKernelParams.d; // mKernelParams.d; - } + auto const layout = mFixedParams.attentionInputLayout; - // O : [TOTAL, 1, h, d] - uint32_t tensor_size_o[4]; - tensor_size_o[0] = mKernelParams.d; - tensor_size_o[1] = mKernelParams.h; - tensor_size_o[2] = 1; + // Q Layout: [total_seqlen, H, D] + const uint32_t tensor_size_q[3] = {d, h, total_q_seqlen}; - // box size for k and v - uint32_t box_size[4]; - // Update this on device? - box_size[2] = 1; - box_size[1] = 1; - box_size[0] = mLaunchParams.padded_d / d_groups; + // Stride size in bytes. Assumes least significant dim is 1 + const uint64_t tensor_stride_q[2] = {d_in_bytes, uint64_t(mKernelParams.q_stride_in_bytes)}; - // stride size in bytes. Assumes least significant dim is 1 (?) - uint64_t tensor_stride_qkv[3]; - tensor_stride_qkv[0] = get_size_in_bytes(tensor_size_qkv[0], mFixedParams.dataType); // d - tensor_stride_qkv[1] = tensor_size_qkv[1] * tensor_stride_qkv[0]; // d*h - tensor_stride_qkv[2] = mKernelParams.qkv_stride_in_bytes; + // Starting memory address + char const* q_ptr = reinterpret_cast( + layout == AttentionInputLayout::PACKED_QKV ? mKernelParams.qkv_ptr : mKernelParams.q_ptr); - uint64_t tensor_stride_o[3]; - tensor_stride_o[0] = get_size_in_bytes(tensor_size_o[0], mFixedParams.dataTypeOut); // d - tensor_stride_o[1] = tensor_size_o[1] * tensor_stride_o[0]; // d*h - tensor_stride_o[2] = tensor_size_o[2] * tensor_stride_o[1]; // d*h*1 + // Box size of TMA + const uint32_t box_size_q[3] = {d_per_group, 1, q_step}; - // traversal stride - uint32_t traversal_stride_qkv[4] = {1, 1, 1, 1}; - uint32_t traversal_stride_o[4] = {1, 1, 1, 1}; + // Traversal stride. + const uint32_t traversal_stride[3] = {1, 1, 1}; - // OOB fill zeros - uint32_t oob_fill = 0; + // OOB fill zeros. + const uint32_t oob_fill = 0; - // FP32 to TF32 conversion disabled - uint32_t fp32_to_tf32 = 0; + // FP32 to TF32 conversion disabled. + const uint32_t fp32_to_tf32 = 0; - // gmma descriptor mode - uint32_t const d_bytes_per_group = d_in_bytes / d_groups; + // GMMA descriptor mode. cudaTmaDescSwizzle const swizzle_mode = (d_bytes_per_group > 64 ? cudaTmaDescSwizzle::SWIZZLE_128B : (d_bytes_per_group > 32 ? cudaTmaDescSwizzle::SWIZZLE_64B : cudaTmaDescSwizzle::SWIZZLE_32B)); - uint32_t q_step = 0, kv_step = 0; - xmmaKernel->getStepSize(q_step, kv_step, mKernelParams, mLaunchParams); - - // QKV [TOTAL, 3, h, d] - // NOTE: we may need to use actual seqlen to set oob_value - auto const* qkv_ptr = static_cast(mKernelParams.qkv_ptr); - tensor_size_qkv[3] = mLaunchParams.total_q_seqlen; - // O [TOTAL, 1, h, d] - auto* o_ptr = static_cast(mKernelParams.o_ptr); - tensor_size_o[3] = mLaunchParams.total_q_seqlen; - - // Q: STEP_Q - box_size[3] = q_step; // Desc Format (data type). cudaTmaDescFormat const desc_format = (get_size_in_bytes(mFixedParams.dataType) == 1) ? cudaTmaDescFormat::U8 : cudaTmaDescFormat::F16_RN; - qkv_tma_descriptor.set_tma_desctriptor(qkv_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED, - swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_qkv, tensor_stride_qkv, - traversal_stride_qkv, box_size, oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_q); - // K/V: STEP_KV - box_size[3] = kv_step; - qkv_tma_descriptor.set_tma_desctriptor(qkv_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED, - swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_qkv, tensor_stride_qkv, - traversal_stride_qkv, box_size, oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_kv); + Multiple_tma_descriptor<3> qo_tma_descriptor; - // Separate TMA descriptor for V when d != dv in packed qkv input layout, e.g. MLA + 192/128 dims - if (mKernelParams.d != mKernelParams.dv) - { - // view V as [total_seq_len, 1, h, dv] - tensor_size_qkv[0] = mKernelParams.dv; - tensor_size_qkv[1] = mKernelParams.h; - tensor_size_qkv[2] = 1; - - tensor_stride_qkv[0] = get_size_in_bytes(tensor_size_qkv[0], mFixedParams.dataType); - tensor_stride_qkv[1] = 0; // not used - - size_t v_offset = 2 * mKernelParams.h * mKernelParams.d * get_size_in_bytes(mFixedParams.dataType); - qkv_tma_descriptor.set_tma_desctriptor(qkv_ptr + v_offset, desc_format, - cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, - tensor_size_qkv, tensor_stride_qkv, traversal_stride_qkv, box_size, oob_fill, fp32_to_tf32, - &mKernelParams.tma_desc_v); - } + // Q + qo_tma_descriptor.set_tma_desctriptor(q_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, + cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_q, tensor_stride_q, traversal_stride, box_size_q, + oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_q); - // O: 16 - // Note: sliding window causal kernel currently has reg spill when TMA store is enabled - box_size[3] = 16; + // O if ((get_size_in_bytes(mFixedParams.dataTypeOut) == 1) && mLaunchParams.attention_mask_type != ContextAttentionMaskType::SLIDING_OR_CHUNKED_CAUSAL) { - qkv_tma_descriptor.set_tma_desctriptor(o_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED, - swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_o, tensor_stride_o, traversal_stride_o, - box_size, oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_o); - } -} + // O Layout: [total_seqlen, H, DV] + const uint32_t tensor_size_o[3] = {dv, h, total_q_seqlen}; -//////////////////////////////////////////////////////////////////////////////////////////////////// + const uint64_t tensor_stride_o[2] + = {get_size_in_bytes(dv, mFixedParams.dataTypeOut), uint64_t(mKernelParams.o_stride_in_bytes)}; -// Contiguous in the shape of [B, S, H, D]. -// Contiguous KV in the shape of [B, S, 2, H, D]. -// Paged KV has [B, 2, NumBlocksPerSequence] buffers, -// and each points to the contiguous buffer with shape [H, TokensPerBlock, D] -// TMA descriptors need cudaMemcpyAsync since we need multiple tma descriptors in device memory. -void FusedMHARunnerV2::setSeparateQKvTmaDescriptors(MHARunnerParams runnerParams) -{ - // split D into multiple groups in order to match the TMA swizzle mode (128B) - uint32_t const d_in_bytes = get_size_in_bytes(mLaunchParams.padded_d, mFixedParams.dataType); - uint32_t const d_groups = d_in_bytes > 128 ? d_in_bytes / 128 : 1; + char* o_ptr = reinterpret_cast(mKernelParams.o_ptr); - uint32_t q_step = 0, kv_step = 0; - xmmaKernel->getStepSize(q_step, kv_step, mKernelParams, mLaunchParams); + // Box size of TMA + const uint32_t box_size_o[3] = {d_per_group, 1, 16}; - // Separate q, and paged kv tma descriptors. - Multiple_tma_descriptor<4> qo_tma_descriptor; - Multiple_tma_descriptor<4> kv_tma_descriptor; - // Contiguous Q - // query tensor size [B x S, 1, H, D] - uint32_t tensor_size_qo[4]; - tensor_size_qo[3] = mLaunchParams.total_q_seqlen; - tensor_size_qo[2] = 1; - tensor_size_qo[1] = mKernelParams.h; - tensor_size_qo[0] = mKernelParams.d; - - // box size for q and o - uint32_t box_size_qo[4]; - box_size_qo[3] = q_step; - box_size_qo[2] = 1; - box_size_qo[1] = 1; - box_size_qo[0] = mLaunchParams.padded_d / d_groups; - - // stride size in bytes. - uint64_t tensor_stride_qo[3]; - tensor_stride_qo[0] = get_size_in_bytes(tensor_size_qo[0], mFixedParams.dataType); - tensor_stride_qo[1] = tensor_size_qo[1] * tensor_stride_qo[0]; - tensor_stride_qo[2] = tensor_size_qo[2] * tensor_stride_qo[1]; - - // traversal stride - uint32_t traversal_stride[4] = {1, 1, 1, 1}; - - // OOB fill zeros - uint32_t oob_fill = 0; - - // FP32 to TF32 conversion disabled - uint32_t fp32_to_tf32 = 0; + // dataTypeOut may be different with dataType, so desc_format and swizzle_mode + // may be incorrect. For example, QKV are in bf16 while O is in fp8. + // Luckily, this case doesn't exist so far. But we should keep one eye on it. + qo_tma_descriptor.set_tma_desctriptor(o_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED, + swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_o, tensor_stride_o, traversal_stride, + box_size_o, oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_o); + } - // Desc Format (data type). - cudaTmaDescFormat const desc_format - = (get_size_in_bytes(mFixedParams.dataType) == 1) ? cudaTmaDescFormat::U8 : cudaTmaDescFormat::F16_RN; + if (layout == AttentionInputLayout::Q_PAGED_KV) + { + // KV in q_paged_kv uses 4D tensor + // Layout: [INT32_MAX, H_KV, TokensPerBlock, D] + const uint32_t tokens_per_block = mKernelParams.paged_kv_cache.mTokensPerBlock; + const uint32_t tensor_size_k[4] = {d, tokens_per_block, h_kv, INT_MAX}; + const uint32_t tensor_size_v[4] = {dv, tokens_per_block, h_kv, INT_MAX}; - // gmma descriptor mode - uint32_t const d_bytes_per_group = d_in_bytes / d_groups; - cudaTmaDescSwizzle const swizzle_mode = (d_bytes_per_group > 64 - ? cudaTmaDescSwizzle::SWIZZLE_128B - : (d_bytes_per_group > 32 ? cudaTmaDescSwizzle::SWIZZLE_64B : cudaTmaDescSwizzle::SWIZZLE_32B)); + const uint64_t tensor_stride_k[3] = {uint64_t(mKernelParams.k_stride_in_bytes / tokens_per_block), // d + uint64_t(mKernelParams.k_stride_in_bytes), // d * 64 + uint64_t(mKernelParams.paged_kv_cache.mBytesPerBlock)}; + const uint64_t tensor_stride_v[3] + = {// we cannot use dv * Kernel_traits::ELEMENT_BYTES because V may be padded (MLA) + uint64_t(mKernelParams.v_stride_in_bytes / tokens_per_block), // dv + uint64_t(mKernelParams.v_stride_in_bytes), // dv * 64 + uint64_t(mKernelParams.paged_kv_cache.mBytesPerBlock)}; - // Q ptr. - auto const* q_ptr = static_cast(mKernelParams.q_ptr); + char const* kv_ptr = reinterpret_cast(runnerParams.pagedKvCache.mPrimaryPoolPtr); - // Q: STEP_Q. - qo_tma_descriptor.set_tma_desctriptor(q_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, - cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_qo, tensor_stride_qo, traversal_stride, box_size_qo, - oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_q); + const uint32_t box_size_kv[4] = {d_per_group, std::min(tokens_per_block, kv_step), 1, 1}; - // O ptr. - auto const* o_ptr = static_cast(mKernelParams.o_ptr); - // Note (added by Yuxin): TMA descriptor for o here might be problematic if d and dv are different. + TLLM_CHECK(kv_step % tokens_per_block == 0 || tokens_per_block % kv_step == 0); + mKernelParams.blocks_per_tma_load = std::max(1, kv_step / tokens_per_block); + mKernelParams.blocks_per_tma_load_log2 = log2(mKernelParams.blocks_per_tma_load); - // O: 16. Reuse - box_size_qo[3] = 16; - if ((get_size_in_bytes(mFixedParams.dataTypeOut) == 1) - && mLaunchParams.attention_mask_type != ContextAttentionMaskType::SLIDING_OR_CHUNKED_CAUSAL) + const uint32_t traversal_stride[4] = {1, 1, 1, 1}; + + Multiple_tma_descriptor<4> kv_tma_descriptor; + // K + kv_tma_descriptor.set_tma_desctriptor(kv_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED, + swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_k, tensor_stride_k, traversal_stride, + box_size_kv, oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_k); + // V + kv_tma_descriptor.set_tma_desctriptor(kv_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED, + swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_v, tensor_stride_v, traversal_stride, + box_size_kv, oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_v); + } + else { - qo_tma_descriptor.set_tma_desctriptor(o_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED, - swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_qo, tensor_stride_qo, traversal_stride, - box_size_qo, oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_o); - } - - // Contiguous KV layout [B, S, 2, H, D]. - if (mFixedParams.attentionInputLayout == AttentionInputLayout::Q_CONTIGUOUS_KV) - { - // Per batch tensor size. - uint32_t tensor_size_kv[4]; - // Maximum number of blocks in this device. - tensor_size_kv[3] = mLaunchParams.total_kv_seqlen; - tensor_size_kv[2] = 2; - tensor_size_kv[1] = mKernelParams.h_kv; - tensor_size_kv[0] = mKernelParams.d; - - // Box size for k and v. - uint32_t box_size_kv[4]; - box_size_kv[3] = kv_step; - box_size_kv[2] = 1; - box_size_kv[1] = 1; - box_size_kv[0] = mLaunchParams.padded_d / d_groups; - - // Stride size in bytes. - uint64_t tensor_stride_kv[3]; - tensor_stride_kv[0] = get_size_in_bytes(tensor_size_kv[0], mFixedParams.dataType); - tensor_stride_kv[1] = tensor_size_kv[1] * tensor_stride_kv[0]; - tensor_stride_kv[2] = tensor_size_kv[2] * tensor_stride_kv[1]; - - // Set the paged_kv tma descriptor. - kv_tma_descriptor.set_tma_desctriptor(runnerParams.kvPtr, desc_format, - cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, - tensor_size_kv, tensor_stride_kv, traversal_stride, box_size_kv, oob_fill, fp32_to_tf32, - &mKernelParams.tma_desc_kv); - } - else if (mFixedParams.attentionInputLayout == AttentionInputLayout::Q_PAGED_KV) - { - // Paged KV - // Per batch tensor size. - uint32_t tokens_per_block = uint32_t(mKernelParams.paged_kv_cache.mTokensPerBlock); - uint32_t tensor_size_kv[4]; - // Maximum number of blocks in this device. - tensor_size_kv[3] = mLaunchParams.total_device_memory / mKernelParams.paged_kv_cache.mBytesPerBlock; - tensor_size_kv[2] = mKernelParams.h_kv; - tensor_size_kv[1] = tokens_per_block; - tensor_size_kv[0] = mKernelParams.d; - - // Box size for k and v. - uint32_t box_size_kv[4]; - box_size_kv[3] = 1; - box_size_kv[2] = 1; - box_size_kv[1] = std::min(tokens_per_block, kv_step); - box_size_kv[0] = mLaunchParams.padded_d / d_groups; - - TLLM_CHECK_WITH_INFO( - tokens_per_block % 2 == 0, "FMHA with paged kv cache needs tokens_per_block to be power of 2 !"); - mKernelParams.blocks_per_tma_load = std::max(1, int32_t(kv_step / tokens_per_block)); - mKernelParams.blocks_per_tma_load_log2 = log2(mKernelParams.blocks_per_tma_load); + // Otherwise KV uses 3D tensor + const uint32_t tensor_size_k[3] = {d, h_kv, total_kv_seqlen}; + const uint32_t tensor_size_v[3] = {dv, h_kv, total_kv_seqlen}; - // Stride size in bytes. - uint64_t tensor_stride_kv[3]; - tensor_stride_kv[0] = get_size_in_bytes(tensor_size_kv[0], mFixedParams.dataType); - tensor_stride_kv[1] = tensor_size_kv[1] * tensor_stride_kv[0]; - tensor_stride_kv[2] = tensor_size_kv[2] * tensor_stride_kv[1]; + const uint64_t tensor_stride_k[2] = {d_in_bytes, uint64_t(mKernelParams.k_stride_in_bytes)}; + const uint64_t tensor_stride_v[2] = {dv_in_bytes, uint64_t(mKernelParams.v_stride_in_bytes)}; - // Set the paged_kv tma descriptor. - kv_tma_descriptor.set_tma_desctriptor(runnerParams.pagedKvCache.mPrimaryPoolPtr, desc_format, - cudaTmaDescInterleave::INTERLEAVE_DISABLED, swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, - tensor_size_kv, tensor_stride_kv, traversal_stride, box_size_kv, oob_fill, fp32_to_tf32, - &mKernelParams.tma_desc_kv); + const uint32_t box_size_kv[3] = {d_per_group, 1, kv_step}; + + char const *k_ptr, *v_ptr; + + if (layout == AttentionInputLayout::PACKED_QKV) + { + // Layout: [total_seqlen, (H, D) + (H_KV, D) + (H_KV, DV)] + k_ptr = q_ptr + h * d_in_bytes; + v_ptr = k_ptr + h_kv * d_in_bytes; + } + else if (layout == AttentionInputLayout::Q_CONTIGUOUS_KV) + { + // Layout, [B, S, H_kv * D + H_kv * Dv]. + k_ptr = reinterpret_cast(mKernelParams.kv_ptr); + v_ptr = k_ptr + h_kv * d_in_bytes; + } + + Multiple_tma_descriptor<3> kv_tma_descriptor; + // K + kv_tma_descriptor.set_tma_desctriptor(k_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED, + swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_k, tensor_stride_k, traversal_stride, + box_size_kv, oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_k); + // V + kv_tma_descriptor.set_tma_desctriptor(v_ptr, desc_format, cudaTmaDescInterleave::INTERLEAVE_DISABLED, + swizzle_mode, cudaTmaDescPromotion::PROMOTION_DISABLED, tensor_size_v, tensor_stride_v, traversal_stride, + box_size_kv, oob_fill, fp32_to_tf32, &mKernelParams.tma_desc_v); } } @@ -734,13 +632,7 @@ void FusedMHARunnerV2::run(MHARunnerParams runnerParams) // Need to set tma descriptors additionally. if (mSM == kSM_90 && mLaunchParams.use_tma) { - switch (mFixedParams.attentionInputLayout) - { - case AttentionInputLayout::PACKED_QKV: setPackedQkvTmaDescriptors(runnerParams); break; - case AttentionInputLayout::Q_CONTIGUOUS_KV: - case AttentionInputLayout::Q_PAGED_KV: setSeparateQKvTmaDescriptors(runnerParams); break; - default: TLLM_CHECK_WITH_INFO(false, "Unsupported attention input layout."); - } + setTmaDescriptors(runnerParams); } // Check if the sliding window size is valid or not. if (mFixedParams.attentionInputLayout == AttentionInputLayout::Q_PAGED_KV diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.h b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.h index ac25da6d0555..afa8eb949a66 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.h +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.h @@ -71,11 +71,8 @@ class FusedMHARunnerV2 // Set the launch params to select kernels. void setupLaunchParams(MHARunnerParams runnerParams); - // Set the tma descriptors for packed qkv input. - void setPackedQkvTmaDescriptors(MHARunnerParams runnerParams); - - // Set the tma descriptors for separate q and kv input. - void setSeparateQKvTmaDescriptors(MHARunnerParams runnerParams); + // Set the tma descriptors. + void setTmaDescriptors(MHARunnerParams runnerParams); // Check if it is a valid sequence length (only used by non-flash-attention kernels). bool isValidS(int s) const; diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h index 9e000f9c872d..96435cca5286 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h @@ -342,6 +342,10 @@ struct Fused_multihead_attention_params_v2 void const* qkv_ptr; // The separate Q matrice. void const* q_ptr; + // The separate K matrice. + void const* k_ptr; + // The separate V matrice. + void const* v_ptr; // The separate KV matrice. void const* kv_ptr; // The separate paged kv cache. @@ -353,14 +357,12 @@ struct Fused_multihead_attention_params_v2 // The Softmax stats vector of layout [2, B, S, H], including softmax_sum and softmax_max void* softmax_stats_ptr; - // The stride between rows of the Q, K and V matrices. - int64_t qkv_stride_in_bytes; - // The stride between rows of the separate Q matrice. + // The stride between rows of Q. int64_t q_stride_in_bytes; - // The stride between rows of the separate KV matrice. - int64_t kv_stride_in_bytes; - // The stride between rows of the separate V matrice, set if it is not same as that of K. - int64_t v_stride_in_bytes = 0; + // The stride between rows of K. + int64_t k_stride_in_bytes; + // The stride between rows of V. + int64_t v_stride_in_bytes; // The stride between matrices of packed mask. int64_t packed_mask_stride_in_bytes; // The stride between rows of O. @@ -375,7 +377,8 @@ struct Fused_multihead_attention_params_v2 // Kv in packed qkv layout: [B, S, 3, H, D] // Contiguous kv layout: [B, 2, H, S, D]. // Paged kv layout: [UINT32_MAX, H, Tokens_per_block, D]. - cudaTmaDesc tma_desc_kv; + cudaTmaDesc tma_desc_k; + cudaTmaDesc tma_desc_v; // Tma descriptor for o cudaTmaDesc tma_desc_o; @@ -433,10 +436,6 @@ struct Fused_multihead_attention_params_v2 float* scales; } q, k, v; } sage; - - // Separate TMA descriptor for V when d != dv in packed qkv input layout, e.g. MLA + 192/128 dims - // We need to add this parameter in the tail of the struct for cubin compatibility - cudaTmaDesc tma_desc_v; }; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -450,7 +449,7 @@ struct Launch_params int total_q_seqlen = 0; // total kv sequence length. int total_kv_seqlen = 0; - // padded head size (new power of 2) for tma descriptors. + // padded head size for tma descriptors. int padded_d = 0; // flags to control small batch kernel choice // true: never unroll diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/allreduce_gemm/allreduce_gemm_impl_sm100.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/allreduce_gemm/allreduce_gemm_impl_sm100.h index ed18541d0ace..a4be82607a81 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/allreduce_gemm/allreduce_gemm_impl_sm100.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/allreduce_gemm/allreduce_gemm_impl_sm100.h @@ -221,9 +221,6 @@ class GemmAllReduceImplTwoshot_Sm100 : public GemmAllReduceImplInterface { MPI_group_barrier(_ranks); } - - TLLM_CUDA_CHECK(cudaStreamCreate(&_memcpy_stream)); - TLLM_CUDA_CHECK(cudaEventCreate(&_fork_join_event)); } int free() override @@ -267,8 +264,6 @@ class GemmAllReduceImplTwoshot_Sm100 : public GemmAllReduceImplInterface DeviceAllocationNvls _tile_barriers; DeviceAllocationNvls _completion_barriers; DeviceAllocationNvls _stage_buf; - cudaStream_t _memcpy_stream; - cudaEvent_t _fork_join_event; }; GemmAllReduceImplTwoshot_Sm100() diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/allreduce_gemm/allreduce_gemm_impl_sm90.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/allreduce_gemm/allreduce_gemm_impl_sm90.h index ab867b69a87b..fb446b451d8d 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/allreduce_gemm/allreduce_gemm_impl_sm90.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/allreduce_gemm/allreduce_gemm_impl_sm90.h @@ -186,9 +186,6 @@ class GemmAllReduceImplTwoshot_Sm90 : public GemmAllReduceImplInterface { MPI_group_barrier(_ranks); } - - TLLM_CUDA_CHECK(cudaStreamCreate(&_memcpy_stream)); - TLLM_CUDA_CHECK(cudaEventCreate(&_fork_join_event)); } int free() override @@ -232,8 +229,6 @@ class GemmAllReduceImplTwoshot_Sm90 : public GemmAllReduceImplInterface DeviceAllocationNvls _tile_barriers; DeviceAllocationNvls _completion_barriers; DeviceAllocationNvls _stage_buf; - cudaStream_t _memcpy_stream; - cudaEvent_t _fork_join_event; }; GemmAllReduceImplTwoshot_Sm90() 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 912c3553bb00..c7c9a55b9590 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h @@ -845,10 +845,10 @@ struct GemmProfilerBackend mWType = wtype; mOType = otype; mNumExperts = num_experts; - mNumExpertsPerNode = num_experts / (parallelism_config.ep_size * parallelism_config.tp_size); + mNumExpertsPerNode = num_experts / parallelism_config.ep_size; mK = k; mExpertHiddenSize = hidden_size; - mExpertInterSize = inter_size; + mExpertInterSize = inter_size; // Already divided by tp_size mGroupSize = group_size; mActivationType = activation_type; mBias = bias; diff --git a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp index 7eb6682ec7a7..52471c70d7f1 100644 --- a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp +++ b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp @@ -56,7 +56,8 @@ FmhaDispatcher::FmhaDispatcher(MHARunnerFixedParams fixedParams) else { TLLM_CHECK_WITH_INFO(mFixedParams.dataType == mFixedParams.dataTypeKv, - "KV cache data type should be the same as input data type."); + "KV cache data type %s is not the same as input data type %s.", + data_type_to_string(mFixedParams.dataTypeKv).c_str(), data_type_to_string(mFixedParams.dataType).c_str()); // For FP8 MLA generation, the output type is BF16, which could be different from the input type. // So we shouldn't do this check anymore. diff --git a/cpp/tensorrt_llm/kernels/moePrepareKernels.cu b/cpp/tensorrt_llm/kernels/moePrepareKernels.cu index 5914ce14ee0b..6ca40a948aa3 100644 --- a/cpp/tensorrt_llm/kernels/moePrepareKernels.cu +++ b/cpp/tensorrt_llm/kernels/moePrepareKernels.cu @@ -319,19 +319,19 @@ __global__ void computeCumsumDevice(int* sendCountsCumsum, int* recvCountsCumsum } } -template +template class PacketPipeline { public: __device__ __inline__ PacketPipeline( - void* bufferBase, STEP_COMMUNICATOR_TYPE* stepCommunicator, int* sharedNewStepPtr, bool isSender) + void* bufferBase, StepCommunicatorBase* stepCommunicator, int* sharedNewStepPtr, bool isSender) : bufferBase(bufferBase) , stepCommunicator(stepCommunicator) , shared_new_step(sharedNewStepPtr) { step = 0; needRelease = false; - packetId = isSender ? 0 : PACKET_PER_STEP - 1; + packetId = isSender ? 0 : PipelineConfig::PACKET_PER_STEP - 1; } __device__ __forceinline__ void* getFirstSendPacket() @@ -343,9 +343,10 @@ public: { packetId++; - if (packetId < PACKET_PER_STEP) + if (packetId < PipelineConfig::PACKET_PER_STEP) { - return acquireNewStep ? bufferBase + step * PACKET_PER_STEP * PACKET_SIZE + packetId * PACKET_SIZE + return acquireNewStep ? bufferBase + step * PipelineConfig::PACKET_PER_STEP * PipelineConfig::PACKET_SIZE + + packetId * PipelineConfig::PACKET_SIZE : nullptr; } @@ -365,7 +366,7 @@ public: { step = *(shared_new_step); packetId = 0; - return bufferBase + step * PACKET_SIZE * PACKET_PER_STEP; + return bufferBase + step * PipelineConfig::PACKET_SIZE * PipelineConfig::PACKET_PER_STEP; } return nullptr; @@ -382,9 +383,10 @@ public: __device__ __inline__ void* getNewRecvPacket() { packetId++; - if (packetId < PACKET_PER_STEP) + if (packetId < PipelineConfig::PACKET_PER_STEP) { - return bufferBase + step * PACKET_PER_STEP * PACKET_SIZE + packetId * PACKET_SIZE; + return bufferBase + step * PipelineConfig::PACKET_PER_STEP * PipelineConfig::PACKET_SIZE + + packetId * PipelineConfig::PACKET_SIZE; } __syncthreads(); @@ -401,7 +403,7 @@ public: __syncthreads(); packetId = 0; step = *(shared_new_step); - void* packetPtr = bufferBase + step * PACKET_SIZE * PACKET_PER_STEP; + void* packetPtr = bufferBase + step * PipelineConfig::PACKET_SIZE * PipelineConfig::PACKET_PER_STEP; return packetPtr; } @@ -415,14 +417,14 @@ public: } void* bufferBase; - STEP_COMMUNICATOR_TYPE* stepCommunicator; + StepCommunicatorBase* stepCommunicator; int step; int packetId; bool needRelease; int* shared_new_step; }; -template +template __global__ void allToAllMetadataDevice(int* sendExperts, int* recvExperts, float* sendScales, float* recvScales, int* localExpertStatics, int* gatheredExpertStatics, MoeCommWorkspace workspace, int* sendCountsCumsum, int* localSendIndice, int* recvCountsCumsum, int* localRecvIndice, int tokenCount, int maxTokenCountPerRank, @@ -431,22 +433,21 @@ __global__ void allToAllMetadataDevice(int* sendExperts, int* recvExperts, float bool isSender = (blockIdx.y == 0); int targetRankId = blockIdx.x; int slotCountPerRank = slotCount / rankCount; - int groupSize = topK / UNIT_SIZE; - int groupId = threadIdx.x % groupSize; + int groupSize = topK / PipelineConfig::UNIT_SIZE; __shared__ int sharedNewStep; - __align__(16) int experts[UNIT_SIZE]; - __align__(16) float scales[UNIT_SIZE]; + __align__(16) int experts[PipelineConfig::UNIT_SIZE]; + __align__(16) float scales[PipelineConfig::UNIT_SIZE]; uint8_t* bufferBase = (uint8_t*) (workspace.getFifoBasePtr(isSender, rankId, targetRankId, 0, 1)); - STEP_COMMUNICATOR_TYPE stepCommunicator(workspace.getFifoConnInfo(isSender, rankId, targetRankId, 0, rankCount, 1)); - PacketPipeline pipeline(bufferBase, &stepCommunicator, &sharedNewStep, isSender); + StepCommunicatorBase stepCommunicator(workspace.getFifoConnInfo(isSender, rankId, targetRankId, 0, rankCount, 1)); + PacketPipeline pipeline(bufferBase, &stepCommunicator, &sharedNewStep, isSender); if (isSender) { int baseCumsum = targetRankId == 0 ? 0 : *(sendCountsCumsum + targetRankId - 1); int sendTokenCount = *(sendCountsCumsum + targetRankId) - baseCumsum; - int unitCount = sendTokenCount * topK / UNIT_SIZE; + int unitCount = sendTokenCount * topK / PipelineConfig::UNIT_SIZE; void* packPtr = pipeline.getFirstSendPacket(); int indexBase = 0; @@ -457,13 +458,15 @@ __global__ void allToAllMetadataDevice(int* sendExperts, int* recvExperts, float if (threadIdx.x < UNIT_PER_ITER) { int index = indexBase + threadIdx.x; + int groupId = index % groupSize; if (index < unitCount) { int tokenId = *(localSendIndice + maxTokenCountPerRank * targetRankId + (index / groupSize)); - *((int4*) (experts)) = *(int4*) (sendExperts + tokenId * topK + groupId * UNIT_SIZE); + *((ExpertType*) (experts)) + = *(ExpertType*) (sendExperts + tokenId * topK + groupId * PipelineConfig::UNIT_SIZE); #pragma unroll - for (int j = 0; j < UNIT_SIZE; j++) + for (int j = 0; j < PipelineConfig::UNIT_SIZE; j++) { int expertId = experts[j]; if (expertId / slotCountPerRank != targetRankId) @@ -472,14 +475,15 @@ __global__ void allToAllMetadataDevice(int* sendExperts, int* recvExperts, float } } - int* expertsPtr = (int*) (packPtr) + threadIdx.x * UNIT_SIZE; - *((int4*) (expertsPtr)) = *((int4*) (experts)); + int* expertsPtr = (int*) (packPtr) + threadIdx.x * PipelineConfig::UNIT_SIZE; + *((ExpertType*) (expertsPtr)) = *((ExpertType*) (experts)); if (sendScales != nullptr) { - *((float4*) (scales)) = *(float4*) (sendScales + tokenId * topK + groupId * UNIT_SIZE); - float* scaleBasePtr = (float*) (packPtr + SCALE_OFFSET); - float* scalesPtr = (float*) (scaleBasePtr) + threadIdx.x * UNIT_SIZE; - *((float4*) (scalesPtr)) = *((float4*) (scales)); + *((ScaleType*) (scales)) + = *(ScaleType*) (sendScales + tokenId * topK + groupId * PipelineConfig::UNIT_SIZE); + float* scaleBasePtr = (float*) (packPtr + PipelineConfig::SCALE_OFFSET); + float* scalesPtr = (float*) (scaleBasePtr) + threadIdx.x * PipelineConfig::UNIT_SIZE; + *((ScaleType*) (scalesPtr)) = *((ScaleType*) (scales)); } } } @@ -488,7 +492,7 @@ __global__ void allToAllMetadataDevice(int* sendExperts, int* recvExperts, float int staticCopyIdx = threadIdx.x - UNIT_PER_ITER; if (staticCopyBase + staticCopyIdx * 4 < expertCount) { - int4* staticBasePtr = (int4*) (packPtr + STATIC_COPY_OFFSET); + int4* staticBasePtr = (int4*) (packPtr + PipelineConfig::STATIC_COPY_OFFSET); int4 staticData = *(int4*) (localExpertStatics + staticCopyBase + staticCopyIdx * 4); *(staticBasePtr + staticCopyIdx) = staticData; } @@ -521,18 +525,21 @@ __global__ void allToAllMetadataDevice(int* sendExperts, int* recvExperts, float if (threadIdx.x < packetUnitCount) { int tokenId = baseCumsum + (unitIdBase + threadIdx.x) / groupSize; - int* expertsPtr = (int*) (packetPtr) + threadIdx.x * UNIT_SIZE; - *((int4*) (experts)) = *((int4*) (expertsPtr)); - int4* dstExpertsPtr = (int4*) (recvExperts + tokenId * topK + groupId * UNIT_SIZE); - *dstExpertsPtr = *((int4*) (experts)); + int groupId = (unitIdBase + threadIdx.x) % groupSize; + int* expertsPtr = (int*) (packetPtr) + threadIdx.x * PipelineConfig::UNIT_SIZE; + *((ExpertType*) (experts)) = *((ExpertType*) (expertsPtr)); + ExpertType* dstExpertsPtr + = (ExpertType*) (recvExperts + tokenId * topK + groupId * PipelineConfig::UNIT_SIZE); + *dstExpertsPtr = *((ExpertType*) (experts)); if (recvScales != nullptr) { - float* scaleBasePtr = (float*) (packetPtr + SCALE_OFFSET); - float* scalesPtr = scaleBasePtr + threadIdx.x * UNIT_SIZE; - *((float4*) (scales)) = *((float4*) (scalesPtr)); - float4* dstScalesPtr = (float4*) (recvScales + tokenId * topK + groupId * UNIT_SIZE); - *dstScalesPtr = *((float4*) (scales)); + float* scaleBasePtr = (float*) (packetPtr + PipelineConfig::SCALE_OFFSET); + float* scalesPtr = scaleBasePtr + threadIdx.x * PipelineConfig::UNIT_SIZE; + *((ScaleType*) (scales)) = *((ScaleType*) (scalesPtr)); + ScaleType* dstScalesPtr + = (ScaleType*) (recvScales + tokenId * topK + groupId * PipelineConfig::UNIT_SIZE); + *dstScalesPtr = *((ScaleType*) (scales)); } } } @@ -541,7 +548,7 @@ __global__ void allToAllMetadataDevice(int* sendExperts, int* recvExperts, float int staticCopyIdx = threadIdx.x - UNIT_PER_ITER; if (staticCopyBase + staticCopyIdx * 4 < expertCount) { - int4* staticBasePtr = (int4*) (packetPtr + STATIC_COPY_OFFSET); + int4* staticBasePtr = (int4*) (packetPtr + PipelineConfig::STATIC_COPY_OFFSET); int4 staticData = *(staticBasePtr + staticCopyIdx); *(int4*) (gatheredExpertStatics + targetRankId * expertCount + staticCopyBase + staticCopyIdx * 4) = staticData; @@ -630,10 +637,28 @@ void allToAllMetadata(int* sendExperts, int* recvExperts, float* sendScales, flo dim3 block(block_size); dim3 grid(rankCount, 2); - allToAllMetadataDevice<<>>(sendExperts, recvExperts, sendScales, - recvScales, localExpertStatics, gatheredExpertStatics, workspace, sendCountsCumsum, localSendIndice, - recvCountsCumsum, localRecvIndice, tokenCount, maxTokenCountPerRank, topK, expertCount, slotCount, rankId, - rankCount); + if (topK % 4 == 0) + { + using PipelineConfig = PipelineConfig<4, 16>; + static_assert( + PipelineConfig::PACKET_SIZE_IN_U64 * PipelineConfig::PACKET_PER_STEP * STEP_DEPTH <= FIFO_SIZE_IN_U64, + "FIFO size is too small"); + allToAllMetadataDevice<<>>(sendExperts, recvExperts, + sendScales, recvScales, localExpertStatics, gatheredExpertStatics, workspace, sendCountsCumsum, + localSendIndice, recvCountsCumsum, localRecvIndice, tokenCount, maxTokenCountPerRank, topK, expertCount, + slotCount, rankId, rankCount); + } + else + { + using PipelineConfig = PipelineConfig<1, 64>; + static_assert( + PipelineConfig::PACKET_SIZE_IN_U64 * PipelineConfig::PACKET_PER_STEP * STEP_DEPTH <= FIFO_SIZE_IN_U64, + "FIFO size is too small"); + allToAllMetadataDevice<<>>(sendExperts, recvExperts, + sendScales, recvScales, localExpertStatics, gatheredExpertStatics, workspace, sendCountsCumsum, + localSendIndice, recvCountsCumsum, localRecvIndice, tokenCount, maxTokenCountPerRank, topK, expertCount, + slotCount, rankId, rankCount); + } int smCount = tensorrt_llm::common::getMultiProcessorCount(); memsetExpertIdsDevice<<>>( @@ -642,7 +667,7 @@ void allToAllMetadata(int* sendExperts, int* recvExperts, float* sendScales, flo size_t getMoePrepareWorkspaceSize(int epSize) { - return (STEP_DEPTH * PACKET_PER_STEP * PACKET_SIZE + StepCommunicatorBase::META_SIZE) * epSize; + return (FIFO_SIZE_IN_U64 * 8 + StepCommunicatorBase::META_SIZE) * epSize; } } // namespace moe_prepare diff --git a/cpp/tensorrt_llm/kernels/moePrepareKernels.h b/cpp/tensorrt_llm/kernels/moePrepareKernels.h index ce5a156d361b..0635397970fb 100644 --- a/cpp/tensorrt_llm/kernels/moePrepareKernels.h +++ b/cpp/tensorrt_llm/kernels/moePrepareKernels.h @@ -29,7 +29,6 @@ namespace moe_prepare { #define STEP_DEPTH 2 -#define PACKET_PER_STEP 16 #define THREADS_PER_UNIT 1 #define UNIT_PER_PIPELINE 128 #define PIPELINE_PER_CTA 4 @@ -39,21 +38,26 @@ namespace moe_prepare #define BYTES_COUNTER 8 #define CUMSUM_THREADS_PER_BLOCK 128 -#define UNIT_SIZE 4 #define UNIT_PER_ITER 256 #define STATIC_COPY_PER_ITER 128 -#define MAX_TOKEN_SIZE 8192 -static constexpr int UNIT_BYTES_SIZE = EXPERT_BYTES_PER_UNIT + SCALE_BYTES_PER_UNIT; static constexpr int THREADS_PER_PIPELINE = THREADS_PER_UNIT * UNIT_PER_PIPELINE; static constexpr int THREADS_PER_CTA = THREADS_PER_PIPELINE * PIPELINE_PER_CTA; -static constexpr int SCALE_OFFSET = UNIT_SIZE * UNIT_PER_ITER * sizeof(int); -static constexpr int STATIC_COPY_OFFSET = UNIT_SIZE * UNIT_PER_ITER * (sizeof(int) + sizeof(float)); -static constexpr int PACKET_SIZE - = UNIT_SIZE * UNIT_PER_ITER * (sizeof(int) + sizeof(float)) + STATIC_COPY_PER_ITER * 4 * sizeof(int); -static constexpr int PACKET_SIZE_IN_U64 = (PACKET_SIZE / 8); -static constexpr int FIFO_SIZE_IN_U64 = PACKET_SIZE_IN_U64 * PACKET_PER_STEP * STEP_DEPTH; +template +struct PipelineConfig +{ + static constexpr int UNIT_SIZE = UNIT_SIZE_INPUT; + static constexpr int PACKET_PER_STEP = PACKET_PER_STEP_INPUT; + static constexpr int UNIT_BYTES_SIZE = UNIT_SIZE * UNIT_PER_ITER * (sizeof(int) + sizeof(float)); + static constexpr int SCALE_OFFSET = UNIT_SIZE * UNIT_PER_ITER * sizeof(int); + static constexpr int STATIC_COPY_OFFSET = UNIT_SIZE * UNIT_PER_ITER * (sizeof(int) + sizeof(float)); + static constexpr int PACKET_SIZE = UNIT_BYTES_SIZE + STATIC_COPY_PER_ITER * 4 * sizeof(int); + static constexpr int PACKET_SIZE_IN_U64 = (PACKET_SIZE / 8); +}; + +// 1MB FIFO size +static constexpr int FIFO_SIZE_IN_U64 = 1024 * 1024 / 8; #ifdef __CUDACC__ #define ALIGN_256 __align__(256) diff --git a/cpp/tensorrt_llm/kernels/topkLastDim.cu b/cpp/tensorrt_llm/kernels/topkLastDim.cu index 285a10fd9ff9..3371ab4a0f2a 100644 --- a/cpp/tensorrt_llm/kernels/topkLastDim.cu +++ b/cpp/tensorrt_llm/kernels/topkLastDim.cu @@ -1459,13 +1459,23 @@ template size_t invokeComputeTopkLastDimWorkspaceSize( SizeType32 batchSize, SizeType32 inputLength, SizeType32 k, bool is_largest) { + using idxT = SizeType32; + size_t buf_size = 0; void* workspace = nullptr; T const* in = nullptr; T* out_val = nullptr; - SizeType32* out_idx = nullptr; - standalone_stable_radix_11bits( - workspace, buf_size, in, batchSize, inputLength, k, out_val, out_idx, is_largest, 0); + idxT* out_idx = nullptr; + + constexpr int block_dim = 512; + constexpr bool fused_last_filter = false; + constexpr bool sorted = true; + + int sm_cnt = tensorrt_llm::common::getMultiProcessorCount(); + unsigned grid_dim = air_topk_stable::calc_grid_dim(batchSize, inputLength, sm_cnt); + + standalone_stable_radix_topk_(workspace, buf_size, in, static_cast(nullptr), + batchSize, inputLength, k, out_val, out_idx, !is_largest, fused_last_filter, grid_dim, 0, sorted); return buf_size; } diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu index ac0684ff18e4..ad5cd15fddad 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu @@ -16,10 +16,21 @@ #include "DevKernel.h" +#include "cutlass/array.h" +#include "cutlass/numeric_conversion.h" +#include #include #include -#include +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Helper function for array conversion +template +__host__ __device__ constexpr static U arrayConvert(T const& input) +{ + cutlass::NumericArrayConverter converter; + return converter(input); +} //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -518,6 +529,83 @@ __global__ void finalizeKernel(KernelParams params) } } +constexpr static int FINALIZE_THREADS_PER_BLOCK = 256; + +__device__ float4 vectorizedLoadPtx(float4 const* ptr) +{ + float4 ret; + asm volatile("ld.global.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(ret.x), "=f"(ret.y), "=f"(ret.z), "=f"(ret.w) + : "l"(ptr)); + return ret; +} + +// Final kernel to unpermute and scale +// This kernel unpermutes the original data, does the k-way reduction and performs the final skip connection. + +template +__global__ void finalizeKernelVecLoad(KernelParams params) +{ + using Type = typename KernelParams::Type; + using TypeExpW = typename KernelParams::TypeExpW; + + int const hiddenDimBits = params.hiddenDim * cutlass::sizeof_bits::value; + assert(hiddenDimBits % 128 == 0); + + // Load 128-bits per thread, according to the smallest data type we read/write + constexpr int64_t FINALIZE_ELEM_PER_THREAD = 128 / cutlass::sizeof_bits::value; + using InputElem = cutlass::Array; + using OutputElem = cutlass::Array; + using ComputeElem = cutlass::Array; + + int64_t const tokenIdx = blockIdx.x; + int64_t const startOffset = threadIdx.x; + int64_t const stride = FINALIZE_THREADS_PER_BLOCK; + int64_t const numElemsInCol = params.hiddenDim / FINALIZE_ELEM_PER_THREAD; + + auto const offset = tokenIdx * params.hiddenDim; + Type* outputPtr = params.outPtr + offset; + auto* outElemPtr = reinterpret_cast(outputPtr); + auto const* inElemPtr = reinterpret_cast(params.inPtr); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // wait on primary kernel when using PDL + if constexpr (KernelParams::UsePdl) + { + cudaGridDependencySynchronize(); + } +#endif + + for (int elemIndex = startOffset; elemIndex < numElemsInCol; elemIndex += stride) + { + ComputeElem threadOutput; + threadOutput.fill(0); + for (int k = 0; k < params.topK; ++k) + { + int const expandedIdx = tokenIdx * params.topK + k; + int const permutedIdx = params.expandedIdxToPermutedIdx[expandedIdx]; + if (permutedIdx == -1) + { + continue; + } + + float const scale + = (params.expertWeightsPtr != nullptr) ? static_cast(params.expertWeightsPtr[expandedIdx]) : 1.f; + + auto const* inputPermutedPtr = inElemPtr + permutedIdx * numElemsInCol; + + float4 input = vectorizedLoadPtx(reinterpret_cast(&inputPermutedPtr[elemIndex])); + InputElem inputPermutedElem = *reinterpret_cast(&input); + ComputeElem expertResult = arrayConvert(inputPermutedElem); + + threadOutput = threadOutput + scale * expertResult; + } + + OutputElem outputElem = arrayConvert(threadOutput); + outElemPtr[elemIndex] = outputElem; + } +} + //////////////////////////////////////////////////////////////////////////////////////////////////// template @@ -552,7 +640,9 @@ __global__ void finalizeDeepSeekKernel(KernelParams params) int const expandedIdx = tokenIdx * params.topK + k; int const permutedIdx = params.expandedIdxToPermutedIdx[expandedIdx]; if (permutedIdx == -1) + { continue; + } int const totalNumPaddedTokens = params.totalNumPaddedTokens[0]; int const scaleIdx = permutedIdx + totalNumPaddedTokens * (hiddenIdx / 128); float const blockScale = params.inDqSfsPtr ? params.inDqSfsPtr[scaleIdx] : 1; @@ -591,7 +681,6 @@ __global__ void finalizeDeepSeekKernel(KernelParams params) } //////////////////////////////////////////////////////////////////////////////////////////////////// - void run(Data const& data, void* stream) { if (data.mUseDeepSeekFp8) @@ -610,9 +699,22 @@ void run(Data const& data, void* stream) int const numBlocksX = (data.hiddenDim - 1 + numThreads) / numThreads; // Capped at rather arbitrary 8192 to avoid gridDim exceeding 65535 specified by CUDA. int const numBlocksY = std::min(8192, data.numTokens); - dim3 numBlocks(numBlocksX, numBlocksY); - LAUNCH_EXPW(data, finalizeKernel, numBlocks, numThreads, 0, stream); + if (numBlocksX * numBlocksY < 1184) + { + // The number 1184 comes from 148 * 8, where 148 is the number of SMs (Streaming Multiprocessors) in the + // Blackwell architecture, + // and the value 8 means that each Streaming Multiprocessor (SM) can hold up to 8 blocks for this kernel. + // This limitation is intended to ensure that when the number of waves is greater than 1, we choose to use + // the kernel with vectorized loading. + dim3 numBlocks(numBlocksX, numBlocksY); + LAUNCH_EXPW(data, finalizeKernel, numBlocks, numThreads, 0, stream); + } + else + { + LAUNCH_EXPW(data, finalizeKernelVecLoad, /*numBlocks=*/data.numTokens, + /*numThreads=*/FINALIZE_THREADS_PER_BLOCK, 0, stream); + } } } diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index c06fda8e4943..32413eb26a29 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -413,9 +413,13 @@ class TllmGenFmhaKernel return std::make_tuple(numCtasPerSeqQ, numCtasPerSeqKv, numCtasX, numCtasY, numCtasZ, clusterDimX); } - // Compute the seqLenPerCtaKv for selecting the MLA generation kernel. - int computeSeqLenPerCtaKv(RunnerParams const& params) const + // Determine if we should use the SwapsMmaAbForGeneration kernel for MLA generation. + bool useSwapsMmaAbMlaGenKernel(RunnerParams const& params) const { + // Use the SwapsMmaAbForGeneration kernel for MLA generation when the following conditions are met: + // 1. The seqLenPerCtaKv <= 1024 based on the benchmark results (this might be fine-tuned later). + // 2. The numCtas (after splitting the heads across multiple CTAs) <= params.mMultiProcessorCount. + // The maximum number Ctas per Kv sequence, which makes sure that each CtaKv has work to do. // Here we assume the stepKv is 256. int const maxNumCtasPerSeqKv = (params.mMaxSeqLenKv + 256 - 1) / 256; @@ -427,8 +431,8 @@ class TllmGenFmhaKernel = std::min(maxNumCtasPerSeqKv, std::max(1, int32_t(params.mMultiProcessorCount / numCtas))); // Compute the seqLenPerCtaKv. int const seqLenPerCtaKv = (params.mMaxSeqLenKv + numCtasPerSeqKv - 1) / numCtasPerSeqKv; - // Return the seqLenPerCtaKv. - return seqLenPerCtaKv; + // Whether we should use the SwapsMmaAbForGeneration kernel for MLA generation. + return seqLenPerCtaKv <= 1024 && numCtas <= params.mMultiProcessorCount; } std::pair hashFromRunnerParams( @@ -442,10 +446,11 @@ class TllmGenFmhaKernel // We use the low-latency kernel (SwapsMmaAbForGeneration with tileSizeQ = 16) when any of the following // conditions are met: // 1. The number of headsQPerKv is <= 32. - // 2. The seqLenPerCtaKv <= 1024 based on the benchmark results (this might be fine-tuned later). + // 2. The seqLenPerCtaKv <= 1024 based on the benchmark results (this might be fine-tuned later) and + // the numCtas (after splitting the heads across multiple CTAs) <= params.mMultiProcessorCount. // Check the conditions. - if (params.mNumHeadsQPerKv <= 32 || computeSeqLenPerCtaKv(params) <= 1024) + if (params.mNumHeadsQPerKv <= 32 || useSwapsMmaAbMlaGenKernel(params)) { kernelType = FmhaKernelType::SwapsMmaAbForGeneration; } diff --git a/cpp/tensorrt_llm/nanobind/CMakeLists.txt b/cpp/tensorrt_llm/nanobind/CMakeLists.txt index d2e7eac20c28..aa5b3cf45daf 100755 --- a/cpp/tensorrt_llm/nanobind/CMakeLists.txt +++ b/cpp/tensorrt_llm/nanobind/CMakeLists.txt @@ -3,7 +3,22 @@ set(TRTLLM_NB_MODULE ${TRTLLM_NB_MODULE} PARENT_SCOPE) -set(SRCS ../runtime/ipcNvlsMemory.cu bindings.cpp) +set(SRCS + batch_manager/algorithms.cpp + batch_manager/bindings.cpp + batch_manager/cacheTransceiver.cpp + batch_manager/kvCacheManager.cpp + batch_manager/llmRequest.cpp + executor/bindings.cpp + executor/executor.cpp + executor/executorConfig.cpp + executor/request.cpp + runtime/bindings.cpp + testing/modelSpecBinding.cpp + runtime/moeBindings.cpp + userbuffers/bindings.cpp + ../runtime/ipcNvlsMemory.cu + bindings.cpp) include_directories(${PROJECT_SOURCE_DIR}/include) @@ -14,20 +29,29 @@ set_property(TARGET ${TRTLLM_NB_MODULE} PROPERTY POSITION_INDEPENDENT_CODE ON) target_link_directories(${TRTLLM_NB_MODULE} PUBLIC "${TORCH_INSTALL_PREFIX}/lib") +if(ENABLE_NVSHMEM) + target_link_libraries(${TRTLLM_NB_MODULE} PUBLIC nvshmem::nvshmem_host + nvshmem::nvshmem_device) +endif() + target_link_libraries( ${TRTLLM_NB_MODULE} - PUBLIC ${SHARED_TARGET} ${UNDEFINED_FLAG} ${NO_AS_NEEDED_FLAG} - ${Python3_LIBRARIES} ${TORCH_LIBRARIES} torch_python) - + PUBLIC ${SHARED_TARGET} + ${UNDEFINED_FLAG} + ${NO_AS_NEEDED_FLAG} + ${Python3_LIBRARIES} + ${TORCH_LIBRARIES} + torch_python + ${CUDA_NVML_LIB}) target_compile_definitions( ${TRTLLM_NB_MODULE} PUBLIC TRTLLM_NB_MODULE=${TRTLLM_NB_MODULE} - NB_DETAILED_ERROR_MESSAGES=1) + PYBIND11_DETAILED_ERROR_MESSAGES=1) if(NOT WIN32) set_target_properties( ${TRTLLM_NB_MODULE} PROPERTIES LINK_FLAGS - "-Wl,-rpath,'$ORIGIN/libs' -Wl,-rpath,'$ORIGIN/../nvidia/nccl/lib' ${AS_NEEDED_FLAG} ${UNDEFINED_FLAG}" + "-Wl,-rpath,'$ORIGIN/libs' -Wl,-rpath,'$ORIGIN/../nvidia/nccl/lib' -Wl,-rpath,'${CUDA_TOOLKIT_ROOT_DIR}/targets/x86_64-linux/lib/stubs' ${AS_NEEDED_FLAG} ${UNDEFINED_FLAG}" ) endif() diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp new file mode 100644 index 000000000000..e5bc7dcebf0c --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp @@ -0,0 +1,178 @@ +/* + * 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 "algorithms.h" +#include "tensorrt_llm/batch_manager/allocateKvCache.h" +#include "tensorrt_llm/batch_manager/assignReqSeqSlots.h" +#include "tensorrt_llm/batch_manager/capacityScheduler.h" +#include "tensorrt_llm/batch_manager/createNewDecoderRequests.h" +#include "tensorrt_llm/batch_manager/handleContextLogits.h" +#include "tensorrt_llm/batch_manager/handleGenerationLogits.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/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/batch_manager/runtimeBuffers.h" +#include "tensorrt_llm/batch_manager/updateDecoderBuffers.h" +#include "tensorrt_llm/nanobind/common/customCasters.h" +#include "tensorrt_llm/runtime/decoderState.h" +#include "tensorrt_llm/runtime/torch.h" +#include "tensorrt_llm/runtime/torchView.h" + +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace nb = nanobind; + +namespace tr = tensorrt_llm::runtime; +using namespace tensorrt_llm::batch_manager; + +void tensorrt_llm::nanobind::batch_manager::algorithms::initBindings(nb::module_& m) +{ + nb::class_(m, CapacityScheduler::name) + .def(nb::init(), + nb::arg("max_num_requests"), nb::arg("capacity_scheduler_policy"), nb::arg("has_kv_cache_manager"), + nb::arg("two_step_lookahead") = false, nb::arg("no_schedule_until_state") = LlmRequestState::kCONTEXT_INIT, + nb::arg("no_schedule_after_state") = LlmRequestState::kGENERATION_COMPLETE) + .def("__call__", &CapacityScheduler::operator(), nb::arg("active_requests"), + nb::arg("kv_cache_manager") = nullptr, nb::arg("peft_cache_manager") = nullptr, + nb::arg("cross_kv_cache_manager") = nullptr) + .def("name", [](CapacityScheduler const&) { return CapacityScheduler::name; }); + + nb::class_(m, MicroBatchScheduler::name) + .def(nb::init, std::optional, LlmRequestState, + LlmRequestState>(), + nb::arg("ctx_chunk_config") = std::nullopt, nb::arg("max_context_length") = std::nullopt, + nb::arg("no_schedule_until_state") = LlmRequestState::kCONTEXT_INIT, + nb::arg("no_schedule_after_state") = LlmRequestState::kGENERATION_COMPLETE) + .def("__call__", &MicroBatchScheduler::operator(), nb::arg("active_requests"), nb::arg("inflight_req_ids"), + nb::arg("max_batch_size_runtime"), nb::arg("max_num_tokens_runtime")) + .def("name", [](MicroBatchScheduler const&) { return MicroBatchScheduler::name; }); + + nb::class_(m, PauseRequests::name) + .def(nb::init(), nb::arg("max_input_len")) + .def("__call__", &PauseRequests::operator(), nb::arg("requests_to_pause"), nb::arg("inflight_req_ids"), + nb::arg("req_ids_to_pause"), nb::arg("pause_flagged"), nb::arg("seq_slot_manager"), + nb::arg("kv_cache_manager") = std::nullopt, nb::arg("cross_kv_cache_manager") = std::nullopt, + nb::arg("peft_cache_manager") = std::nullopt) + .def("name", [](PauseRequests const&) { return PauseRequests::name; }); + + nb::class_(m, AssignReqSeqSlots::name) + .def(nb::init<>()) + .def("__call__", &AssignReqSeqSlots::operator(), nb::arg("seq_slot_manager"), nb::arg("context_requests"), + nb::arg("generation_requests")) + .def("name", [](AssignReqSeqSlots const&) { return AssignReqSeqSlots::name; }); + + nb::class_(m, AllocateKvCache::name) + .def(nb::init<>()) + .def("__call__", &AllocateKvCache::operator(), nb::arg("kv_cache_manager"), nb::arg("context_requests"), + nb::arg("generation_requests"), nb::arg("model_config"), nb::arg("cross_kv_cache_manager") = std::nullopt) + .def("name", [](AllocateKvCache const&) { return AllocateKvCache::name; }); + + nb::class_(m, HandleContextLogits::name) + .def(nb::init<>()) + .def( + "__call__", + [](HandleContextLogits const& self, DecoderInputBuffers& inputBuffers, RequestVector const& contextRequests, + at::Tensor const& logits, std::vector const& numContextLogitsVec, + tr::ModelConfig const& modelConfig, tr::BufferManager const& manager, + OptionalRef medusaBuffers = std::nullopt) + { + return self(inputBuffers, contextRequests, tr::TorchView::of(logits), numContextLogitsVec, modelConfig, + manager, medusaBuffers); + }, + nb::arg("decoder_input_buffers"), nb::arg("context_requests"), nb::arg("logits"), + nb::arg("num_context_logits"), nb::arg("model_config"), nb::arg("buffer_manager"), + nb::arg("medusa_buffers") = std::nullopt) + .def("name", [](HandleContextLogits const&) { return HandleContextLogits::name; }); + + nb::class_(m, HandleGenerationLogits::name) + .def(nb::init<>()) + .def( + "__call__", + [](HandleGenerationLogits const& self, DecoderInputBuffers& inputBuffers, + RequestVector const& generationRequests, at::Tensor const& logits, tr::SizeType32 logitsIndex, + tr::ModelConfig const& modelConfig, tr::BufferManager const& manager, + OptionalRef genRuntimeBuffers = std::nullopt, + OptionalRef medusaBuffers = std::nullopt) + { + self(inputBuffers, generationRequests, tr::TorchView::of(logits), logitsIndex, modelConfig, manager, + genRuntimeBuffers, medusaBuffers); + }, + nb::arg("decoder_input_buffers"), nb::arg("generation_requests"), nb::arg("logits"), + nb::arg("logits_index"), nb::arg("model_config"), nb::arg("buffer_manager"), + nb::arg("gen_runtime_buffers") = std::nullopt, nb::arg("medusa_buffers") = std::nullopt) + .def("name", [](HandleGenerationLogits const&) { return HandleGenerationLogits::name; }); + + nb::class_(m, MakeDecodingBatchInputOutput::name) + .def(nb::init<>()) + .def("__call__", &MakeDecodingBatchInputOutput::operator(), nb::arg("decoder_input_buffers"), + nb::arg("decoder_state"), nb::arg("model_config"), nb::arg("max_num_sequences"), + nb::arg("fused_runtime_buffers") = std::nullopt) + .def("name", [](MakeDecodingBatchInputOutput const&) { return MakeDecodingBatchInputOutput::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")) + .def( + "__call__", + [](CreateNewDecoderRequests& self, tr::ModelConfig const& modelConfig, tr::WorldConfig const& worldConfig, + executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, + tr::BufferManager const& bufferManager, nvinfer1::DataType logitsType, + DecoderInputBuffers& inputBuffers, runtime::decoder::DecoderState& decoderState, + tensorrt_llm::runtime::CudaStream const& runtimeStream, + tensorrt_llm::runtime::CudaStream const& decoderStream, SizeType32 maxSequenceLength, + SizeType32 beamWidth, OptionalRef medusaBuffers = std::nullopt) + { + auto [batchSlots, samplingConfigs, lookaheadPrompt, lookaheadAlgoConfigs] = self(modelConfig, + worldConfig, decodingConfig, contextRequests, bufferManager, logitsType, inputBuffers, decoderState, + runtimeStream, decoderStream, maxSequenceLength, beamWidth, medusaBuffers); + + return std::tuple{runtime::Torch::tensor(batchSlots), std::move(samplingConfigs), + std::move(lookaheadPrompt), std::move(lookaheadAlgoConfigs)}; + }, + nb::arg("model_config"), nb::arg("world_config"), nb::arg("decoding_config"), nb::arg("context_requests"), + nb::arg("buffer_manager"), nb::arg("logits_type"), nb::arg("decoder_input_buffers"), + nb::arg("decoder_state"), nb::arg("runtime_stream"), nb::arg("decoder_stream"), + nb::arg("max_sequence_length"), nb::arg("beam_width"), nb::arg("medusa_buffers") = std::nullopt) + .def("name", [](CreateNewDecoderRequests const&) { return CreateNewDecoderRequests::name; }); + + nb::class_(m, UpdateDecoderBuffers::name) + .def(nb::init<>()) + .def("__call__", &UpdateDecoderBuffers::operator(), nb::arg("model_config"), nb::arg("decoder_output_buffers"), + nb::arg("copy_buffer_manager"), nb::arg("decoder_state"), nb::arg("return_log_probs"), + nb::arg("decoder_finish_event")) + .def("name", [](UpdateDecoderBuffers const&) { return UpdateDecoderBuffers::name; }); +} diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.h b/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.h new file mode 100644 index 000000000000..cac81d73f275 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.h @@ -0,0 +1,29 @@ +/* + * 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::batch_manager::algorithms +{ + +void initBindings(nb::module_& m); + +} diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp new file mode 100644 index 000000000000..151b33b11953 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp @@ -0,0 +1,525 @@ +/* + * 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 "bindings.h" +#include "tensorrt_llm/nanobind/common/customCasters.h" + +#include "tensorrt_llm/batch_manager/common.h" +#include "tensorrt_llm/batch_manager/decoderBuffers.h" +#include "tensorrt_llm/batch_manager/medusaBuffers.h" +#include "tensorrt_llm/batch_manager/microBatchScheduler.h" +#include "tensorrt_llm/batch_manager/peftCacheManager.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/nanobind/common/bindTypes.h" +#include "tensorrt_llm/runtime/gptDecoderBatched.h" +#include "tensorrt_llm/runtime/runtimeKernels.h" +#include "tensorrt_llm/runtime/torch.h" +#include "tensorrt_llm/runtime/torchView.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace nb = nanobind; +namespace tb = tensorrt_llm::batch_manager; +namespace tle = tensorrt_llm::executor; +namespace tr = tensorrt_llm::runtime; + +using namespace tensorrt_llm::runtime; + +namespace tensorrt_llm::nanobind::batch_manager +{ + +void initBindings(nb::module_& m) +{ + using GenLlmReq = tb::GenericLlmRequest; + + // Create and register exceptions in module scope + static nb::object peft_exc = nb::exception(m, "PeftTaskNotCachedException"); + static nb::object lora_exc = nb::exception(m, "LoraCacheFullException"); + + // Register with no captures + nb::register_exception_translator( + [](std::exception_ptr const& p, void*) + { + try + { + if (p) + std::rethrow_exception(p); + } + catch (const tb::PeftTaskNotCachedException& e) + { + PyErr_SetString(peft_exc.ptr(), e.what()); + } + catch (const tr::LoraCacheFullException& e) + { + PyErr_SetString(lora_exc.ptr(), e.what()); + } + }); + + NanobindUtils::bindSet(m, "ReqIdsSet"); + + nb::enum_(m, "LlmRequestType") + .value("LLMREQUEST_TYPE_CONTEXT_AND_GENERATION", tb::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION) + .value("LLMREQUEST_TYPE_CONTEXT_ONLY", tb::LLMREQUEST_TYPE_CONTEXT_ONLY) + .value("LLMREQUEST_TYPE_GENERATION_ONLY", tb::LLMREQUEST_TYPE_GENERATION_ONLY) + .export_values(); + + nb::class_(m, "ContextChunkingConfig") + .def(nb::init(), nb::arg("chunking_policy"), + nb::arg("chunk_unit_size")) + .def_rw("chunking_policy", &tb::batch_scheduler::ContextChunkingConfig::chunkingPolicy) + .def_rw("chunk_unit_size", &tb::batch_scheduler::ContextChunkingConfig::chunkUnitSize); + + nb::class_(m, "GenericLlmRequest") + .def("set_exclude_input_from_output", &GenLlmReq::setExcludeInputFromOutput, nb::arg("exclude")) + .def("get_num_tokens", &GenLlmReq::getNumTokens, nb::arg("beam")) + .def_prop_ro("max_beam_num_tokens", &GenLlmReq::getMaxBeamNumTokens) + .def("get_token", &GenLlmReq::getToken, nb::arg("beam"), nb::arg("pos")) + .def("get_tokens", nb::overload_cast(&GenLlmReq::getTokens, nb::const_), nb::arg("beam")) + .def("get_tokens", nb::overload_cast<>(&GenLlmReq::getTokens, nb::const_)) + .def("get_last_tokens", nb::overload_cast(&GenLlmReq::getLastTokens), nb::arg("beam")) + .def("get_last_tokens", nb::overload_cast<>(&GenLlmReq::getLastTokens)) + .def("get_beam_width_by_iter", &GenLlmReq::getBeamWidthByIter, nb::arg("for_next_iteration") = false) + .def_prop_ro("max_num_generated_tokens", &GenLlmReq::getMaxNumGeneratedTokens) + .def("add_new_token", &GenLlmReq::addNewToken, nb::arg("token"), nb::arg("beam")) + .def("add_new_tokens", &GenLlmReq::addNewTokens, nb::arg("beam_tokens")) + .def_prop_ro("num_draft_tokens", &GenLlmReq::getNumDraftTokens) + .def("set_generated_tokens", &GenLlmReq::setGeneratedTokens, nb::arg("generated_beam_tokens")) + .def("pause", &GenLlmReq::pause, nb::arg("max_input_len")) + .def_prop_rw("max_sent_token_len", &GenLlmReq::getMaxSentTokenLen, &GenLlmReq::setMaxSentTokenLen) + .def_prop_ro("prompt_embedding_table", &GenLlmReq::getPromptEmbeddingTable) + .def_prop_ro("multimodal_embedding", &GenLlmReq::getMultimodalEmbedding) + .def_prop_ro("mrope_rotary_cos_sin", &GenLlmReq::getMropeRotaryCosSin) + .def_prop_ro("bad_words_list", &GenLlmReq::getBadWordsList) + .def_prop_rw("draft_logits", &GenLlmReq::getDraftLogits, &GenLlmReq::setDraftLogits) + .def_prop_ro("embedding_bias", &GenLlmReq::getEmbeddingBias) + .def_prop_rw("lora_config", &GenLlmReq::getLoraConfig, &GenLlmReq::setLoraConfig) + .def_prop_rw("lora_weights", &GenLlmReq::getLoraWeights, &GenLlmReq::setLoraWeights) + .def_prop_ro("stop_words_list", &GenLlmReq::getStopWordsList) + .def_prop_ro("context_logits", &GenLlmReq::getContextLogitsHost) + .def_prop_ro("generation_logits", &GenLlmReq::getGenerationLogitsHost) + .def_prop_ro("prompt_vocab_size", &GenLlmReq::getPromptVocabSize) + .def_prop_ro("mrope_position_deltas", &GenLlmReq::getMropePositionDeltas) + .def_prop_ro("lora_task_id", &GenLlmReq::getLoraTaskId) + .def_prop_ro("lookahead_config", &GenLlmReq::getLookaheadConfig) + .def_prop_rw("context_chunk_size", &GenLlmReq::getContextChunkSize, &GenLlmReq::setContextChunkSize) + .def_prop_rw("decoding_iter", &GenLlmReq::getDecodingIter, &GenLlmReq::setDecodingIter) + .def_rw("request_id", &GenLlmReq::mRequestId) + .def_rw("prompt_len", &GenLlmReq::mPromptLen) + .def_rw("max_new_tokens", &GenLlmReq::mMaxNewTokens) + .def_rw("sampling_config", &GenLlmReq::mSamplingConfig) + .def_prop_rw("state", &GenLlmReq::getState, &GenLlmReq::setState) + .def_prop_rw("streaming", &GenLlmReq::isStreaming, &GenLlmReq::setStreaming) + .def_rw("end_id", &GenLlmReq::mEndId) + .def_rw("pad_id", &GenLlmReq::mPadId) + .def_rw("seq_slot", &GenLlmReq::mSeqSlot) + .def_prop_ro("return_log_probs", &GenLlmReq::returnLogProbs) + .def_prop_ro("return_context_logits", &GenLlmReq::getReturnContextLogits) + .def_prop_ro("return_generation_logits", &GenLlmReq::getReturnGenerationLogits) + .def_prop_ro("log_probs", nb::overload_cast<>(&GenLlmReq::getLogProbs, nb::const_)) + .def("get_log_probs", nb::overload_cast(&GenLlmReq::getLogProbs, nb::const_)) + .def("set_log_probs", &GenLlmReq::setLogProbs, nb::arg("log_probs"), nb::arg("beam")) + .def("set_return_encoder_output", &GenLlmReq::setReturnEncoderOutput, nb::arg("return_encoder_output")) + .def("get_return_encoder_output", &GenLlmReq::getReturnEncoderOutput) + .def("priority", nb::overload_cast<>(&GenLlmReq::priority, nb::const_)) + .def("set_priority", nb::overload_cast(&GenLlmReq::setPriority)) + .def_prop_ro("cum_log_probs", &GenLlmReq::getCumLogProbs) + .def("set_cum_log_prob", &GenLlmReq::setCumLogProb, nb::arg("cum_log_prob"), nb::arg("beam")) + .def("update_num_tokens_per_iteration", &GenLlmReq::updateNumTokensPerIteration, + nb::arg("num_tokens_per_iteration"), nb::arg("model_config")) + .def_prop_ro("orig_prompt_len", &GenLlmReq::getOrigPromptLen) + .def("has_draft_tokens", &GenLlmReq::hasDraftTokens) + .def("move_to_next_context_chunk", &GenLlmReq::moveToNextContextChunk) + .def_prop_ro("is_last_context_chunk", &GenLlmReq::isLastContextChunk) + .def_prop_ro("is_first_context_chunk", &GenLlmReq::isFirstContextChunk) + .def_prop_ro("context_remaining_length", &GenLlmReq::getContextRemainingLength) + .def_prop_ro("context_logits", &GenLlmReq::getContextLogitsHost) + .def_prop_ro("num_draft_tokens", &GenLlmReq::getNumDraftTokens) + .def("set_finished_reason", &GenLlmReq::setFinishedReason, nb::arg("finish_reason"), nb::arg("beam")) + .def_prop_ro("is_finished", &GenLlmReq::isFinished) + .def_prop_ro("is_finished_due_to_length", &GenLlmReq::isFinishedDueToLength) + .def_prop_rw( + "context_current_position", &GenLlmReq::getContextCurrentPosition, &GenLlmReq::setContextCurrentPosition) + .def_prop_ro("prepopulated_prompt_len", &GenLlmReq::getPrepopulatedPromptLen) + .def_prop_rw("guided_decoding_params", &GenLlmReq::getGuidedDecodingParams, &GenLlmReq::setGuidedDecodingParams) + .def_prop_ro("context_phase_params", &GenLlmReq::getContextPhaseParams) + .def_prop_ro("is_context_only_request", &GenLlmReq::isContextOnlyRequest) + .def_prop_ro("is_generation_only_request", &GenLlmReq::isGenerationOnlyRequest) + .def_prop_ro("is_generation_complete_state", &GenLlmReq::isGenerationCompleteState) + .def_prop_ro("is_context_finished", &GenLlmReq::isContextFinished) + .def_prop_ro("is_disagg_generation_init_state", &GenLlmReq::isDisaggGenerationInitState) + .def_prop_ro("is_disagg_generation_transmission_complete", &GenLlmReq::isDisaggGenerationTransmissionComplete) + .def_prop_ro( + "is_disagg_generation_transmission_in_progress", &GenLlmReq::isDisaggGenerationTransmissionInProgress) + .def_prop_ro("is_context_init_state", &GenLlmReq::isContextInitState) + .def_prop_ro("is_generation_in_progress_state", &GenLlmReq::isGenerationInProgressState) + .def_prop_ro("is_disagg_context_transmission_state", &GenLlmReq::isDisaggContextTransmissionState) + .def_prop_ro("is_disagg_context_complete_state", &GenLlmReq::isDisaggContextCompleteState) + .def_prop_ro("stage", &GenLlmReq::getRequestStage) + .def_prop_ro("kv_cache_transfer_time_ms", &GenLlmReq::getKvCacheTransferTimeMS) + .def_prop_ro("kv_cache_size", &GenLlmReq::getKvCacheSize) + .def_prop_ro("avg_decoded_tokens_per_iter", &GenLlmReq::getAvgDecodedTokensPerIter) + .def_prop_ro("alloc_total_blocks", &GenLlmReq::getAllocTotalBlocksPerRequest) + .def_prop_ro("alloc_new_blocks", &GenLlmReq::getAllocNewBlocksPerRequest) + .def("alloc_context_logits", &GenLlmReq::allocContextLogitsHost, nb::arg("vocab_size"), nb::arg("logit_dtype")) + .def_prop_ro("reused_blocks", &GenLlmReq::getReusedBlocksPerRequest) + .def_prop_ro("missed_blocks", &GenLlmReq::getMissedBlocksPerRequest) + .def_prop_ro("kv_cache_hit_rate", &GenLlmReq::getKVCacheHitRatePerRequest) + .def_prop_ro("llm_request_type", &GenLlmReq::getLlmRequestType) + .def_prop_ro("multimodal_hashes", + [](GenLlmReq& self) + { + std::optional>> hashes = std::nullopt; + if (self.getMultimodalHashes()) + { + hashes = *self.getMultimodalHashes().value(); + } + return hashes; + }) + .def_prop_ro("multimodal_positions", + [](GenLlmReq& self) + { + std::optional> positions = std::nullopt; + if (self.getMultimodalPositions()) + { + positions = *self.getMultimodalPositions().value(); + } + return positions; + }) + .def_prop_ro("multimodal_lengths", + [](GenLlmReq& self) + { + std::optional> lengths = std::nullopt; + if (self.getMultimodalLengths()) + { + lengths = *self.getMultimodalLengths().value(); + } + return lengths; + }) + .def_prop_ro("position_ids", + [](GenLlmReq& self) + { + std::optional> positionIds = std::nullopt; + if (self.getPositionIds()) + { + positionIds = *self.getPositionIds().value(); + } + return positionIds; + }) + .def_prop_rw( + "draft_tokens", + [](GenLlmReq& self) + { + std::optional draftTokens = std::nullopt; + if (self.hasDraftTokens()) + { + draftTokens = *self.getDraftTokens(); + } + return draftTokens; + }, + [](GenLlmReq& self, std::optional const& draftTokens) + { + if (draftTokens) + { + self.setDraftTokens(std::make_shared(draftTokens.value())); + } + }) + .def_prop_rw("is_dummy_request", &GenLlmReq::isDummyRequest, &GenLlmReq::setIsDummyRequest) + .def_prop_ro("return_perf_metrics", &GenLlmReq::getReturnPerfMetrics); + + nb::class_(m, "LlmRequest", nb::dynamic_attr()) + .def( + "__init__", + [](tb::LlmRequest* self, tb::LlmRequest::RequestIdType request_id, + tb::LlmRequest::SizeType32 max_new_tokens, std::vector input_tokens, + runtime::SamplingConfig sampling_config, bool is_streaming, + std::optional end_id, std::optional pad_id, + std::optional embedding_bias, std::optional bad_words_list, + std::optional stop_words_list, + std::optional> position_ids, + std::optional prompt_embedding_table, + std::optional prompt_vocab_size, + std::optional>> multimodal_hashes, + std::optional> multimodal_positions, + std::optional> multimodal_lengths, + std::optional multimodal_embedding, std::optional mrope_rotary_cos_sin, + std::optional mrope_position_deltas, + std::optional lora_task_id, std::optional lora_weights, + std::optional lora_config, + std::optional lookahead_config, + std::optional kv_cache_retention_config, bool return_log_probs, + bool return_context_logits, bool return_generation_logits, + std::optional draft_tokens, std::optional draft_logits, + bool exclude_input_from_output, + std::optional logits_post_processor, + bool apply_logits_post_processor_batched, std::optional encoder_input_tokens, + bool return_encoder_output, std::optional client_id, + executor::PriorityType priority, std::optional encoder_input_features, + std::optional encoder_output_length, + std::optional cross_attention_mask, tb::LlmRequestType llm_request_type, + std::optional input_token_extra_ids, + tb::LlmRequest::SizeType32 num_return_sequences, std::optional eagle_config, + std::optional skip_cross_attn_blocks, bool return_perf_metrics, + std::optional guided_decoding_params, + std::optional language_adapter_uid, + std::optional allotted_time_ms, + std::optional context_phase_params) + { + auto makeOptionalTensor = [](std::optional const& atTensor, bool unsqueeze = false) + { + std::optional tensorPtr = std::nullopt; + if (atTensor) + { + tensorPtr = tr::TorchView::of(atTensor.value()); + if (unsqueeze) + { + (*tensorPtr)->unsqueeze(0); + } + } + return tensorPtr; + }; + + auto embedding_bias_tensor_ptr = makeOptionalTensor(embedding_bias, true); + auto bad_words_list_tensor_ptr = makeOptionalTensor(bad_words_list, true); + auto stop_words_list_tensor_ptr = makeOptionalTensor(stop_words_list, true); + auto prompt_embedding_table_tensor_ptr = makeOptionalTensor(prompt_embedding_table); + auto multimodal_embedding_tensor_ptr = makeOptionalTensor(multimodal_embedding); + auto lora_weights_tensor_ptr = makeOptionalTensor(lora_weights); + auto mrope_rotary_cos_sin_tensor_ptr = makeOptionalTensor(mrope_rotary_cos_sin); + auto lora_config_tensor_ptr = makeOptionalTensor(lora_config); + auto draft_logits_tensor_ptr = makeOptionalTensor(draft_logits); + auto encoder_input_features_tensor_ptr = makeOptionalTensor(encoder_input_features); + auto cross_attention_mask_tensor_ptr = makeOptionalTensor(cross_attention_mask); + auto skip_cross_attn_blocks_tensor_ptr = makeOptionalTensor(skip_cross_attn_blocks); + + // 49 parameters + new (self) tb::LlmRequest{request_id, max_new_tokens, input_tokens, sampling_config, is_streaming, + end_id, pad_id, embedding_bias_tensor_ptr, bad_words_list_tensor_ptr, stop_words_list_tensor_ptr, + position_ids, prompt_embedding_table_tensor_ptr, prompt_vocab_size, multimodal_hashes, + multimodal_positions, multimodal_lengths, multimodal_embedding_tensor_ptr, + mrope_rotary_cos_sin_tensor_ptr, mrope_position_deltas, lora_task_id, lora_weights_tensor_ptr, + lora_config_tensor_ptr, lookahead_config, kv_cache_retention_config, return_log_probs, + return_context_logits, return_generation_logits, draft_tokens, draft_logits_tensor_ptr, + exclude_input_from_output, logits_post_processor, apply_logits_post_processor_batched, + encoder_input_tokens, return_encoder_output, client_id, priority, encoder_input_features_tensor_ptr, + encoder_output_length, cross_attention_mask_tensor_ptr, llm_request_type, input_token_extra_ids, + num_return_sequences, eagle_config, skip_cross_attn_blocks_tensor_ptr, return_perf_metrics, + guided_decoding_params, language_adapter_uid, allotted_time_ms, context_phase_params}; + }, + nb::arg("request_id"), nb::arg("max_new_tokens"), nb::arg("input_tokens"), nb::arg("sampling_config"), + nb::arg("is_streaming"), nb::arg("end_id") = std::nullopt, nb::arg("pad_id") = std::nullopt, + nb::arg("embedding_bias") = std::nullopt, nb::arg("bad_words_list") = std::nullopt, + nb::arg("stop_words_list") = std::nullopt, nb::arg("position_ids") = std::nullopt, + nb::arg("prompt_embedding_table") = std::nullopt, nb::arg("prompt_vocab_size") = std::nullopt, + nb::arg("multimodal_hashes") = std::nullopt, nb::arg("multimodal_positions") = std::nullopt, + nb::arg("multimodal_lengths") = std::nullopt, nb::arg("multimodal_embedding") = std::nullopt, + nb::arg("mrope_rotary_cos_sin") = std::nullopt, nb::arg("mrope_position_deltas") = std::nullopt, + nb::arg("lora_task_id") = std::nullopt, nb::arg("lora_weights") = std::nullopt, + nb::arg("lora_config") = std::nullopt, nb::arg("lookahead_config") = std::nullopt, + nb::arg("kv_cache_retention_config") = std::nullopt, nb::arg("return_log_probs") = false, + nb::arg("return_context_logits") = false, nb::arg("return_generation_logits") = false, + nb::arg("draft_tokens") = std::nullopt, nb::arg("draft_logits") = std::nullopt, + nb::arg("exclude_input_from_output") = false, nb::arg("logits_post_processor") = std::nullopt, + nb::arg("apply_logits_post_processor_batched") = false, nb::arg("encoder_input_tokens") = std::nullopt, + nb::arg("return_encoder_output") = false, nb::arg("client_id") = std::nullopt, + nb::arg("priority") = executor::Request::kDefaultPriority, nb::arg("encoder_input_features") = std::nullopt, + nb::arg("encoder_output_len") = std::nullopt, nb::arg("cross_attention_mask") = std::nullopt, + nb::arg("llm_request_type") = tb::LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, + nb::arg("input_token_extra_ids") = std::nullopt, nb::arg("num_return_sequences") = 1, + nb::arg("eagle_config") = std::nullopt, nb::arg("skip_cross_attn_blocks") = std::nullopt, + nb::arg("return_perf_metrics") = false, nb::arg("guided_decoding_params") = std::nullopt, + nb::arg("language_adapter_uid") = std::nullopt, nb::arg("allotted_time_ms") = std::nullopt, + nb::arg("context_phase_params") = std::nullopt) + .def("validate", &tb::LlmRequest::validate, nb::arg("max_input_len"), nb::arg("max_seq_len"), + nb::arg("max_draft_len"), nb::arg("vocab_size_padded"), nb::arg("max_endocer_input_len") = std::nullopt, + nb::arg("enable_kv_cache_reuse") = false) + .def("create_response", &tb::LlmRequest::createResponse, nb::arg("use_fast_logits") = false, + nb::arg("mpi_world_rank") = 0) + .def("create_result", &tb::LlmRequest::createResult, nb::arg("use_fast_logits") = false, + nb::arg("mpi_world_rank") = 0) + .def("create_serialized_result", + [](tb::LlmRequest& self, bool use_fast_logits = false, int mpi_world_rank = 0) + { + std::vector serialized_result; + bool is_final = false; + self.createSerializedResult(serialized_result, is_final, use_fast_logits, mpi_world_rank); + return std::make_tuple(nb::bytes(serialized_result.data(), serialized_result.size()), is_final); + }) + .def("move_prompt_embedding_table_to_gpu", &tb::LlmRequest::movePromptEmbeddingTableToGpu, nb::arg("manager")) + .def("move_lora_weights_to_gpu", &tb::LlmRequest::moveLoraWeightsToGpu, nb::arg("manager")) + .def("finish_by_reason", &tb::LlmRequest::finishByReason, nb::arg("finish_reason")) + .def("set_first_scheduled_time", &tb::LlmRequest::setFirstScheduledTime) + .def("update_perf_metrics", &tb::LlmRequest::updatePerfMetrics, nb::arg("iter_counter")); + + nb::class_(m, "SequenceSlotManager") + .def(nb::init(), nb::arg("max_num_slots"), + nb::arg("max_sequence_idle_microseconds")) + .def("get_sequence_slot", &tb::SequenceSlotManager::getSequenceSlot, nb::arg("start_flag"), + nb::arg("sequence_id")) + .def("free_sequence_slot", &tb::SequenceSlotManager::freeSequenceSlot, nb::arg("sequence_id")) + .def("free_idle_sequence_slots", &tb::SequenceSlotManager::freeIdleSequenceSlots); + + nb::class_(m, "RnnStateManager") + .def(nb::init(), + nb::arg("max_num_sequences"), nb::arg("model_config"), nb::arg("world_config"), nb::arg("buffer_manager")); + + nb::class_(m, "DecoderInputBuffers") + .def(nb::init(), nb::arg("max_batch_size"), + nb::arg("max_tokens_per_engine_step"), nb::arg("manager")) + .def_rw("setup_batch_slots", &tb::DecoderInputBuffers::setupBatchSlots) + .def_rw("setup_batch_slots_device", &tb::DecoderInputBuffers::setupBatchSlotsDevice) + .def_rw("fill_values", &tb::DecoderInputBuffers::fillValues) + .def_rw("fill_values_device", &tb::DecoderInputBuffers::fillValuesDevice) + .def_rw("inputs_ids", &tb::DecoderInputBuffers::inputsIds) + .def_rw("forward_batch_slots", &tb::DecoderInputBuffers::forwardBatchSlots) + .def_rw("logits", &tb::DecoderInputBuffers::logits) + .def_rw("decoder_requests", &tb::DecoderInputBuffers::decoderRequests); + + nb::class_(m, "DecoderOutputBuffers") + .def_rw("sequence_lengths_host", &tb::DecoderOutputBuffers::sequenceLengthsHost) + .def_rw("finished_sum_host", &tb::DecoderOutputBuffers::finishedSumHost) + .def_prop_ro("new_output_tokens_host", + [](tb::DecoderOutputBuffers& self) { return tr::Torch::tensor(self.newOutputTokensHost); }) + .def_rw("cum_log_probs_host", &tb::DecoderOutputBuffers::cumLogProbsHost) + .def_rw("log_probs_host", &tb::DecoderOutputBuffers::logProbsHost) + .def_rw("finish_reasons_host", &tb::DecoderOutputBuffers::finishReasonsHost); + + nb::class_(m, "SlotDecoderBuffers") + .def(nb::init(), + nb::arg("max_beam_width"), nb::arg("max_seq_len"), nb::arg("buffer_manager")) + .def_rw("output_ids", &tb::SlotDecoderBuffers::outputIds) + .def_rw("output_ids_host", &tb::SlotDecoderBuffers::outputIdsHost) + .def_rw("sequence_lengths_host", &tb::SlotDecoderBuffers::sequenceLengthsHost) + .def_rw("cum_log_probs", &tb::SlotDecoderBuffers::cumLogProbs) + .def_rw("cum_log_probs_host", &tb::SlotDecoderBuffers::cumLogProbsHost) + .def_rw("log_probs", &tb::SlotDecoderBuffers::logProbs) + .def_rw("log_probs_host", &tb::SlotDecoderBuffers::logProbsHost) + .def_rw("finish_reasons_host", &tb::SlotDecoderBuffers::finishReasonsHost); + + nb::class_(m, "MedusaBuffers") + .def(nb::init(), + nb::arg("max_beam_width"), nb::arg("max_seq_len"), nb::arg("buffer_manager"), nb::arg("model_config"), + nb::arg("world_config"), nb::arg("decoding_config"), nb::arg("runtime")); + + m.def( + "add_new_tokens_to_requests", + [](std::vector>& requests, + std::vector const& tokens, int beam_idx) + { + TLLM_CHECK_WITH_INFO(requests.size() == tokens.size(), "Expected the same number of requests and tokens."); + + for (int i = 0; i < requests.size(); ++i) + { + requests[i]->addNewToken(tokens[i], beam_idx); + } + }, + nb::arg("requests"), nb::arg("tokens"), nb::arg("beam_idx"), + "Add new tokens to multiple LLM requests. The tokens vector should contain tokens for beam beam_idx of all " + "requests in order."); + + m.def( + "make_decoding_batch_input", + [](std::vector>& contextRequests, + std::vector>& genRequests, tr::ITensor::SharedPtr logits, int beamWidth, + std::vector const& numContextLogitsPrefixSum, tb::DecoderInputBuffers const& decoderInputBuffers, + runtime::decoder::DecoderState& decoderState, tr::BufferManager const& manager) + { + std::vector activeSlots; + std::vector generationSteps; + std::vector> logitsVec = {{}}; + + for (int i = 0; i < contextRequests.size(); ++i) + { + if (contextRequests[i]->isLastContextChunk()) + { + activeSlots.push_back(*contextRequests[i]->mSeqSlot); + generationSteps.push_back(contextRequests[i]->getDecodingIter()); + auto contextLogitsOffset = numContextLogitsPrefixSum[i + 1] - 1; + tr::ITensor::SharedPtr logitsView = ITensor::slice(logits, contextLogitsOffset, 1); + + if (beamWidth > 1) + { + // Tile logits of context requests + auto const logitsShape = logitsView->getShape(); + auto const logitsType = logitsView->getDataType(); + auto decoderLogits = manager.gpu(ITensor::makeShape({beamWidth, logitsShape.d[1]}), logitsType); + tensorrt_llm::runtime::kernels::tileTensor( + *decoderLogits, *logitsView, beamWidth, manager.getStream()); + decoderLogits->unsqueeze(0); + logitsVec[0].push_back(std::move(decoderLogits)); + } + else + { + logitsView->unsqueeze(1); + logitsVec[0].push_back(std::move(logitsView)); + } + } + } + + auto genLogitsOffset = numContextLogitsPrefixSum.back(); + for (int i = 0; i < genRequests.size(); ++i) + { + if (genRequests[i]->isGenerationInProgressState()) + { + activeSlots.push_back(*genRequests[i]->mSeqSlot); + generationSteps.push_back(genRequests[i]->getDecodingIter()); + + auto logitsOffset = genLogitsOffset + i * beamWidth; + auto numberOfLogits = beamWidth; + tr::ITensor::SharedPtr logitsView = ITensor::slice(logits, logitsOffset, numberOfLogits); + logitsView->unsqueeze(0); + logitsVec[0].push_back(std::move(logitsView)); + } + } + + auto& batchSlots = decoderInputBuffers.forwardBatchSlots; + batchSlots[0]->resize(activeSlots.size()); + auto batchSlotsRange = tr::BufferRange(*batchSlots[0]); + for (int i = 0; i < activeSlots.size(); ++i) + { + batchSlotsRange[i] = activeSlots[i]; + } + + auto decodingInput = std::make_unique(logitsVec, 1); + decodingInput->batchSlots = batchSlots; + + auto const maxBeamWidth = decoderState.getMaxBeamWidth(); + if (maxBeamWidth > 1) + { + // For Variable-Beam-Width-Search + decoderState.getJointDecodingInput().generationSteps = generationSteps; + } + + return decodingInput; + }, + nb::arg("context_requests"), nb::arg("generation_requests"), nb::arg("logits"), nb::arg("beam_width"), + nb::arg("num_context_logits_prefix_sum"), nb::arg("decoder_input_buffers"), nb::arg("decoder_state"), + nb::arg("buffer_manager"), "Make decoding batch input."); +} + +} // namespace tensorrt_llm::nanobind::batch_manager diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.h b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.h new file mode 100644 index 000000000000..3d5a0f5d5b2b --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.h @@ -0,0 +1,28 @@ +/* + * 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::batch_manager +{ + +void initBindings(nb::module_& m); + +} diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp new file mode 100644 index 000000000000..8a7f73f3b067 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -0,0 +1,104 @@ +/* + * 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 "cacheTransceiver.h" +#include "tensorrt_llm/batch_manager/cacheTransceiver.h" +#include "tensorrt_llm/batch_manager/kvCacheManager.h" +#include "tensorrt_llm/executor/executor.h" +#include "tensorrt_llm/nanobind/common/customCasters.h" +#include +#include +#include +#include +#include +#include +#include + +using SizeType32 = tensorrt_llm::runtime::SizeType32; + +namespace tb = tensorrt_llm::batch_manager; +namespace nb = nanobind; + +namespace +{ + +class PyCacheTransceiver : public tb::BaseCacheTransceiver +{ +public: + // using BaseCacheTransceiver::BaseCacheTransceiver; // Inherit constructors + NB_TRAMPOLINE(tb::BaseCacheTransceiver, 6); + + void respondAndSendAsync(tb::LlmRequest* llmRequest) override + { + NB_OVERRIDE_PURE(respondAndSendAsync, llmRequest); + } + + void requestAndReceiveSync(tb::LlmRequest* llmRequest) override + { + NB_OVERRIDE_PURE(requestAndReceiveSync, llmRequest); + } + + void requestAndReceiveAsync(tb::LlmRequest* llmRequest) override + { + NB_OVERRIDE_PURE(requestAndReceiveAsync, llmRequest); + } + + void checkContextTransferStatus(std::optional const& atLeastRequestNum = std::nullopt) override + { + NB_OVERRIDE_PURE(checkContextTransferStatus, atLeastRequestNum); + } + + void checkGenTransferStatus(std::optional const& atLeastRequestNum = std::nullopt) override + { + NB_OVERRIDE_PURE(checkGenTransferStatus, atLeastRequestNum); + } + + bool checkGenTransferComplete() const override + { + NB_OVERRIDE_PURE(checkGenTransferComplete); + } +}; +} // namespace + +void tb::CacheTransceiverBindings::initBindings(nb::module_& m) +{ + nb::class_(m, "BaseCacheTransceiver") + .def("respond_and_send_async", &BaseCacheTransceiver::respondAndSendAsync) + .def("request_and_receive_sync", &BaseCacheTransceiver::requestAndReceiveSync) + .def("request_and_receive_async", &BaseCacheTransceiver::requestAndReceiveAsync) + .def("check_context_transfer_status", &BaseCacheTransceiver::checkContextTransferStatus) + .def("check_gen_transfer_status", &BaseCacheTransceiver::checkGenTransferStatus) + .def("check_gen_transfer_complete", &BaseCacheTransceiver::checkGenTransferComplete); + + nb::enum_(m, "AttentionType") + .value("DEFAULT", executor::kv_cache::CacheState::AttentionType::kDEFAULT) + .value("MLA", executor::kv_cache::CacheState::AttentionType::kMLA); + + nb::class_(m, "CacheTransceiver") + .def(nb::init, SizeType32, SizeType32, + runtime::WorldConfig, nvinfer1::DataType, executor::kv_cache::CacheState::AttentionType, + std::optional>(), + nb::arg("cache_manager"), nb::arg("num_kv_heads_per_layer"), nb::arg("size_per_head"), + nb::arg("tokens_per_block"), nb::arg("world_config"), nb::arg("dtype"), nb::arg("attention_type"), + nb::arg("cache_transceiver_config") = std::nullopt); + + nb::class_(m, "CacheTransBufferManager") + .def(nb::init>(), nb::arg("cache_manager"), + nb::arg("max_num_tokens") = std::nullopt) + .def_static("pre_alloc_buffer_size", &tb::kv_cache_manager::CacheTransBufferManager::preAllocBufferSize, + nb::arg("cache_size_bytes_per_token_per_window"), nb::arg("cache_transceiver_config") = nb::none()); +} diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.h b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.h new file mode 100644 index 000000000000..90fc63d4fdea --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.h @@ -0,0 +1,29 @@ +/* + * 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::batch_manager +{ +class CacheTransceiverBindings +{ +public: + static void initBindings(nb::module_& m); +}; +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp new file mode 100644 index 000000000000..74049eaf96ba --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -0,0 +1,490 @@ +/* + * 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 "kvCacheManager.h" +#include "tensorrt_llm/batch_manager/kvCacheManager.h" +#include "tensorrt_llm/batch_manager/peftCacheManager.h" +#include "tensorrt_llm/nanobind/common/bindTypes.h" +#include "tensorrt_llm/nanobind/common/customCasters.h" +#include "tensorrt_llm/runtime/torch.h" +#include "tensorrt_llm/runtime/torchView.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tb = tensorrt_llm::batch_manager; +namespace tbk = tensorrt_llm::batch_manager::kv_cache_manager; +namespace tr = tensorrt_llm::runtime; +namespace nb = nanobind; +using BlockKey = tbk::BlockKey; +using VecUniqueTokens = tensorrt_llm::runtime::VecUniqueTokens; +using SizeType32 = tensorrt_llm::runtime::SizeType32; +using TokenIdType = tensorrt_llm::runtime::TokenIdType; +using VecTokens = std::vector; +using CudaStreamPtr = std::shared_ptr; +using CacheBlockIds = std::vector>; + +NB_MAKE_OPAQUE(CacheBlockIds); + +namespace +{ +std::optional from_torch(std::optional torchPtr) +{ + if (torchPtr) + { + return tr::TorchView::of(torchPtr.value()); + } + return std::nullopt; +} + +class PyKvCacheManager : public tbk::BaseKVCacheManager +{ +public: + NB_TRAMPOLINE(tbk::BaseKVCacheManager, 28); + + // using BaseKVCacheManager::BaseKVCacheManager; // Inherit constructors + void allocatePools(bool useUvm = false) override + { + NB_OVERRIDE_PURE(allocatePools, useUvm); + } + + void releasePools() override + { + NB_OVERRIDE_PURE(releasePools); + } + + void startScheduling() override + { + NB_OVERRIDE_PURE(startScheduling); + } + + SizeType32 getTokensPerBlock() const override + { + NB_OVERRIDE_PURE(getTokensPerBlock); + } + + SizeType32 getMaxNumBlocks() const override + { + NB_OVERRIDE_PURE(getMaxNumBlocks); + } + + SizeType32 getNumPools() const override + { + NB_OVERRIDE_PURE(getNumPools); + } + + tbk::KvCacheStats getKvCacheStats() const override + { + NB_OVERRIDE_PURE(getKvCacheStats); + } + + void addToken(tb::LlmRequest::RequestIdType requestId) override + { + NB_OVERRIDE_PURE(addToken, requestId); + } + + void addSequence(tb::LlmRequest::RequestIdType requestId, SizeType32 inputLength, SizeType32 beamWidth, + tensorrt_llm::common::OptionalRef llmRequest = std::nullopt) override + { + NB_OVERRIDE_PURE(addSequence, requestId, inputLength, beamWidth, llmRequest); + } + + void removeSequence(tb::LlmRequest::RequestIdType requestId, + tensorrt_llm::common::OptionalRef llmRequest = std::nullopt) override + { + NB_OVERRIDE_PURE(removeSequence, requestId, llmRequest); + } + + tbk::GenerationRequest const& getSequence(tb::LlmRequest::RequestIdType requestId) const override + { + NB_OVERRIDE_PURE(getSequence, requestId); + } + + void schedulingRemoveSequence(tb::LlmRequest::RequestIdType requestId) override + { + NB_OVERRIDE_PURE(schedulingRemoveSequence, requestId); + } + + tensorrt_llm::runtime::ITensor::SharedPtr getBlockPoolPointers() const override + { + NB_OVERRIDE_PURE(getBlockPoolPointers); + } + + tensorrt_llm::runtime::ITensor::SharedPtr getLayerToPoolMapping() const override + { + NB_OVERRIDE_PURE(getLayerToPoolMapping); + } + + void getBlockOffsetsOfBatch(tensorrt_llm::runtime::ITensor& output, SizeType32 firstBatchSlotIdx, + SizeType32 batchSize, SizeType32 beamWidth) const override + { + NB_OVERRIDE_PURE(getBlockOffsetsOfBatch, output, firstBatchSlotIdx, batchSize, beamWidth); + } + + SizeType32 copyBlockOffsets(tensorrt_llm::runtime::ITensor& output, SizeType32 outputSlotOffset, + tb::LlmRequest::RequestIdType requestId) const override + { + NB_OVERRIDE_PURE(copyBlockOffsets, output, outputSlotOffset, requestId); + } + + bool isEnableBlockReuse() const override + { + NB_OVERRIDE_PURE(isEnableBlockReuse); + } + + void rewindKVCache(tb::LlmRequest::RequestIdType requestId, SizeType32 rewindLengths) override + { + NB_OVERRIDE_PURE(rewindKVCache, requestId, rewindLengths); + } + + bool isCrossKv() const override + { + NB_OVERRIDE_PURE(isCrossKv); + } + + std::optional findNewContextBlock( + VecUniqueTokens const& uniqueTokens, tb::LlmRequest const& llmRequest) const override + { + NB_OVERRIDE_PURE(findNewContextBlock, uniqueTokens, llmRequest); + } + + void storeContextBlocks(tb::LlmRequest const& llmRequest) override + { + NB_OVERRIDE_PURE(storeContextBlocks, llmRequest); + } + + std::vector> const& getCacheBlockIds( + tb::LlmRequest::RequestIdType requestId, SizeType32 windowSize) const override + { + NB_OVERRIDE_PURE(getCacheBlockIds, requestId, windowSize); + } + + std::vector>> getBatchCacheBlockIds( + std::vector const& requestIds, SizeType32 windowSize) const override + { + NB_OVERRIDE_PURE(getBatchCacheBlockIds, requestIds, windowSize); + } + + std::vector getNewlyAllocatedBlockIds( + tb::LlmRequest::RequestIdType requestId, SizeType32 windowSize) const override + { + NB_OVERRIDE_PURE(getNewlyAllocatedBlockIds, requestId, windowSize); + } + + SizeType32 getUsedNumBlocks() const override + { + NB_OVERRIDE_PURE(getUsedNumBlocks); + } + + SizeType32 getNumFreeBlocks() const override + { + NB_OVERRIDE_PURE(getNumFreeBlocks); + } + + tbk::BlockManager const& getBlockManager() const override + { + NB_OVERRIDE_PURE(getBlockManager); + } + + std::deque getLatestEvents( + std::optional timeout = std::nullopt) const override + { + NB_OVERRIDE_PURE(getLatestEvents, timeout); + } + + tensorrt_llm::runtime::ITensor::SharedPtr getPrimaryPool(SizeType32 layer_idx) const override + { + NB_OVERRIDE_PURE(getPrimaryPool, layer_idx); + } + + SizeType32 getPoolLayerIdx(SizeType32 layer_idx) const override + { + NB_OVERRIDE_PURE(getPoolLayerIdx, layer_idx); + } + + void refreshBlocks() override + { + NB_OVERRIDE_PURE(refreshBlocks); + } + + void flushIterationEvents() override + { + NB_OVERRIDE_PURE(flushIterationEvents); + } +}; + +// TODO: Deduplicate executor bindings KvCacheStats +class PyBasePeftCacheManager : public tb::BasePeftCacheManager +{ +public: + ~PyBasePeftCacheManager() override = default; + + NB_TRAMPOLINE(tb::BasePeftCacheManager, 8); + + void addRequestPeft(tb::BasePeftCacheManager::LlmRequestPtr llmRequest, bool tryGpuCache = true) override + { + NB_OVERRIDE_PURE(addRequestPeft, llmRequest, tryGpuCache); + } + + tb::BasePeftCacheManager::PeftTable ensureBatch(tb::RequestVector const& contextRequests, + tb::RequestVector const& generationRequests, bool resetGpuCache = false) override + { + NB_OVERRIDE_PURE(ensureBatch, contextRequests, generationRequests, resetGpuCache); + } + + void resetDeviceCache() override + { + NB_OVERRIDE_PURE(resetDeviceCache); + } + + void markRequestDone(tb::LlmRequest const& llmReq, bool pause = false) override + { + NB_OVERRIDE_PURE(markRequestDone, llmReq, pause); + } + + tr::SizeType32 getMaxDevicePages() const override + { + NB_OVERRIDE_PURE(getMaxDevicePages); + } + + tr::SizeType32 getMaxHostPages() const override + { + NB_OVERRIDE_PURE(getMaxHostPages); + } + + tr::SizeType32 determineNumPages(std::shared_ptr llmRequest) const override + { + NB_OVERRIDE_PURE(determineNumPages, llmRequest); + } + + bool enabled() const override + { + NB_OVERRIDE_PURE(enabled); + } +}; +} // namespace + +void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) +{ + nb::class_(m, "KvCacheStats") + .def(nb::init<>()) + .def_rw("max_num_blocks", &tbk::KvCacheStats::maxNumBlocks) + .def_rw("free_num_blocks", &tbk::KvCacheStats::freeNumBlocks) + .def_rw("used_num_blocks", &tbk::KvCacheStats::usedNumBlocks) + .def_rw("tokens_per_block", &tbk::KvCacheStats::toksPerBlock) + .def_rw("alloc_total_blocks", &tbk::KvCacheStats::allocTotalBlocks) + .def_rw("alloc_new_blocks", &tbk::KvCacheStats::allocNewBlocks) + .def_rw("reused_blocks", &tbk::KvCacheStats::reusedBlocks) + .def_rw("missed_blocks", &tbk::KvCacheStats::missedBlocks) + .def_rw("cache_hit_rate", &tbk::KvCacheStats::cacheHitRate) + .def_rw("num_free_blocks_per_window_size", &tbk::KvCacheStats::numFreeBlocksPerWindowSize); + + nb::class_(m, "TempAttentionWindowInputs") + .def(nb::init<>()) + .def_rw("paged_context_fmha", &tbk::TempAttentionWindowInputs::pagedContextFMHA) + .def_rw("max_input_len", &tbk::TempAttentionWindowInputs::maxInputLen) + .def_rw("max_num_tokens", &tbk::TempAttentionWindowInputs::maxNumTokens); + + nb::class_(m, "BlockKey") + .def(nb::init<>()) + .def(nb::init>(), nb::arg("tokens"), + nb::arg("lora_task_id") = std::nullopt) + .def(nb::init, VecUniqueTokens const&>(), nb::arg("uses_extra_ids"), + nb::arg("lora_task_id"), nb::arg("unique_tokens")) + .def_ro("uses_extra_ids", &tbk::BlockKey::usesExtraIds) + .def_ro("lora_task_id", &tbk::BlockKey::loraTaskId) + .def_ro("unique_tokens", &tbk::BlockKey::uniqueTokens); + + nb::class_(m, "BlockKeyHasher") + .def_static("hash", &tbk::BlockKeyHasher::hash, nb::arg("block_key"), nb::arg("parent_hash") = 0); + + nb::class_(m, "KVCacheEventManager") + .def(nb::init(), nb::arg("max_kv_event_entries")); + + nb::class_(m, "BaseKVCacheManager") + .def_static("calculate_max_num_blocks", &tbk::BaseKVCacheManager::calculateMaxNumBlocks, nb::arg("config"), + nb::arg("is_cross_attention"), nb::arg("dtype"), nb::arg("model_config"), nb::arg("world_config"), + nb::arg("window_size_to_layers"), nb::arg("allotted_primary_mem_bytes"), + nb::arg("allotted_secondary_mem_bytes"), nb::arg("extra_cost_memory"), nb::arg("kv_factor")) + .def("allocate_pools", &BaseKVCacheManager::allocatePools) + .def("release_pools", &BaseKVCacheManager::releasePools) + .def("start_scheduling", &BaseKVCacheManager::startScheduling) + .def_prop_ro("tokens_per_block", &BaseKVCacheManager::getTokensPerBlock) + .def_prop_ro("max_num_blocks", &BaseKVCacheManager::getMaxNumBlocks) + .def_prop_ro("num_pools", &BaseKVCacheManager::getNumPools) + .def("get_kv_cache_stats", &BaseKVCacheManager::getKvCacheStats) + .def_prop_ro("max_blocks_per_seq", + [](tbk::BaseKVCacheManager& self) { return self.getOffsetTableDimensions().maxBlocksPerSeq; }) + .def("get_needed_blocks_one_step", &BaseKVCacheManager::getNeededBlocksOneStep) + .def("get_remaining_blocks_to_completion", &BaseKVCacheManager::getRemainingBlocksToCompletion) + .def("add_token", &BaseKVCacheManager::addToken) + .def("add_sequence", &BaseKVCacheManager::addSequence) + .def("remove_sequence", &BaseKVCacheManager::removeSequence) + .def("scheduling_remove_sequence", &BaseKVCacheManager::schedulingRemoveSequence) + .def("get_block_pool_pointers", + [](tbk::BaseKVCacheManager& self) + { + std::optional block_pool_pointers{std::nullopt}; + auto tensor = self.getBlockPoolPointers(); + if (tensor) + { + std::shared_ptr _tensor = std::move(tensor); + block_pool_pointers = tr::Torch::tensor(_tensor); + } + return block_pool_pointers; + }) + .def("get_layer_to_pool_mapping", + [](tbk::BaseKVCacheManager& self) + { + std::optional layer_to_pool_mapping{std::nullopt}; + auto tensor = self.getLayerToPoolMapping(); + if (tensor) + { + std::shared_ptr _tensor = std::move(tensor); + layer_to_pool_mapping = tr::Torch::tensor(_tensor); + } + return layer_to_pool_mapping; + }) + .def("get_primary_pool_data", + [](tbk::BaseKVCacheManager& self, SizeType32 layer_idx) -> at::Tensor + { + auto pool = tr::Torch::tensor(self.getPrimaryPool(layer_idx)); + auto pool_layer_idx = self.getPoolLayerIdx(layer_idx); + return pool.index({torch::indexing::Slice(), pool_layer_idx}); + }) + .def("get_block_offsets_of_batch", + [](tbk::BaseKVCacheManager& self, at::Tensor output, SizeType32 firstBatchSlotIdx, SizeType32 batchSize, + SizeType32 beamWidth) + { + auto _output = from_torch(output); + TLLM_CHECK_WITH_INFO(_output.has_value(), "Invalid output tensor."); + self.getBlockOffsetsOfBatch(*(_output.value()), firstBatchSlotIdx, batchSize, beamWidth); + }) + .def("copy_block_offsets", + [](tbk::BaseKVCacheManager& self, at::Tensor output, SizeType32 outputSlotOffset, + tb::LlmRequest::RequestIdType requestId) + { + auto _output = from_torch(output); + TLLM_CHECK_WITH_INFO(_output.has_value(), "Invalid output tensor."); + auto maxBlockCount = self.copyBlockOffsets(*(_output.value()), outputSlotOffset, requestId); + return maxBlockCount; + }) + .def("copy_batch_block_offsets", + [](tbk::BaseKVCacheManager& self, at::Tensor output, + std::vector const& requestIds, SizeType32 const beamWidth, + SizeType32 const offset) + { + auto _output = from_torch(output); + TLLM_CHECK_WITH_INFO(_output.has_value(), "Invalid output tensor."); + for (size_t i = 0; i < requestIds.size(); ++i) + { + self.copyBlockOffsets(*(_output.value()), i * beamWidth + offset, requestIds[i]); + } + }) + .def( + "get_latest_events", + [](tbk::BaseKVCacheManager& self, std::optional timeout_ms = std::nullopt) + { + if (timeout_ms) + { + return self.getLatestEvents(std::chrono::milliseconds(static_cast(*timeout_ms))); + } + return self.getLatestEvents(std::nullopt); + }, + nb::arg("timeout_ms") = std::nullopt) + .def_prop_ro("enable_block_reuse", &BaseKVCacheManager::isEnableBlockReuse) + .def("rewind_kv_cache", &BaseKVCacheManager::rewindKVCache) + .def_prop_ro("cross_kv", &BaseKVCacheManager::isCrossKv) + .def("store_context_blocks", &BaseKVCacheManager::storeContextBlocks) + .def("get_cache_block_ids", &BaseKVCacheManager::getCacheBlockIds) + .def("get_batch_cache_block_ids", &BaseKVCacheManager::getBatchCacheBlockIds) + .def("get_newly_allocated_block_ids", &BaseKVCacheManager::getNewlyAllocatedBlockIds) + .def("flush_iteration_events", &BaseKVCacheManager::flushIterationEvents); + + nb::bind_vector(m, "CacheBlockIds") + .def("__getstate__", [](CacheBlockIds const& v) { return nb::make_tuple(v); }) + .def("__setstate__", + [](CacheBlockIds& self, nb::tuple const& t) + { + if (t.size() != 1) + throw std::runtime_error("Invalid state!"); + new (&self) CacheBlockIds(nb::cast>>(t[0])); + }); + + nb::enum_(m, "CacheType") + .value("SELF", tbk::CacheType::kSELF) + .value("CROSS", tbk::CacheType::kCROSS) + .value("SELFKONLY", tbk::CacheType::kSELFKONLY); + + nb::class_(m, "KVCacheManager") + .def(nb::init const&, SizeType32, SizeType32, + std::map> const&, SizeType32, SizeType32, + std::vector const&, std::optional const&, + nvinfer1::DataType, SizeType32, int64_t, std::optional, bool, bool, + tbk::CacheType, std::optional, + std::shared_ptr, bool, bool>(), + nb::arg("num_kv_heads_per_layer"), nb::arg("size_per_head"), nb::arg("tokens_per_block"), + nb::arg("blocks_per_window"), nb::arg("max_num_sequences"), nb::arg("max_beam_width"), + nb::arg("max_attention_window_vec"), nb::arg("temp_attention_window_inputs").none(), nb::arg("dtype"), + nb::arg("sink_token_length"), nb::arg("stream"), nb::arg("max_sequence_length").none(), + nb::arg("enable_block_reuse") = false, nb::arg("onboard_blocks") = true, + nb::arg("cache_type") = tbk::CacheType::kSELF, nb::arg("secondary_offload_min_priority") = std::nullopt, + nb::arg("event_manager") = nullptr, nb::arg("enable_partial_reuse") = true, + nb::arg("copy_on_partial_reuse") = true); +} + +void tb::BasePeftCacheManagerBindings::initBindings(nb::module_& m) +{ + nb::class_(m, "BasePeftCacheManager") + .def("add_request_peft", &tb::BasePeftCacheManager::addRequestPeft, nb::arg("request"), + nb::arg("try_gpu_cache") = true) + .def( + "ensure_batch", + [](tb::BasePeftCacheManager& self, tb::RequestVector const& contextRequests, + tb::RequestVector const& generationRequests, bool resetGpuCache) + { + nb::gil_scoped_release release; + return self.ensureBatch(contextRequests, generationRequests, resetGpuCache); + }, + nb::arg("context_requests"), nb::arg("generation_requests"), nb::arg("reset_gpu_cache") = false) + .def("reset_device_cache", &tb::BasePeftCacheManager::resetDeviceCache) + .def("mark_request_done", &tb::BasePeftCacheManager::markRequestDone, nb::arg("request"), + nb::arg("pause") = false) + .def_prop_ro("max_device_pages", &tb::BasePeftCacheManager::getMaxDevicePages) + .def_prop_ro("max_host_pages", &tb::BasePeftCacheManager::getMaxHostPages) + .def("determine_num_pages", &tb::BasePeftCacheManager::determineNumPages, nb::arg("request")) + .def_prop_ro("enabled", &tb::BasePeftCacheManager::enabled); + + nb::class_(m, "PeftCacheManager") + .def(nb::init(), + nb::arg("config"), nb::arg("model_config"), nb::arg("world_config"), nb::arg("buffer_manager")) + .def("is_task_cached", &tb::PeftCacheManager::isTaskCached, nb::arg("taskId")); + + nb::class_(m, "NoOpPeftCacheManager").def(nb::init<>()); +} diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.h b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.h new file mode 100644 index 000000000000..786c0d391df5 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.h @@ -0,0 +1,39 @@ +/* + * 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::batch_manager::kv_cache_manager +{ +class KVCacheManagerBindings +{ +public: + static void initBindings(nb::module_& m); +}; +} // namespace tensorrt_llm::batch_manager::kv_cache_manager + +namespace tensorrt_llm::batch_manager +{ +class BasePeftCacheManagerBindings +{ +public: + static void initBindings(nb::module_& m); +}; +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.cpp new file mode 100644 index 000000000000..d8f45cb865f3 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.cpp @@ -0,0 +1,131 @@ +/* + * 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 "llmRequest.h" +#include "tensorrt_llm/nanobind/common/customCasters.h" + +#include "tensorrt_llm/batch_manager/llmRequest.h" +#include "tensorrt_llm/nanobind/common/bindTypes.h" +#include "tensorrt_llm/runtime/torch.h" +#include "tensorrt_llm/runtime/torchUtils.h" +#include "tensorrt_llm/runtime/torchView.h" + +#include +#include + +#include + +namespace tb = tensorrt_llm::batch_manager; +namespace tr = tensorrt_llm::runtime; +namespace tle = tensorrt_llm::executor; + +using namespace tensorrt_llm::nanobind::batch_manager; + +using LlmRequestPtr = std::shared_ptr; +using RequestList = std::list; + +namespace +{ + +std::optional from_torch(std::optional torchPtr) +{ + if (torchPtr) + { + return tr::TorchView::of(torchPtr.value()); + } + return std::nullopt; +} + +} // namespace + +std::optional LlmRequest::callbackAdapter( + std::optional callback) +{ + if (!callback) + { + return std::nullopt; + } + + return [callback](RequestIdType reqId, tr::ITensor::SharedPtr& tensor, tb::LlmRequest::BeamTokens const& tokens, + tr::BufferManager::CudaStreamPtr stream, std::optional clientId) + { + at::Tensor atTensor = tr::Torch::tensor(tensor); + callback.value()(reqId, atTensor, tokens, runtime::TorchUtils::stream(*stream).unwrap(), clientId); + }; +} + +std::shared_ptr LlmRequest::toTrtLlm() const +{ + + auto const draftTokens = std::make_shared>(*mDraftTokens.get()); + auto const optDraftTokens = std::optional>>(draftTokens); + auto const encoderInputTokens = mEncoderTokens.has_value() + ? std::make_shared>(*mEncoderTokens.value().get()) + : nullptr; + auto const optEncoderInputTokens = std::optional>>(encoderInputTokens); + // 49 parameters + return std::make_shared( // + mRequestId, // + mMaxNewTokens, // + std::make_shared>(mTokens.at(0)), // + mSamplingConfig, // + mIsStreaming, // + mEndId, // + mPadId, // + from_torch(mEmbeddingBias), // + from_torch(mBadWordsList), // + from_torch(mStopWordsList), // + mPositionIds, // + from_torch(mPromptEmbeddingTable), // + mPromptVocabSize, // + mMultimodalHashes, // + mMultimodalPositions, // + mMultimodalLengths, // + from_torch(mMultimodalEmbedding), // + from_torch(mMropeRotaryCosSin), // + mMropePositionDeltas, // + mLoraTaskId, // + from_torch(mLoraWeights), // + from_torch(mLoraConfig), // + mLookaheadConfig, // + mKvCacheRetentionConfig, // + mReturnLogProbs, // + mReturnContextLogits, // + mReturnGenerationLogits, // + optDraftTokens, // + from_torch(mDraftLogits), // + mExcludeInputFromOutput, // + callbackAdapter(mLogitsPostProcessor), // + mApplyLogitsPostProcessorBatched, // + optEncoderInputTokens, // + mReturnEncoderOutput, // + mClientId, // + mPriority, // + from_torch(mEncoderInputFeatures), // + mEncoderOutputLength, // + from_torch(mCrossAttentionMask), // + getLlmRequestType(), // + std::nullopt, // inputTokenExtraIds + mNumReturnSequences, // + mEagleConfig, // + from_torch(mSkipCrossAttnBlocks), // + false, // returnPerfMetrics + mGuidedDecodingParams, // + mLanguageAdapterUid, // + mAllottedTimeMs, // + mContextPhaseParams // + ); +} diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.h b/cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.h new file mode 100644 index 000000000000..624dc55112d7 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.h @@ -0,0 +1,160 @@ +/* + * 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/batch_manager/llmRequest.h" + +#include +#include +#include +#include +#include + +namespace nb = nanobind; + +namespace tensorrt_llm::nanobind::batch_manager +{ + +namespace tb = tensorrt_llm::batch_manager; + +/* Unfortunately, torch's default nanobind bindings don't know about c10::cuda::CUDAStream, + * so we have to pass the more generic c10::Stream, and convert it back to a full-fledged + * torch.cuda.Stream in python. See example in test/bindings/test_gpt_manager.py + */ +class LlmRequest : public tb::GenericLlmRequest +{ +public: + using Base = GenericLlmRequest; + using TensorPtr = Base::TensorPtr; + using SizeType32 = Base::SizeType32; + using TokenIdType = Base::TokenIdType; + using RequestIdType = Base::RequestIdType; + using LoraTaskIdType = Base::LoraTaskIdType; + using VecLogProbs = Base::VecLogProbs; + using BeamTokens = Base::BeamTokens; + using VecTokens = Base::VecTokens; + using VecTokenExtraIds = Base::VecTokenExtraIds; + using LogitsPostProcessor = Base::LogitsPostProcessor; + + // 49 parameters + LlmRequest(RequestIdType requestId, SizeType32 maxNewTokens, std::vector inputTokens, + runtime::SamplingConfig samplingConfig, bool isStreaming, std::optional endId = std::nullopt, + std::optional padId = std::nullopt, std::optional embeddingBias = std::nullopt, + std::optional badWordsList = std::nullopt, std::optional stopWordsList = std::nullopt, + std::optional> positionIds = std::nullopt, + std::optional promptEmbeddingTable = std::nullopt, + std::optional promptVocabSize = std::nullopt, + std::optional>> multimodalHashes = std::nullopt, + std::optional> multimodalPositions = std::nullopt, + std::optional> multimodalLengths = std::nullopt, + std::optional multimodalEmbedding = std::nullopt, + std::optional mropeRotaryCosSin = std::nullopt, + std::optional mropePositionDeltas = std::nullopt, + std::optional loraTaskId = std::nullopt, std::optional loraWeights = std::nullopt, + std::optional loraConfig = std::nullopt, + std::optional lookaheadConfig = std::nullopt, + std::optional kvCacheRetentionConfig = std::nullopt, + bool returnLogProbs = false, bool returnContextLogits = false, bool returnGenerationLogits = false, + std::optional draftTokens = std::nullopt, std::optional draftLogits = std::nullopt, + bool excludeInputFromOutput = false, std::optional logitsPostProcessor = std::nullopt, + bool applyLogitsPostProcessorBatched = false, std::optional encoderInputTokens = std::nullopt, + bool returnEncoderOutput = false, std::optional clientId = std::nullopt, + executor::PriorityType priority = executor::Request::kDefaultPriority, + std::optional encoderInputFeatures = std::nullopt, + std::optional encoderOutputLength = std::nullopt, + std::optional crossAttentionMask = std::nullopt, + tb::LlmRequestType llmRequestType = tb::LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, + std::optional inputTokenExtraIds = std::nullopt, SizeType32 numReturnSequences = 1, + std::optional eagleConfig = std::nullopt, + std::optional skipCrossAttnBlocks = std::nullopt, bool returnPerfMetrics = false, + std::optional guidedDecodingParams = std::nullopt, + std::optional languageAdapterUid = std::nullopt, + std::optional allottedTimeMs = std::nullopt, + std::optional const& contextPhaseParams = std::nullopt) + : Base(requestId, // + maxNewTokens, // + std::make_shared>(std::move(inputTokens)), // + samplingConfig, // + isStreaming, // + endId, // + padId, // + embeddingBias, // + badWordsList, // + stopWordsList, // + positionIds.has_value() ? std::make_shared>(std::move(positionIds.value())) // + : std::optional>>(std::nullopt), // + promptEmbeddingTable, // + promptVocabSize, // + multimodalHashes.has_value() + ? std::make_optional( + std::make_shared>>(std::move(multimodalHashes.value()))) // + : std::optional>>>(std::nullopt), // + multimodalPositions.has_value() + ? std::make_shared>(std::move(multimodalPositions.value())) // + : std::optional>>(std::nullopt), // + multimodalLengths.has_value() + ? std::make_shared>(std::move(multimodalLengths.value())) // + : std::optional>>(std::nullopt), // + multimodalEmbedding, // + mropeRotaryCosSin, // + mropePositionDeltas, // + loraTaskId, // + loraWeights, // + loraConfig, // + lookaheadConfig, // + kvCacheRetentionConfig, // + returnLogProbs, // + returnContextLogits, // + returnGenerationLogits, // + draftTokens.has_value() ? std::make_shared(std::move(draftTokens.value())) // + : std::make_shared(), // + draftLogits, // + excludeInputFromOutput, // + logitsPostProcessor, // + applyLogitsPostProcessorBatched, // + encoderInputTokens ? std::make_optional(std::make_shared(std::move(*encoderInputTokens))) // + : std::optional>(std::nullopt), // + returnEncoderOutput, // + clientId, // + priority, // + encoderInputFeatures, // + encoderOutputLength, // + crossAttentionMask, // + llmRequestType, // + inputTokenExtraIds // + ? std::make_optional(std::make_shared(std::move(*inputTokenExtraIds))) // + : std::optional>(std::nullopt), // + numReturnSequences, // + eagleConfig, // + skipCrossAttnBlocks, // + returnPerfMetrics, // + guidedDecodingParams, // + languageAdapterUid, // + allottedTimeMs, // + contextPhaseParams // + ) + { + } + + static std::optional callbackAdapter( + std::optional callback); + + [[nodiscard]] std::shared_ptr toTrtLlm() const; +}; + +} // namespace tensorrt_llm::nanobind::batch_manager diff --git a/cpp/tensorrt_llm/nanobind/bindings.cpp b/cpp/tensorrt_llm/nanobind/bindings.cpp index adc82587433d..43a985658ddf 100644 --- a/cpp/tensorrt_llm/nanobind/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/bindings.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * 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"); @@ -15,14 +15,484 @@ * limitations under the License. */ +#include "tensorrt_llm/nanobind/common/customCasters.h" #include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "tensorrt_llm/batch_manager/peftCacheManagerConfig.h" +#include "tensorrt_llm/common/quantization.h" +#include "tensorrt_llm/nanobind/batch_manager/algorithms.h" +#include "tensorrt_llm/nanobind/batch_manager/bindings.h" +#include "tensorrt_llm/nanobind/batch_manager/cacheTransceiver.h" +#include "tensorrt_llm/nanobind/batch_manager/kvCacheManager.h" +#include "tensorrt_llm/nanobind/batch_manager/llmRequest.h" +#include "tensorrt_llm/nanobind/executor/bindings.h" +#include "tensorrt_llm/nanobind/runtime/bindings.h" +#include "tensorrt_llm/nanobind/testing/modelSpecBinding.h" +#include "tensorrt_llm/nanobind/userbuffers/bindings.h" +#include "tensorrt_llm/runtime/common.h" +#include "tensorrt_llm/runtime/cudaStream.h" +#include "tensorrt_llm/runtime/gptJsonConfig.h" +#include "tensorrt_llm/runtime/ipcNvlsMemory.h" +#include "tensorrt_llm/runtime/memoryCounters.h" +#include "tensorrt_llm/runtime/samplingConfig.h" +#include "tensorrt_llm/runtime/utils/mpiUtils.h" + +namespace nb = nanobind; +namespace tb = tensorrt_llm::batch_manager; +namespace tbk = tensorrt_llm::batch_manager::kv_cache_manager; +namespace tpb = tensorrt_llm::nanobind::batch_manager; +namespace tc = tensorrt_llm::common; +namespace tr = tensorrt_llm::runtime; +namespace tle = tensorrt_llm::executor; +using SizeType32 = tr::SizeType32; +using TokenIdType = tr::TokenIdType; +template +using OptVec = std::optional>; #if not defined(TRTLLM_NB_MODULE) #error "TRTLLM_NB_MODULE must be defined" #endif +namespace +{ +tr::SamplingConfig makeSamplingConfig(std::vector const& configs) +{ + return tr::SamplingConfig(configs); +} +} // namespace + NB_MODULE(TRTLLM_NB_MODULE, m) { m.doc() = "TensorRT-LLM Python bindings for C++ runtime"; m.attr("binding_type") = "nanobind"; + nb::set_leak_warnings(false); + + // Create MpiComm binding first since it's used in the executor bindings + nb::class_(m, "MpiComm") + .def_static("rank", + []() + { + auto& session = tensorrt_llm::mpi::MpiComm::session(); + return session.tensorrt_llm::mpi::MpiComm::getRank(); + }) + .def_static("size", + []() + { + auto& session = tensorrt_llm::mpi::MpiComm::session(); + return session.tensorrt_llm::mpi::MpiComm::getSize(); + }) + .def_static("local_size", + []() + { + auto& session = tensorrt_llm::mpi::MpiComm::localSession(); + return session.tensorrt_llm::mpi::MpiComm::getSize(); + }) + .def_static("local_init", []() { tensorrt_llm::mpi::MpiComm::localSession(); }) + .def_static("set_raw_mpi_session_by_fortran_handle", + [](int64_t fortran_handle) { tensorrt_llm::mpi::MpiComm::setRawSessionByFortran(fortran_handle); }) + .def_static("split", + [](size_t color, size_t rank) + { + auto& world = tensorrt_llm::mpi::MpiComm::world(); + tensorrt_llm::mpi::MpiComm::setSession(world.split(color, rank)); + }); + + nb::class_(m, "CudaStream") + .def( + "__init__", + [](tr::CudaStream* self, nb::object py_stream) + { + cudaStream_t stream = reinterpret_cast(nb::cast(py_stream)); + new (self) tr::CudaStream{stream}; + }, + nb::arg("stream_ptr")) + .def("get_device", &tr::CudaStream::getDevice); + + // Create submodule for executor bindings. + auto mExecutor = m.def_submodule("executor", "Executor bindings"); + auto mInternal = m.def_submodule("internal", "Internal submodule of TRTLLM runtime"); + auto mInternalRuntime = mInternal.def_submodule("runtime", "Runtime internal bindings"); + auto mInternalTesting = mInternal.def_submodule("testing", "Testing internal bindings"); + auto mInternalBatchManager = mInternal.def_submodule("batch_manager", "Batch manager internal bindings"); + + tensorrt_llm::nanobind::executor::initBindings(mExecutor); + tensorrt_llm::nanobind::runtime::initBindingsEarly(mInternalRuntime); + + auto buildInfo = m.def_submodule("BuildInfo"); + buildInfo.attr("ENABLE_MULTI_DEVICE") = nb::int_(ENABLE_MULTI_DEVICE); + + nb::class_(m, "PeftCacheManagerConfig") + .def(nb::init, std::optional, std::optional>(), + nb::arg("num_host_module_layer") = 0, nb::arg("num_device_module_layer") = 0, + nb::arg("optimal_adapter_size") = 8, nb::arg("max_adapter_size") = 64, nb::arg("num_put_workers") = 1, + nb::arg("num_ensure_workers") = 1, nb::arg("num_copy_streams") = 1, + nb::arg("max_pages_per_block_host") = 24, nb::arg("max_pages_per_block_device") = 8, + nb::arg("device_cache_percent") = std::nullopt, nb::arg("host_cache_size") = std::nullopt, + nb::arg("lora_prefetch_dir") = std::nullopt) + .def_rw("num_host_module_layer", &tb::PeftCacheManagerConfig::numHostModuleLayer) + .def_rw("num_device_module_layer", &tb::PeftCacheManagerConfig::numDeviceModuleLayer) + .def_rw("optimal_adapter_size", &tb::PeftCacheManagerConfig::optimalAdapterSize) + .def_rw("max_adapter_size", &tb::PeftCacheManagerConfig::maxAdapterSize) + .def_rw("num_put_workers", &tb::PeftCacheManagerConfig::numPutWorkers) + .def_rw("num_ensure_workers", &tb::PeftCacheManagerConfig::numEnsureWorkers) + .def_rw("num_copy_streams", &tb::PeftCacheManagerConfig::numCopyStreams) + .def_rw("max_pages_per_block_host", &tb::PeftCacheManagerConfig::maxPagesPerBlockHost) + .def_rw("max_pages_per_block_device", &tb::PeftCacheManagerConfig::maxPagesPerBlockDevice) + .def_rw("device_cache_percent", &tb::PeftCacheManagerConfig::deviceCachePercent) + .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) + .export_values(); + + nb::enum_(m, "GptModelVariant") + .value("GPT", tr::ModelConfig::ModelVariant::kGpt) + .value("GLM", tr::ModelConfig::ModelVariant::kGlm) + .value("CHATGLM", tr::ModelConfig::ModelVariant::kChatGlm) + .value("MAMBA", tr::ModelConfig::ModelVariant::kMamba) + .value("RECURRENTGEMMA", tr::ModelConfig::ModelVariant::kRecurrentGemma); + + nb::enum_(m, "KVCacheType") + .value("CONTINUOUS", tr::ModelConfig::KVCacheType::kCONTINUOUS) + .value("PAGED", tr::ModelConfig::KVCacheType::kPAGED) + .value("DISABLED", tr::ModelConfig::KVCacheType::kDISABLED) + .def("from_string", tr::ModelConfig::KVCacheTypeFromString); + + nb::enum_(m, "LayerType") + .value("ATTENTION", tr::ModelConfig::LayerType::kATTENTION) + .value("RECURRENT", tr::ModelConfig::LayerType::kRECURRENT); + + nb::enum_(m, "LoraModuleType") + .value("INVALID", tr::LoraModule::ModuleType::kINVALID) + .value("ATTN_QKV", tr::LoraModule::ModuleType::kATTN_QKV) + .value("ATTN_Q", tr::LoraModule::ModuleType::kATTN_Q) + .value("ATTN_K", tr::LoraModule::ModuleType::kATTN_K) + .value("ATTN_V", tr::LoraModule::ModuleType::kATTN_V) + .value("ATTN_DENSE", tr::LoraModule::ModuleType::kATTN_DENSE) + .value("MLP_H_TO_4H", tr::LoraModule::ModuleType::kMLP_H_TO_4H) + .value("MLP_4H_TO_H", tr::LoraModule::ModuleType::kMLP_4H_TO_H) + .value("MLP_GATE", tr::LoraModule::ModuleType::kMLP_GATE) + .value("CROSS_ATTN_QKV", tr::LoraModule::ModuleType::kCROSS_ATTN_QKV) + .value("CROSS_ATTN_Q", tr::LoraModule::ModuleType::kCROSS_ATTN_Q) + .value("CROSS_ATTN_K", tr::LoraModule::ModuleType::kCROSS_ATTN_K) + .value("CROSS_ATTN_V", tr::LoraModule::ModuleType::kCROSS_ATTN_V) + .value("CROSS_ATTN_DENSE", tr::LoraModule::ModuleType::kCROSS_ATTN_DENSE) + .value("MOE_H_TO_4H", tr::LoraModule::ModuleType::kMOE_H_TO_4H) + .value("MOE_4H_TO_H", tr::LoraModule::ModuleType::kMOE_4H_TO_H) + .value("MOE_GATE", tr::LoraModule::ModuleType::kMOE_GATE) + .value("MOE_ROUTER", tr::LoraModule::ModuleType::kMOE_ROUTER) + .value("MLP_ROUTER", tr::LoraModule::ModuleType::kMLP_ROUTER) + .value("MLP_GATE_UP", tr::LoraModule::ModuleType::kMLP_GATE_UP); + + nb::class_(m, "LoraModule") + .def(nb::init(), + nb::arg("module_type"), nb::arg("in_dim"), nb::arg("out_dim"), nb::arg("in_dim_first"), + nb::arg("out_dim_first"), nb::arg("in_tp_split_dim"), nb::arg("out_tp_split_dim")) + .def_prop_ro("module_type", &tr::LoraModule::name) + .def_prop_ro("in_dim", &tr::LoraModule::inDim) + .def_prop_ro("out_dim", &tr::LoraModule::outDim) + .def_prop_ro("in_dim_first", &tr::LoraModule::inDimFirst) + .def_prop_ro("out_dim_first", &tr::LoraModule::outDimFirst) + .def_prop_ro("in_tp_split_dim", &tr::LoraModule::inTpSplitDim) + .def_prop_ro("out_tp_split_dim", &tr::LoraModule::outTpSplitDim) + .def_static("create_lora_modules", &tr::LoraModule::createLoraModules, nb::arg("lora_module_names"), + nb::arg("hidden_size"), nb::arg("mlp_hidden_size"), nb::arg("num_attention_heads"), + nb::arg("num_kv_attention_heads"), nb::arg("attention_head_size"), nb::arg("tp_size") = 1, + nb::arg("num_experts") = 0); + + nb::class_(m, "QuantMode") + .def_static("none", &tc::QuantMode::none) + .def_static("int4_weights", &tc::QuantMode::int4Weights) + .def_static("int8_weights", &tc::QuantMode::int8Weights) + .def_static("activations", &tc::QuantMode::activations) + .def_static("per_channel_scaling", &tc::QuantMode::perChannelScaling) + .def_static("per_token_scaling", &tc::QuantMode::perTokenScaling) + .def_static("per_group_scaling", &tc::QuantMode::perGroupScaling) + .def_static("int8_kv_cache", &tc::QuantMode::int8KvCache) + .def_static("fp8_kv_cache", &tc::QuantMode::fp8KvCache) + .def_static("fp8_qdq", &tc::QuantMode::fp8Qdq) + .def_prop_ro("value", &tc::QuantMode::value) + .def("is_set", &tc::QuantMode::isSet, nb::arg("mode")) + .def_prop_ro("has_int4_weights", &tc::QuantMode::hasInt4Weights) + .def_prop_ro("has_int8_weights", &tc::QuantMode::hasInt8Weights) + .def_prop_ro("has_activations", &tc::QuantMode::hasActivations) + .def_prop_ro("has_per_channel_scaling", &tc::QuantMode::hasPerChannelScaling) + .def_prop_ro("has_per_token_scaling", &tc::QuantMode::hasPerTokenScaling) + .def_prop_ro("has_per_group_scaling", &tc::QuantMode::hasPerGroupScaling) + .def_prop_ro("has_static_activation_scaling", &tc::QuantMode::hasStaticActivationScaling) + .def_prop_ro("has_int8_kv_cache", &tc::QuantMode::hasInt8KvCache) + .def_prop_ro("has_fp8_kv_cache", &tc::QuantMode::hasFp8KvCache) + .def_prop_ro("has_fp8_qdq", &tc::QuantMode::hasFp8Qdq) + .def_prop_ro("has_nvfp4", &tc::QuantMode::hasNvfp4) + .def_prop_ro("has_w4a8_mxfp4_fp8", &tc::QuantMode::hasW4a8Mxfp4Fp8) + .def_prop_ro("has_kv_cache_quant", &tc::QuantMode::hasKvCacheQuant) + .def_static("from_description", &tc::QuantMode::fromDescription, nb::arg("quantize_weights"), + nb::arg("quantize_activations"), nb::arg("per_token"), nb::arg("per_channel"), nb::arg("per_group"), + nb::arg("use_int4_weights"), nb::arg("use_int8_kv_cache"), nb::arg("use_fp8_kv_kache"), + nb::arg("use_fp8_qdq"), nb::arg("use_fp8_rowwise"), nb::arg("use_w4a8_qserve"), nb::arg("use_nvfp4"), + nb::arg("use_fp8_block_scales"), nb::arg("use_w4a8_mxfp4_fp8")) + .def_static("use_smooth_quant", &tc::QuantMode::useSmoothQuant, nb::arg("per_token") = false, + nb::arg("per_channel") = false) + .def_static("use_weight_only", &tc::QuantMode::useWeightOnly, nb::arg("use_int4_weights") = false, + nb::arg("per_group") = false) + .def_static("from_quant_algo", &tc::QuantMode::fromQuantAlgo, nb::arg("quant_algo") = nb::none(), + nb::arg("kv_cache_quant_algo") = nb::none()) + .def(nb::self + nb::self) + .def(nb::self += nb::self) + .def(nb::self - nb::self) + .def(nb::self -= nb::self) + .def(nb::self == nb::self) + .def(nb::self != nb::self); + + nb::class_(m, "ModelConfig") + .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) + .def("vocab_size_padded", &tr::ModelConfig::getVocabSizePadded, nb::arg("world_size")) + .def("num_layers", &tr::ModelConfig::getNbLayers, nb::arg("pipeline_parallelism") = 1, + nb::arg("pipeline_parallelism_rank") = 0) + .def("num_attention_layers", &tr::ModelConfig::getNbAttentionLayers, nb::arg("pipeline_parallelism") = 1, + nb::arg("pipeline_parallelism_rank") = 0) + .def("num_rnn_layers", &tr::ModelConfig::getNbRnnLayers, nb::arg("pipeline_parallelism") = 1, + nb::arg("pipeline_parallelism_rank") = 0) + .def("num_kv_heads", &tr::ModelConfig::getNbKvHeads, nb::arg("layer_idx")) + .def("set_num_kv_heads", &tr::ModelConfig::setNbKvHeads, nb::arg("num_kv_heads")) + .def_prop_ro("num_heads", &tr::ModelConfig::getNbHeads) + .def_prop_ro("hidden_size", &tr::ModelConfig::getHiddenSize) + .def_prop_ro("size_per_head", &tr::ModelConfig::getSizePerHead) + .def_prop_ro("data_type", &tr::ModelConfig::getDataType) + .def_prop_ro("speculative_decoding_mode", &tr::ModelConfig::getSpeculativeDecodingMode) + .def_prop_rw("head_size", &tr::ModelConfig::getSizePerHead, &tr::ModelConfig::setSizePerHead) + .def_prop_rw( + "num_kv_heads_per_layer", &tr::ModelConfig::getNumKvHeadsPerLayer, &tr::ModelConfig::setNumKvHeadsPerLayer) + .def_prop_rw("use_gpt_attention_plugin", + nb::overload_cast<>(&tr::ModelConfig::useGptAttentionPlugin, nb::const_), + nb::overload_cast(&tr::ModelConfig::useGptAttentionPlugin)) + .def_prop_rw("use_packed_input", nb::overload_cast<>(&tr::ModelConfig::usePackedInput, nb::const_), + nb::overload_cast(&tr::ModelConfig::usePackedInput)) + .def_prop_rw("kv_cache_type", nb::overload_cast<>(&tr::ModelConfig::getKVCacheType, nb::const_), + nb::overload_cast(&tr::ModelConfig::setKVCacheType)) + .def_prop_rw("tokens_per_block", &tr::ModelConfig::getTokensPerBlock, &tr::ModelConfig::setTokensPerBlock) + .def_prop_rw("quant_mode", &tr::ModelConfig::getQuantMode, &tr::ModelConfig::setQuantMode) + .def_prop_ro("supports_inflight_batching", &tr::ModelConfig::supportsInflightBatching) + .def_prop_rw("max_batch_size", &tr::ModelConfig::getMaxBatchSize, &tr::ModelConfig::setMaxBatchSize) + .def_prop_rw("max_beam_width", &tr::ModelConfig::getMaxBeamWidth, &tr::ModelConfig::setMaxBeamWidth) + .def_prop_rw("max_input_len", &tr::ModelConfig::getMaxInputLen, &tr::ModelConfig::setMaxInputLen) + .def_prop_rw("max_seq_len", &tr::ModelConfig::getMaxSequenceLen, &tr::ModelConfig::setMaxSequenceLen) + .def_prop_rw("max_num_tokens", &tr::ModelConfig::getMaxNumTokens, &tr::ModelConfig::setMaxNumTokens) + .def_prop_rw("max_prompt_embedding_table_size", &tr::ModelConfig::getMaxPromptEmbeddingTableSize, + &tr::ModelConfig::setMaxPromptEmbeddingTableSize) + .def_prop_ro("use_prompt_tuning", &tr::ModelConfig::usePromptTuning) + .def_prop_ro("use_mrope", &tr::ModelConfig::useMrope) + .def_prop_rw("use_lora_plugin", nb::overload_cast<>(&tr::ModelConfig::useLoraPlugin, nb::const_), + nb::overload_cast(&tr::ModelConfig::useLoraPlugin)) + .def_prop_rw("layer_types", &tr::ModelConfig::getLayerTypes, &tr::ModelConfig::setLayerTypes) + .def_prop_rw("compute_context_logits", nb::overload_cast<>(&tr::ModelConfig::computeContextLogits, nb::const_), + nb::overload_cast(&tr::ModelConfig::computeContextLogits)) + .def_prop_rw("compute_generation_logits", + nb::overload_cast<>(&tr::ModelConfig::computeGenerationLogits, nb::const_), + nb::overload_cast(&tr::ModelConfig::computeGenerationLogits)) + .def_prop_rw("model_variant", &tr::ModelConfig::getModelVariant, &tr::ModelConfig::setModelVariant) + .def_prop_rw("use_cross_attention", &tr::ModelConfig::useCrossAttention, &tr::ModelConfig::setUseCrossAttention) + .def_prop_rw("lora_modules", &tr::ModelConfig::getLoraModules, &tr::ModelConfig::setLoraModules) + .def_prop_rw("max_lora_rank", &tr::ModelConfig::getMaxLoraRank, &tr::ModelConfig::setMaxLoraRank) + .def_prop_rw("mlp_hidden_size", &tr::ModelConfig::getMlpHiddenSize, &tr::ModelConfig::setMlpHiddenSize) + .def_prop_rw("size_per_head", &tr::ModelConfig::getSizePerHead, &tr::ModelConfig::setSizePerHead); + + nb::class_(m, "WorldConfig") + .def(nb::init> const&, bool>(), + nb::arg("tensor_parallelism") = 1, nb::arg("pipeline_parallelism") = 1, nb::arg("context_parallelism") = 1, + nb::arg("rank") = 0, nb::arg("gpus_per_node") = tr::WorldConfig::kDefaultGpusPerNode, + nb::arg("device_ids") = nb::none(), nb::arg("enable_attention_dp") = false) + .def_prop_ro("size", &tr::WorldConfig::getSize) + .def_prop_ro("tensor_parallelism", &tr::WorldConfig::getTensorParallelism) + .def_prop_ro("pipeline_parallelism", &tr::WorldConfig::getPipelineParallelism) + .def_prop_ro("context_parallelism", &tr::WorldConfig::getContextParallelism) + .def_prop_ro("is_tensor_parallel", &tr::WorldConfig::isTensorParallel) + .def_prop_ro("is_pipeline_parallel", &tr::WorldConfig::isPipelineParallel) + .def_prop_ro("is_context_parallel", &tr::WorldConfig::isContextParallel) + .def_prop_ro("rank", &tr::WorldConfig::getRank) + .def_prop_ro("local_rank", &tr::WorldConfig::getLocalRank) + .def_prop_ro("node_rank", &tr::WorldConfig::getNodeRank) + .def_prop_ro("gpus_per_node", &tr::WorldConfig::getGpusPerNode) + .def_prop_ro("gpus_per_group", &tr::WorldConfig::getGpusPerGroup) + .def_prop_ro("device", &tr::WorldConfig::getDevice) + .def_prop_ro("pipeline_parallel_rank", &tr::WorldConfig::getPipelineParallelRank) + .def_prop_ro("tensor_parallel_rank", &tr::WorldConfig::getTensorParallelRank) + .def_prop_ro("context_parallel_rank", &tr::WorldConfig::getContextParallelRank) + .def_prop_ro("enable_attention_dp", &tr::WorldConfig::enableAttentionDP) + .def_static("mpi", + nb::overload_cast, std::optional, + std::optional, std::optional> const&, bool>(&tr::WorldConfig::mpi), + nb::arg("gpus_per_node") = tr::WorldConfig::kDefaultGpusPerNode, nb::arg("tensor_parallelism") = nb::none(), + nb::arg("pipeline_parallelism") = nb::none(), nb::arg("context_parallelism") = nb::none(), + nb::arg("device_ids") = nb::none(), nb::arg("enable_attention_dp") = false); + + auto SamplingConfigGetState = [](tr::SamplingConfig const& config) -> nb::tuple + { + return nb::make_tuple(config.beamWidth, config.temperature, config.minLength, config.repetitionPenalty, + config.presencePenalty, config.frequencyPenalty, config.topK, config.topP, config.randomSeed, + config.topPDecay, config.topPMin, config.topPResetIds, config.beamSearchDiversityRate, config.lengthPenalty, + config.earlyStopping, config.noRepeatNgramSize, config.numReturnSequences, config.minP, + config.beamWidthArray); + }; + auto SamplingConfigSetState = [](tr::SamplingConfig& self, nb::tuple t) + { + if (t.size() != 19) + { + throw std::runtime_error("Invalid SamplingConfig state!"); + } + + tr::SamplingConfig config; + config.beamWidth = nb::cast(t[0]); + config.temperature = nb::cast>(t[1]); + config.minLength = nb::cast>(t[2]); + config.repetitionPenalty = nb::cast>(t[3]); + config.presencePenalty = nb::cast>(t[4]); + config.frequencyPenalty = nb::cast>(t[5]); + config.topK = nb::cast>(t[6]); + config.topP = nb::cast>(t[7]); + config.randomSeed = nb::cast>(t[8]); + config.topPDecay = nb::cast>(t[9]); + config.topPMin = nb::cast>(t[10]); + config.topPResetIds = nb::cast>(t[11]); + config.beamSearchDiversityRate = nb::cast>(t[12]); + config.lengthPenalty = nb::cast>(t[13]); + config.earlyStopping = nb::cast>(t[14]); + config.noRepeatNgramSize = nb::cast>(t[15]); + config.numReturnSequences = nb::cast(t[16]); + config.minP = nb::cast>(t[17]); + config.beamWidthArray = nb::cast>>(t[18]); + + new (&self) tr::SamplingConfig(config); + }; + + nb::class_(m, "SamplingConfig") + .def(nb::init(), nb::arg("beam_width") = 1) + .def(nb::init>(), + nb::arg("executor_sample_config"), nb::arg("external_draft_tokens_config") = std::nullopt) + .def_rw("beam_width", &tr::SamplingConfig::beamWidth) + .def_rw("temperature", &tr::SamplingConfig::temperature) + .def_rw("min_length", &tr::SamplingConfig::minLength) + .def_rw("repetition_penalty", &tr::SamplingConfig::repetitionPenalty) + .def_rw("presence_penalty", &tr::SamplingConfig::presencePenalty) + .def_rw("frequency_penalty", &tr::SamplingConfig::frequencyPenalty) + .def_rw("top_k", &tr::SamplingConfig::topK) + .def_rw("top_p", &tr::SamplingConfig::topP) + .def_rw("random_seed", &tr::SamplingConfig::randomSeed) + .def_rw("top_p_decay", &tr::SamplingConfig::topPDecay) + .def_rw("top_p_min", &tr::SamplingConfig::topPMin) + .def_rw("top_p_reset_ids", &tr::SamplingConfig::topPResetIds) + .def_rw("beam_search_diversity_rate", &tr::SamplingConfig::beamSearchDiversityRate) + .def_rw("length_penalty", &tr::SamplingConfig::lengthPenalty) + .def_rw("early_stopping", &tr::SamplingConfig::earlyStopping) + .def_rw("no_repeat_ngram_size", &tr::SamplingConfig::noRepeatNgramSize) + .def_rw("num_return_sequences", &tr::SamplingConfig::numReturnSequences) + .def_rw("min_p", &tr::SamplingConfig::minP) + .def_rw("beam_width_array", &tr::SamplingConfig::beamWidthArray) + .def_rw("normalize_log_probs", &tr::SamplingConfig::normalizeLogProbs) + .def("__getstate__", SamplingConfigGetState) + .def("__setstate__", SamplingConfigSetState) + .def("__eq__", &tr::SamplingConfig::operator==); + + nb::bind_vector>(m, "SamplingConfigVector"); + + m.def("make_sampling_config", &makeSamplingConfig, nb::arg("configs")); + + nb::class_(m, "GptJsonConfig") + .def(nb::init>(), + nb::arg("name"), nb::arg("version"), nb::arg("precision"), nb::arg("tensor_parallelism"), + nb::arg("pipeline_parallelism"), nb::arg("context_parallelism"), nb::arg("gpus_per_node"), + nb::arg("model_config"), nb::arg("runtime_defaults") = nb::none()) + .def_static("parse", nb::overload_cast(&tr::GptJsonConfig::parse), nb::arg("json")) + .def_static( + "parse_file", nb::overload_cast(&tr::GptJsonConfig::parse), nb::arg("path")) + .def_prop_ro("model_config", &tr::GptJsonConfig::getModelConfig) + .def_prop_ro("name", &tr::GptJsonConfig::getName) + .def_prop_ro("version", &tr::GptJsonConfig::getVersion) + .def_prop_ro("precision", &tr::GptJsonConfig::getPrecision) + .def_prop_ro("tensor_parallelism", &tr::GptJsonConfig::getTensorParallelism) + .def_prop_ro("pipeline_parallelism", &tr::GptJsonConfig::getPipelineParallelism) + .def_prop_ro("context_parallelism", &tr::GptJsonConfig::getContextParallelism) + .def_prop_ro("gpus_per_node", &tr::GptJsonConfig::getGpusPerNode) + .def_prop_ro("world_size", &tr::GptJsonConfig::getWorldSize) + .def_prop_ro("runtime_defaults", &tr::GptJsonConfig::getRuntimeDefaults) + .def("engine_filename", + nb::overload_cast( + &tr::GptJsonConfig::engineFilename, nb::const_), + nb::arg("world_config"), nb::arg("model")) + .def("engine_filename", + nb::overload_cast(&tr::GptJsonConfig::engineFilename, nb::const_), + nb::arg("world_config")); + + nb::enum_(m, "LlmRequestState") + .value("UNKNOWN", tb::LlmRequestState::kUNKNOWN) + .value("ENCODER_INIT", tb::LlmRequestState::kENCODER_INIT) + .value("CONTEXT_INIT", tb::LlmRequestState::kCONTEXT_INIT) + .value("GENERATION_IN_PROGRESS", tb::LlmRequestState::kGENERATION_IN_PROGRESS) + .value("GENERATION_TO_COMPLETE", tb::LlmRequestState::kGENERATION_TO_COMPLETE) + .value("GENERATION_COMPLETE", tb::LlmRequestState::kGENERATION_COMPLETE) + .value("DISAGG_GENERATION_INIT", tb::LlmRequestState::kDISAGG_GENERATION_INIT) + .value("DISAGG_CONTEXT_TRANS_IN_PROGRESS", tb::LlmRequestState::kDISAGG_CONTEXT_TRANS_IN_PROGRESS) + .value("DISAGG_CONTEXT_COMPLETE", tb::LlmRequestState::kDISAGG_CONTEXT_COMPLETE) + .value("DISAGG_GENERATION_TRANS_IN_PROGRESS", tb::LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS) + .value("DISAGG_GENERATION_TRANS_COMPLETE", tb::LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE) + .value("DISAGG_CONTEXT_INIT_AND_TRANS", tb::LlmRequestState::kDISAGG_CONTEXT_INIT_AND_TRANS); + + nb::class_(m, "MemoryCounters") + .def_static("instance", &tr::MemoryCounters::getInstance, nb::rv_policy::reference) + .def_prop_ro("gpu", &tr::MemoryCounters::getGpu) + .def_prop_ro("cpu", &tr::MemoryCounters::getCpu) + .def_prop_ro("pinned", &tr::MemoryCounters::getPinned) + .def_prop_ro("uvm", &tr::MemoryCounters::getUVM); + + tensorrt_llm::nanobind::runtime::initBindings(mInternalRuntime); + tensorrt_llm::nanobind::testing::initBindings(mInternalTesting); + tpb::initBindings(mInternalBatchManager); + tb::kv_cache_manager::KVCacheManagerBindings::initBindings(mInternalBatchManager); + tb::BasePeftCacheManagerBindings::initBindings(mInternalBatchManager); + tb::CacheTransceiverBindings::initBindings(mInternalBatchManager); + + auto mInternalAlgorithms = mInternal.def_submodule("algorithms", "Algorithms internal bindings"); + tpb::algorithms::initBindings(mInternalAlgorithms); + + auto mUserbuffers = mInternal.def_submodule("userbuffers", "User buffers internal bindings"); + tensorrt_llm::kernels::userbuffers::UserBufferBindings::initBindings(mUserbuffers); + + // NVLS allocators + nb::class_(m, "IpcNvlsHandle") + .def(nb::init<>()) + .def_rw("uc_ptr", &tr::IpcNvlsHandle::uc_ptr) + .def_rw("mc_ptr", &tr::IpcNvlsHandle::mc_ptr) + .def_rw("size", &tr::IpcNvlsHandle::size) + .def("get_ipc_ptrs", + [](tr::IpcNvlsHandle& self) { return reinterpret_cast(self.ipc_uc_ptrs.data()); }); + + m.def("ipc_nvls_allocate", &tr::ipcNvlsAllocate, nb::rv_policy::reference); + m.def("ipc_nvls_free", &tr::ipcNvlsFree); + m.def("ipc_nvls_supported", &tr::ipcNvlsSupported); } diff --git a/cpp/tensorrt_llm/nanobind/common/bindTypes.h b/cpp/tensorrt_llm/nanobind/common/bindTypes.h new file mode 100644 index 000000000000..6312907b88f5 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/common/bindTypes.h @@ -0,0 +1,67 @@ +/* + * 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 + +#include +#include +#include + +namespace NanobindUtils +{ + +namespace nb = nanobind; + +template +void bindSet(nb::module_& m, std::string const& name) +{ + nb::class_(m, name.c_str()) + .def(nb::init<>()) + .def("clear", &T::clear) + .def("size", &T::size) + .def("insert", [](T& s, typename T::value_type const& value) { s.insert(value); }) + .def("erase", nb::overload_cast(&T::erase)) + .def("__len__", [](T const& lst) { return lst.size(); }) + .def("__contains__", [](T const& s, typename T::value_type x) { return s.find(x) != s.end(); }) + .def( + "__iter__", [](T& s) { return nb::make_iterator(nb::type(), "iterator", s.begin(), s.end()); }, + nb::keep_alive<0, 1>()) + .def("__eq__", [](T const& s, T const& other) { return s == other; }) + .def("__getstate__", + [](T const& v) + { + /* Return a tuple that fully encodes the state of the object */ + return nb::make_tuple(std::vector(v.begin(), v.end())); + }) + .def("__setstate__", + [](T& v, nb::tuple const& t) + { + if (t.size() != 1) + throw std::runtime_error("Invalid state!"); + /* Create a new C++ instance */ + T s; + /* Assign any additional state */ + auto state_list = nb::cast>(t[0]); + for (auto& item : state_list) + { + s.insert(item); + } + new (&v) T(s); + }); +} + +} // namespace NanobindUtils diff --git a/cpp/tensorrt_llm/nanobind/common/customCasters.h b/cpp/tensorrt_llm/nanobind/common/customCasters.h new file mode 100644 index 000000000000..2739ccd569ed --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/common/customCasters.h @@ -0,0 +1,292 @@ +/* + * 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 + +#include "tensorrt_llm/batch_manager/common.h" +#include "tensorrt_llm/batch_manager/decoderBuffers.h" +#include "tensorrt_llm/common/optionalRef.h" +#include "tensorrt_llm/runtime/cudaStream.h" +#include "tensorrt_llm/runtime/request.h" +#include "tensorrt_llm/runtime/samplingConfig.h" +#include "tensorrt_llm/runtime/torch.h" +#include "tensorrt_llm/runtime/torchView.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Pybind requires to have a central include in order for type casters to work. +// Opaque bindings add a type caster, so they have the same requirement. +// See the warning in https://pybind11.readthedocs.io/en/stable/advanced/cast/custom.html + +// Opaque bindings +NB_MAKE_OPAQUE(tensorrt_llm::batch_manager::ReqIdsSet) +NB_MAKE_OPAQUE(std::vector) +NB_MAKE_OPAQUE(std::vector) +NB_MAKE_OPAQUE(std::vector) + +namespace nb = nanobind; + +// Custom casters +namespace NB_NAMESPACE +{ + +namespace detail +{ + +template +struct type_caster> +{ + using Type = std::deque; + NB_TYPE_CASTER(Type, const_name("List")); + + bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) noexcept + { + sequence seq(src, nanobind::detail::borrow_t{}); + value.clear(); + make_caster caster; + for (auto const& item : seq) + { + if (!caster.from_python(item, flags, cleanup)) + return false; + value.push_back(caster.operator T&()); + } + return true; + } + + static handle from_cpp(Type const& deque, rv_policy policy, cleanup_list* cleanup) noexcept + { + nb::list list; + + for (auto const& item : deque) + { + nb::object py_item = steal(make_caster::from_cpp(item, policy, cleanup)); + if (!py_item) + return {}; + list.append(py_item); + } + return list.release(); + } +}; + +template +struct type_caster> +{ + using value_conv = make_caster; + + NB_TYPE_CASTER(tensorrt_llm::common::OptionalRef, value_conv::Name); + + bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) + { + if (src.is_none()) + { + // If the Python object is None, create an empty OptionalRef + value = tensorrt_llm::common::OptionalRef(); + return true; + } + + value_conv conv; + if (!conv.from_python(src, flags, cleanup)) + return false; + + // Create an OptionalRef with a reference to the converted value + value = tensorrt_llm::common::OptionalRef(conv); + return true; + } + + static handle from_cpp(tensorrt_llm::common::OptionalRef const& src, rv_policy policy, cleanup_list* cleanup) + { + if (!src.has_value()) + return none().release(); + + return value_conv::from_cpp(*src, policy, cleanup); + } +}; + +template <> +class type_caster +{ +public: + NB_TYPE_CASTER(tensorrt_llm::executor::StreamPtr, const_name("int")); + + bool from_python([[maybe_unused]] handle src, uint8_t flags, cleanup_list* cleanup) + { + auto stream_ptr = nanobind::cast(src); + value = std::make_shared(reinterpret_cast(stream_ptr)); + + return true; + } + + static handle from_cpp( + tensorrt_llm::executor::StreamPtr const& src, rv_policy /* policy */, cleanup_list* /* cleanup */) + { + // Return cudaStream_t as integer. + return PyLong_FromVoidPtr(src->get()); + } +}; + +template <> +struct type_caster +{ +public: + NB_TYPE_CASTER(tensorrt_llm::executor::Tensor, const_name("torch.Tensor")); + + // Convert PyObject(torch.Tensor) -> tensorrt_llm::executor::Tensor + bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) + { + PyObject* obj = src.ptr(); + if (THPVariable_Check(obj)) + { + at::Tensor const& t = THPVariable_Unpack(obj); + value = tensorrt_llm::executor::detail::ofITensor(tensorrt_llm::runtime::TorchView::of(t)); + return true; + } + return false; + } + + // Convert tensorrt_llm::executor::Tensor -> PyObject(torch.Tensor) + static handle from_cpp( + tensorrt_llm::executor::Tensor const& src, rv_policy /* policy */, cleanup_list* /* cleanup */) + { + return THPVariable_Wrap(tensorrt_llm::runtime::Torch::tensor(tensorrt_llm::executor::detail::toITensor(src))); + } +}; + +template <> +struct type_caster +{ +public: + NB_TYPE_CASTER(tensorrt_llm::runtime::ITensor::SharedPtr, const_name("torch.Tensor")); + + // Convert PyObject(torch.Tensor) -> tensorrt_llm::runtime::ITensor::SharedPtr + bool from_python(handle src, uint8_t, cleanup_list*) + { + PyObject* obj = src.ptr(); + if (THPVariable_Check(obj)) + { + at::Tensor const& t = THPVariable_Unpack(obj); + value = std::move(tensorrt_llm::runtime::TorchView::of(t)); + return true; + } + return false; + } + + // Convert tensorrt_llm::runtime::ITensor::SharedPtr -> PyObject(torch.Tensor) + static handle from_cpp( + tensorrt_llm::runtime::ITensor::SharedPtr const& src, rv_policy /* policy */, cleanup_list* /* cleanup */) + { + if (src == nullptr) + { + return none().release(); + } + return THPVariable_Wrap(tensorrt_llm::runtime::Torch::tensor(src)); + } +}; + +template <> +struct type_caster +{ +public: + NB_TYPE_CASTER(tensorrt_llm::runtime::ITensor::SharedConstPtr, const_name("torch.Tensor")); + + // Convert PyObject(torch.Tensor) -> tensorrt_llm::runtime::ITensor::SharedConstPtr + bool from_python(handle src, uint8_t, cleanup_list*) + { + PyObject* obj = src.ptr(); + if (THPVariable_Check(obj)) + { + at::Tensor const& t = THPVariable_Unpack(obj); + value = std::move(tensorrt_llm::runtime::TorchView::of(t)); + return true; + } + return false; + } + + // Convert tensorrt_llm::runtime::ITensor::SharedConstPtr -> PyObject(torch.Tensor) + static handle from_cpp( + tensorrt_llm::runtime::ITensor::SharedConstPtr const& src, rv_policy /* policy */, cleanup_list* /* cleanup */) + { + if (src == nullptr) + { + return none().release(); + } + return THPVariable_Wrap(tensorrt_llm::runtime::Torch::tensor( + reinterpret_cast(src))); + } +}; + +template <> +struct type_caster +{ + NB_TYPE_CASTER(at::Tensor, const_name("torch.Tensor")); + + bool from_python(nb::handle src, uint8_t, cleanup_list*) noexcept + { + PyObject* obj = src.ptr(); + if (THPVariable_Check(obj)) + { + value = THPVariable_Unpack(obj); + return true; + } + return false; + } + + static handle from_cpp(at::Tensor src, rv_policy, cleanup_list*) noexcept + { + return THPVariable_Wrap(src); + } +}; + +template +struct type_caster>> +{ + using VectorType = std::vector>; + + NB_TYPE_CASTER(VectorType, const_name("List[") + make_caster::Name + const_name("]")); + + bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) noexcept + { + // Not needed for our use case since we only convert C++ to Python + return false; + } + + static handle from_cpp(VectorType const& src, rv_policy policy, cleanup_list* cleanup) noexcept + { + + std::vector result; + result.reserve(src.size()); + for (auto const& ref : src) + { + result.push_back(ref.get()); + } + + return make_caster>::from_cpp(result, policy, cleanup); + } +}; +} // namespace detail +} // namespace NB_NAMESPACE diff --git a/cpp/tensorrt_llm/nanobind/executor/bindings.cpp b/cpp/tensorrt_llm/nanobind/executor/bindings.cpp new file mode 100644 index 000000000000..d3f482df8997 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/executor/bindings.cpp @@ -0,0 +1,263 @@ +/* + * 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 "bindings.h" +#include "executor.h" +#include "executorConfig.h" +#include "request.h" +#include "tensorrt_llm/executor/executor.h" +#include "tensorrt_llm/executor/types.h" +#include "tensorrt_llm/nanobind/common/customCasters.h" + +#include +#include +#include +#include +#include +#include + +namespace nb = nanobind; +namespace tle = tensorrt_llm::executor; +using SizeType32 = tle::SizeType32; + +namespace tensorrt_llm::nanobind::executor +{ + +template +void instantiateEventDiff(nb::module_& m, std::string const& name) +{ + nb::class_>(m, ("KVCacheEventDiff" + name).c_str()) + .def_ro("old_value", &tle::KVCacheEventDiff::oldValue) + .def_ro("new_value", &tle::KVCacheEventDiff::newValue); +} + +void initBindings(nb::module_& m) +{ + m.attr("__version__") = tle::version(); + nb::enum_(m, "ModelType") + .value("DECODER_ONLY", tle::ModelType::kDECODER_ONLY) + .value("ENCODER_ONLY", tle::ModelType::kENCODER_ONLY) + .value("ENCODER_DECODER", tle::ModelType::kENCODER_DECODER); + + auto decodingModeGetstate = [](tle::DecodingMode const& self) { return nb::make_tuple(self.getState()); }; + auto decodingModeSetstate = [](tle::DecodingMode& self, nb::tuple const& state) + { + if (state.size() != 1) + { + throw std::runtime_error("Invalid state!"); + } + new (&self) tle::DecodingMode(nb::cast(state[0])); + }; + nb::class_(m, "DecodingMode") + .def("Auto", &tle::DecodingMode::Auto) + .def("TopK", &tle::DecodingMode::TopK) + .def("TopP", &tle::DecodingMode::TopP) + .def("TopKTopP", &tle::DecodingMode::TopKTopP) + .def("BeamSearch", &tle::DecodingMode::BeamSearch) + .def("Medusa", &tle::DecodingMode::Medusa) + .def("Lookahead", &tle::DecodingMode::Lookahead) + .def("ExplicitDraftTokens", &tle::DecodingMode::ExplicitDraftTokens) + .def("Eagle", &tle::DecodingMode::Eagle) + .def("isAuto", &tle::DecodingMode::isAuto) + .def("isTopK", &tle::DecodingMode::isTopK) + .def("isTopP", &tle::DecodingMode::isTopP) + .def("isTopKorTopP", &tle::DecodingMode::isTopKorTopP) + .def("isTopKandTopP", &tle::DecodingMode::isTopKandTopP) + .def("isBeamSearch", &tle::DecodingMode::isBeamSearch) + .def("isMedusa", &tle::DecodingMode::isMedusa) + .def("isLookahead", &tle::DecodingMode::isLookahead) + .def("isExplicitDraftTokens", &tle::DecodingMode::isExplicitDraftTokens) + .def("isEagle", &tle::DecodingMode::isEagle) + .def("useVariableBeamWidthSearch", &tle::DecodingMode::useVariableBeamWidthSearch) + .def_prop_ro("name", &tle::DecodingMode::getName) + .def("__getstate__", decodingModeGetstate) + .def("__setstate__", decodingModeSetstate); + + nb::enum_(m, "CapacitySchedulerPolicy") + .value("MAX_UTILIZATION", tle::CapacitySchedulerPolicy::kMAX_UTILIZATION) + .value("GUARANTEED_NO_EVICT", tle::CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT) + .value("STATIC_BATCH", tle::CapacitySchedulerPolicy::kSTATIC_BATCH); + + nb::enum_(m, "ContextChunkingPolicy") + .value("EQUAL_PROGRESS", tle::ContextChunkingPolicy::kEQUAL_PROGRESS) + .value("FIRST_COME_FIRST_SERVED", tle::ContextChunkingPolicy::kFIRST_COME_FIRST_SERVED); + + nb::enum_(m, "CommunicationType").value("MPI", tle::CommunicationType::kMPI); + + nb::enum_(m, "CommunicationMode") + .value("LEADER", tle::CommunicationMode::kLEADER) + .value("ORCHESTRATOR", tle::CommunicationMode::kORCHESTRATOR); + + nb::class_(m, "KvCacheStats") + .def(nb::init<>()) + .def_rw("max_num_blocks", &tle::KvCacheStats::maxNumBlocks) + .def_rw("free_num_blocks", &tle::KvCacheStats::freeNumBlocks) + .def_rw("used_num_blocks", &tle::KvCacheStats::usedNumBlocks) + .def_rw("tokens_per_block", &tle::KvCacheStats::tokensPerBlock) + .def_rw("alloc_total_blocks", &tle::KvCacheStats::allocTotalBlocks) + .def_rw("alloc_new_blocks", &tle::KvCacheStats::allocNewBlocks) + .def_rw("reused_blocks", &tle::KvCacheStats::reusedBlocks) + .def_rw("missed_blocks", &tle::KvCacheStats::missedBlocks) + .def_rw("cache_hit_rate", &tle::KvCacheStats::cacheHitRate); + + nb::class_(m, "StaticBatchingStats") + .def(nb::init<>()) + .def_rw("num_scheduled_requests", &tle::StaticBatchingStats::numScheduledRequests) + .def_rw("num_context_requests", &tle::StaticBatchingStats::numContextRequests) + .def_rw("num_ctx_tokens", &tle::StaticBatchingStats::numCtxTokens) + .def_rw("num_gen_tokens", &tle::StaticBatchingStats::numGenTokens) + .def_rw("empty_gen_slots", &tle::StaticBatchingStats::emptyGenSlots); + + nb::class_(m, "InflightBatchingStats") + .def(nb::init<>()) + .def_rw("num_scheduled_requests", &tle::InflightBatchingStats::numScheduledRequests) + .def_rw("num_context_requests", &tle::InflightBatchingStats::numContextRequests) + .def_rw("num_gen_requests", &tle::InflightBatchingStats::numGenRequests) + .def_rw("num_paused_requests", &tle::InflightBatchingStats::numPausedRequests) + .def_rw("num_ctx_tokens", &tle::InflightBatchingStats::numCtxTokens) + .def_rw("micro_batch_id", &tle::InflightBatchingStats::microBatchId) + .def_rw("avg_num_decoded_tokens_per_iter", &tle::InflightBatchingStats::avgNumDecodedTokensPerIter); + + nb::class_(m, "SpecDecodingStats") + .def(nb::init<>()) + .def_rw("num_draft_tokens", &tle::SpecDecodingStats::numDraftTokens) + .def_rw("num_accepted_tokens", &tle::SpecDecodingStats::numAcceptedTokens) + .def_rw("num_requests_with_draft_tokens", &tle::SpecDecodingStats::numRequestsWithDraftTokens) + .def_rw("acceptance_length", &tle::SpecDecodingStats::acceptanceLength) + .def_rw("iter_latency_ms", &tle::SpecDecodingStats::iterLatencyMS) + .def_rw("draft_overhead", &tle::SpecDecodingStats::draftOverhead); + + nb::class_(m, "IterationStats") + .def(nb::init<>()) + .def_rw("timestamp", &tle::IterationStats::timestamp) + .def_rw("iter", &tle::IterationStats::iter) + .def_rw("iter_latency_ms", &tle::IterationStats::iterLatencyMS) + .def_rw("new_active_requests_queue_latency_ms", &tle::IterationStats::newActiveRequestsQueueLatencyMS) + .def_rw("num_new_active_requests", &tle::IterationStats::numNewActiveRequests) + .def_rw("num_active_requests", &tle::IterationStats::numActiveRequests) + .def_rw("num_queued_requests", &tle::IterationStats::numQueuedRequests) + .def_rw("num_completed_requests", &tle::IterationStats::numCompletedRequests) + .def_rw("max_num_active_requests", &tle::IterationStats::maxNumActiveRequests) + .def_rw("gpu_mem_usage", &tle::IterationStats::gpuMemUsage) + .def_rw("cpu_mem_usage", &tle::IterationStats::cpuMemUsage) + .def_rw("pinned_mem_usage", &tle::IterationStats::pinnedMemUsage) + .def_rw("kv_cache_stats", &tle::IterationStats::kvCacheStats) + .def_rw("cross_kv_cache_stats", &tle::IterationStats::crossKvCacheStats) + .def_rw("static_batching_stats", &tle::IterationStats::staticBatchingStats) + .def_rw("inflight_batching_stats", &tle::IterationStats::inflightBatchingStats) + .def_rw("specdec_stats", &tle::IterationStats::specDecodingStats) + .def("to_json_str", + [](tle::IterationStats const& iterationStats) + { return tle::JsonSerialization::toJsonStr(iterationStats); }); + + nb::class_(m, "DebugTensorsPerIteration") + .def(nb::init<>()) + .def_rw("iter", &tle::DebugTensorsPerIteration::iter) + .def_rw("debug_tensors", &tle::DebugTensorsPerIteration::debugTensors); + + nb::enum_(m, "RequestStage") + .value("QUEUED", tle::RequestStage::kQUEUED) + .value("ENCODER_IN_PROGRESS", tle::RequestStage::kENCODER_IN_PROGRESS) + .value("CONTEXT_IN_PROGRESS", tle::RequestStage::kCONTEXT_IN_PROGRESS) + .value("GENERATION_IN_PROGRESS", tle::RequestStage::kGENERATION_IN_PROGRESS) + .value("GENERATION_COMPLETE", tle::RequestStage::kGENERATION_COMPLETE); + + nb::class_(m, "DisServingRequestStats") + .def(nb::init<>()) + .def_rw("kv_cache_transfer_ms", &tle::DisServingRequestStats::kvCacheTransferMS) + .def_rw("kv_cache_size", &tle::DisServingRequestStats::kvCacheSize); + + nb::class_(m, "RequestStats") + .def(nb::init<>()) + .def_rw("id", &tle::RequestStats::id) + .def_rw("stage", &tle::RequestStats::stage) + .def_rw("context_prefill_position", &tle::RequestStats::contextPrefillPosition) + .def_rw("num_generated_tokens", &tle::RequestStats::numGeneratedTokens) + .def_rw("avg_num_decoded_tokens_per_iter", &tle::RequestStats::avgNumDecodedTokensPerIter) + .def_rw("scheduled", &tle::RequestStats::scheduled) + .def_rw("paused", &tle::RequestStats::paused) + .def_rw("dis_serving_stats", &tle::RequestStats::disServingStats) + .def_rw("alloc_total_blocks_per_request", &tle::RequestStats::allocTotalBlocksPerRequest) + .def_rw("alloc_new_blocks_per_request", &tle::RequestStats::allocNewBlocksPerRequest) + .def_rw("reused_blocks_per_request", &tle::RequestStats::reusedBlocksPerRequest) + .def_rw("missed_blocks_per_request", &tle::RequestStats::missedBlocksPerRequest) + .def_rw("kv_cache_hit_rate_per_request", &tle::RequestStats::kvCacheHitRatePerRequest) + .def("to_json_str", + [](tle::RequestStats const& iterationStats) { return tle::JsonSerialization::toJsonStr(iterationStats); }); + + nb::class_(m, "RequestStatsPerIteration") + .def(nb::init<>()) + .def_rw("iter", &tle::RequestStatsPerIteration::iter) + .def_rw("request_stats", &tle::RequestStatsPerIteration::requestStats) + .def("to_json_str", + [](tle::RequestStatsPerIteration const& iterationStats) + { return tle::JsonSerialization::toJsonStr(iterationStats); }); + + nb::module_ executor_kv_cache = m.def_submodule("kv_cache", "Executor KV Cache Manager"); + + nb::class_(executor_kv_cache, "KVCacheCreatedData") + .def_ro("num_blocks_per_cache_level", &tle::KVCacheCreatedData::numBlocksPerCacheLevel); + + nb::class_(executor_kv_cache, "UniqueToken") + .def_ro("token_id", &tensorrt_llm::runtime::UniqueToken::tokenId) + .def_ro("token_extra_id", &tensorrt_llm::runtime::UniqueToken::tokenExtraId); + + nb::class_(executor_kv_cache, "KVCacheStoredBlockData") + .def_ro("block_hash", &tle::KVCacheStoredBlockData::blockHash) + .def_ro("tokens", &tle::KVCacheStoredBlockData::tokens) + .def_ro("lora_id", &tle::KVCacheStoredBlockData::loraId) + .def_ro("cache_level", &tle::KVCacheStoredBlockData::cacheLevel) + .def_ro("priority", &tle::KVCacheStoredBlockData::priority); + + nb::class_(executor_kv_cache, "KVCacheStoredData") + .def_ro("parent_hash", &tle::KVCacheStoredData::parentHash) + .def_ro("blocks", &tle::KVCacheStoredData::blocks); + + nb::class_(executor_kv_cache, "KVCacheRemovedData") + .def_ro("block_hashes", &tle::KVCacheRemovedData::blockHashes); + + instantiateEventDiff(executor_kv_cache, "Int"); + + nb::class_(executor_kv_cache, "KVCacheUpdatedData") + .def_ro("block_hash", &tle::KVCacheUpdatedData::blockHash) + .def_ro("cache_level", &tle::KVCacheUpdatedData::cacheLevel) + .def_ro("priority", &tle::KVCacheUpdatedData::priority); + + nb::class_(executor_kv_cache, "KVCacheEvent") + .def_ro("event_id", &tle::KVCacheEvent::eventId) + .def_ro("data", &tle::KVCacheEvent::data) + .def_ro("window_size", &tle::KVCacheEvent::windowSize); + + nb::class_(executor_kv_cache, "KVCacheEventManager") + .def( + "get_latest_events", + [](tle::KVCacheEventManager& self, std::optional timeout_ms = std::nullopt) + { + if (timeout_ms) + { + return self.getLatestEvents(std::chrono::milliseconds(static_cast(*timeout_ms))); + } + return self.getLatestEvents(std::nullopt); + }, + nb::arg("timeout_ms") = std::nullopt); + + 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/bindings.h b/cpp/tensorrt_llm/nanobind/executor/bindings.h new file mode 100644 index 000000000000..4df52c2d34e4 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/executor/bindings.h @@ -0,0 +1,29 @@ +/* + * 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::executor +{ + +// Register bindings for executor API. +void initBindings(nb::module_& m); + +} // namespace tensorrt_llm::nanobind::executor diff --git a/cpp/tensorrt_llm/nanobind/executor/executor.cpp b/cpp/tensorrt_llm/nanobind/executor/executor.cpp new file mode 100644 index 000000000000..5b916c4b1847 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/executor/executor.cpp @@ -0,0 +1,225 @@ +/* + * 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) + .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 new file mode 100644 index 000000000000..22c24abb4bfd --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/executor/executor.h @@ -0,0 +1,129 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/executor/executor.h" +#include "tensorrt_llm/executor/types.h" +#include + +namespace nb = nanobind; +namespace tle = tensorrt_llm::executor; + +namespace tensorrt_llm::nanobind::executor +{ + +class Executor +{ +public: + Executor( + std::filesystem::path const& modelPath, tle::ModelType modelType, tle::ExecutorConfig const& executorConfig); + + Executor(std::filesystem::path const& encoderModelPath, std::filesystem::path const& decoderModelPath, + tle::ModelType modelType, tle::ExecutorConfig const& executorConfig); + + Executor(nb::bytes const& engineBuffer, std::string const& jsonConfigStr, tle::ModelType modelType, + tle::ExecutorConfig const& executorConfig, std::optional managedWeights); + + Executor(std::string const& encoderEngineBuffer, std::string const& encoderJsonConfigStr, + std::string const& decoderEngineBuffer, std::string const& decoderJsonConfigStr, tle::ModelType modelType, + tle::ExecutorConfig const& executorConfig); + + nb::object enter(); + void exit( + [[maybe_unused]] nb::handle type, [[maybe_unused]] nb::handle value, [[maybe_unused]] nb::handle traceback); + void shutdown(); + + [[nodiscard]] tle::IdType enqueueRequest(tle::Request const& request) + { + return mExecutor->enqueueRequest(request); + } + + [[nodiscard]] std::vector enqueueRequests(std::vector const& requests) + { + return mExecutor->enqueueRequests(requests); + } + + [[nodiscard]] std::vector awaitResponses( + std::optional const& timeout = std::nullopt) + { + // Await responses blocks until a response is received. Release GIL so that it can be ran in a background + // thread. + nb::gil_scoped_release release; + return mExecutor->awaitResponses(timeout); + } + + [[nodiscard]] std::vector awaitResponses( + tle::IdType const& requestId, std::optional const& timeout = std::nullopt) + { + // Await responses blocks until a response is received. Release GIL so that it can be ran in a background + // thread. + nb::gil_scoped_release release; + return mExecutor->awaitResponses(requestId, timeout); + } + + [[nodiscard]] std::vector> awaitResponses(std::vector const& requestIds, + std::optional const& timeout = std::nullopt) + { + // Await responses blocks until a response is received. Release GIL so that it can be ran in a background + // thread. + nb::gil_scoped_release release; + return mExecutor->awaitResponses(requestIds, timeout); + } + + [[nodiscard]] tle::SizeType32 getNumResponsesReady(std::optional const& requestId = std::nullopt) const + { + return mExecutor->getNumResponsesReady(requestId); + } + + void cancelRequest(tle::IdType requestId) + { + mExecutor->cancelRequest(requestId); + } + + std::deque getLatestIterationStats() + { + return mExecutor->getLatestIterationStats(); + } + + std::deque getLatestRequestStats() + { + return mExecutor->getLatestRequestStats(); + } + + std::deque getLatestDebugTensors() + { + return mExecutor->getLatestDebugTensors(); + } + + [[nodiscard]] bool canEnqueueRequests() const + { + return mExecutor->canEnqueueRequests(); + } + + [[nodiscard]] std::optional> getKVCacheEventManager() const + { + return mExecutor->getKVCacheEventManager(); + } + + static void initBindings(nb::module_& m); + +private: + std::unique_ptr mExecutor; +}; + +} // namespace tensorrt_llm::nanobind::executor diff --git a/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp new file mode 100644 index 000000000000..6e7adde2cd3f --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp @@ -0,0 +1,639 @@ +/* + * 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 "executorConfig.h" +#include "tensorrt_llm/executor/executor.h" +#include "tensorrt_llm/executor/types.h" +#include "tensorrt_llm/nanobind/common/customCasters.h" +#include "tensorrt_llm/runtime/cudaStream.h" +#include "tensorrt_llm/runtime/utils/mpiUtils.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace nb = nanobind; +namespace tle = tensorrt_llm::executor; +using SizeType32 = tle::SizeType32; +using RuntimeDefaults = tensorrt_llm::runtime::RuntimeDefaults; + +namespace tensorrt_llm::nanobind::executor +{ + +void initConfigBindings(nb::module_& m) +{ + nb::enum_(m, "BatchingType") + .value("STATIC", tle::BatchingType::kSTATIC) + .value("INFLIGHT", tle::BatchingType::kINFLIGHT); + + auto dynamicBatchConfigGetstate = [](tle::DynamicBatchConfig const& self) + { + return nb::make_tuple(self.getEnableBatchSizeTuning(), self.getEnableMaxNumTokensTuning(), + self.getDynamicBatchMovingAverageWindow(), self.getBatchSizeTable()); + }; + auto dynamicBatchConfigSetstate = [](tle::DynamicBatchConfig& self, nb::tuple const& state) + { + if (state.size() != 4) + { + throw std::runtime_error("Invalid state!"); + } + new (&self) tle::DynamicBatchConfig(nb::cast(state[0]), nb::cast(state[1]), + nb::cast(state[2]), nb::cast>>(state[3])); + }; + nb::class_(m, "DynamicBatchConfig") + .def(nb::init(), nb::arg("enable_batch_size_tuning"), + nb::arg("enable_max_num_tokens_tuning"), nb::arg("dynamic_batch_moving_average_window")) + .def_prop_ro("enable_batch_size_tuning", &tle::DynamicBatchConfig::getEnableBatchSizeTuning) + .def_prop_ro("enable_max_num_tokens_tuning", &tle::DynamicBatchConfig::getEnableMaxNumTokensTuning) + .def_prop_ro( + "dynamic_batch_moving_average_window", &tle::DynamicBatchConfig::getDynamicBatchMovingAverageWindow) + .def("__getstate__", dynamicBatchConfigGetstate) + .def("__setstate__", dynamicBatchConfigSetstate); + + auto schedulerConfigSetstate = [](tle::SchedulerConfig& self, nb::tuple const& state) + { + if (state.size() != 3) + { + throw std::runtime_error("Invalid state!"); + } + new (&self) tle::SchedulerConfig(nb::cast(state[0]), + nb::cast>(state[1]), + nb::cast>(state[2])); + }; + auto schedulerConfigGetstate = [](tle::SchedulerConfig const& self) + { + return nb::make_tuple( + self.getCapacitySchedulerPolicy(), self.getContextChunkingPolicy(), self.getDynamicBatchConfig()); + }; + nb::class_(m, "SchedulerConfig") + .def(nb::init, + std::optional>(), + nb::arg("capacity_scheduler_policy") = tle::CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT, + nb::arg("context_chunking_policy") = nb::none(), nb::arg("dynamic_batch_config") = nb::none()) + .def_prop_ro("capacity_scheduler_policy", &tle::SchedulerConfig::getCapacitySchedulerPolicy) + .def_prop_ro("context_chunking_policy", &tle::SchedulerConfig::getContextChunkingPolicy) + .def_prop_ro("dynamic_batch_config", &tle::SchedulerConfig::getDynamicBatchConfig) + .def("__getstate__", schedulerConfigGetstate) + .def("__setstate__", schedulerConfigSetstate); + + nb::class_(m, "RuntimeDefaults") + .def(nb::init>, std::optional>(), + nb::arg("max_attention_window") = nb::none(), nb::arg("sink_token_length") = nb::none()) + .def_ro("max_attention_window", &RuntimeDefaults::maxAttentionWindowVec) + .def_ro("sink_token_length", &RuntimeDefaults::sinkTokenLength); + + auto kvCacheConfigGetstate = [](tle::KvCacheConfig const& self) + { + return nb::make_tuple(self.getEnableBlockReuse(), self.getMaxTokens(), self.getMaxAttentionWindowVec(), + self.getSinkTokenLength(), self.getFreeGpuMemoryFraction(), self.getHostCacheSize(), + self.getOnboardBlocks(), self.getCrossKvCacheFraction(), self.getSecondaryOffloadMinPriority(), + self.getEventBufferMaxSize(), self.getEnablePartialReuse(), self.getCopyOnPartialReuse(), self.getUseUvm()); + }; + auto kvCacheConfigSetstate = [](tle::KvCacheConfig& self, nb::tuple const& state) + { + if (state.size() != 13) + { + throw std::runtime_error("Invalid state!"); + } + new (&self) tle::KvCacheConfig(nb::cast(state[0]), nb::cast>(state[1]), + nb::cast>>(state[2]), nb::cast>(state[3]), + nb::cast>(state[4]), nb::cast>(state[5]), + nb::cast(state[6]), nb::cast>(state[7]), + nb::cast>(state[8]), nb::cast(state[9]), + nb::cast(state[10]), nb::cast(state[11]), nb::cast(state[12])); + }; + nb::class_(m, "KvCacheConfig") + .def(nb::init const&, std::optional> const&, + std::optional const&, std::optional const&, std::optional const&, bool, + std::optional const&, std::optional, size_t const&, bool, bool, bool, + std::optional const&>(), + nb::arg("enable_block_reuse") = true, nb::arg("max_tokens") = nb::none(), + nb::arg("max_attention_window") = nb::none(), nb::arg("sink_token_length") = nb::none(), + nb::arg("free_gpu_memory_fraction") = nb::none(), nb::arg("host_cache_size") = nb::none(), + nb::arg("onboard_blocks") = true, nb::arg("cross_kv_cache_fraction") = nb::none(), + nb::arg("secondary_offload_min_priority") = nb::none(), nb::arg("event_buffer_max_size") = 0, nb::kw_only(), + nb::arg("enable_partial_reuse") = true, nb::arg("copy_on_partial_reuse") = true, nb::arg("use_uvm") = false, + nb::arg("runtime_defaults") = nb::none()) + .def_prop_rw( + "enable_block_reuse", &tle::KvCacheConfig::getEnableBlockReuse, &tle::KvCacheConfig::setEnableBlockReuse) + .def_prop_rw("max_tokens", &tle::KvCacheConfig::getMaxTokens, &tle::KvCacheConfig::setMaxTokens) + .def_prop_rw("max_attention_window", &tle::KvCacheConfig::getMaxAttentionWindowVec, + &tle::KvCacheConfig::setMaxAttentionWindowVec) + .def_prop_rw( + "sink_token_length", &tle::KvCacheConfig::getSinkTokenLength, &tle::KvCacheConfig::setSinkTokenLength) + .def_prop_rw("free_gpu_memory_fraction", &tle::KvCacheConfig::getFreeGpuMemoryFraction, + &tle::KvCacheConfig::setFreeGpuMemoryFraction) + .def_prop_rw("host_cache_size", &tle::KvCacheConfig::getHostCacheSize, &tle::KvCacheConfig::setHostCacheSize) + .def_prop_rw("onboard_blocks", &tle::KvCacheConfig::getOnboardBlocks, &tle::KvCacheConfig::setOnboardBlocks) + .def_prop_rw("cross_kv_cache_fraction", &tle::KvCacheConfig::getCrossKvCacheFraction, + &tle::KvCacheConfig::setCrossKvCacheFraction) + .def_prop_rw("secondary_offload_min_priority", &tle::KvCacheConfig::getSecondaryOffloadMinPriority, + &tle::KvCacheConfig::setSecondaryOffloadMinPriority) + .def_prop_rw("event_buffer_max_size", &tle::KvCacheConfig::getEventBufferMaxSize, + &tle::KvCacheConfig::setEventBufferMaxSize) + .def_prop_rw("enable_partial_reuse", &tle::KvCacheConfig::getEnablePartialReuse, + &tle::KvCacheConfig::setEnablePartialReuse) + .def_prop_rw("copy_on_partial_reuse", &tle::KvCacheConfig::getCopyOnPartialReuse, + &tle::KvCacheConfig::setCopyOnPartialReuse) + .def_prop_rw("use_uvm", &tle::KvCacheConfig::getUseUvm, &tle::KvCacheConfig::setUseUvm) + .def("fill_empty_fields_from_runtime_defaults", &tle::KvCacheConfig::fillEmptyFieldsFromRuntimeDefaults) + .def("__getstate__", kvCacheConfigGetstate) + .def("__setstate__", kvCacheConfigSetstate); + + nb::class_(m, "OrchestratorConfig") + .def(nb::init, bool>(), nb::arg("is_orchestrator") = true, + nb::arg("worker_executable_path") = "", nb::arg("orch_leader_comm").none() = nullptr, + nb::arg("spawn_processes") = true) + .def_prop_rw( + "is_orchestrator", &tle::OrchestratorConfig::getIsOrchestrator, &tle::OrchestratorConfig::setIsOrchestrator) + .def_prop_rw("worker_executable_path", &tle::OrchestratorConfig::getWorkerExecutablePath, + &tle::OrchestratorConfig::setWorkerExecutablePath) + .def_prop_rw("orch_leader_comm", &tle::OrchestratorConfig::getOrchLeaderComm, + &tle::OrchestratorConfig::setOrchLeaderComm) + .def_prop_rw("spawn_processes", &tle::OrchestratorConfig::getSpawnProcesses, + &tle::OrchestratorConfig::setSpawnProcesses); + + auto parallelConfigGetstate = [](tle::ParallelConfig const& self) + { + return nb::make_tuple(self.getCommunicationType(), self.getCommunicationMode(), self.getDeviceIds(), + self.getParticipantIds(), self.getOrchestratorConfig(), self.getNumNodes()); + }; + auto parallelConfigSetstate = [](tle::ParallelConfig& self, nb::tuple const& state) + { + if (state.size() != 6) + { + throw std::runtime_error("Invalid state!"); + } + new (&self) tle::ParallelConfig(nb::cast(state[0]), + nb::cast(state[1]), nb::cast>>(state[2]), + nb::cast>>(state[3]), + nb::cast>(state[4]), nb::cast>(state[5])); + }; + nb::class_(m, "ParallelConfig") + .def(nb::init> const&, + std::optional> const&, std::optional const&, + std::optional const&>(), + nb::arg("communication_type") = tle::CommunicationType::kMPI, + nb::arg("communication_mode") = tle::CommunicationMode::kLEADER, nb::arg("device_ids") = nb::none(), + nb::arg("participant_ids") = nb::none(), nb::arg("orchestrator_config") = nb::none(), + nb::arg("num_nodes") = nb::none()) + .def_prop_rw("communication_type", &tle::ParallelConfig::getCommunicationType, + &tle::ParallelConfig::setCommunicationType) + .def_prop_rw("communication_mode", &tle::ParallelConfig::getCommunicationMode, + &tle::ParallelConfig::setCommunicationMode) + .def_prop_rw("device_ids", &tle::ParallelConfig::getDeviceIds, &tle::ParallelConfig::setDeviceIds) + .def_prop_rw( + "participant_ids", &tle::ParallelConfig::getParticipantIds, &tle::ParallelConfig::setParticipantIds) + .def_prop_rw("orchestrator_config", &tle::ParallelConfig::getOrchestratorConfig, + &tle::ParallelConfig::setOrchestratorConfig) + .def_prop_rw("num_nodes", &tle::ParallelConfig::getNumNodes, &tle::ParallelConfig::setNumNodes) + .def("__getstate__", parallelConfigGetstate) + .def("__setstate__", parallelConfigSetstate); + + auto peftCacheConfigSetstate = [](tle::PeftCacheConfig& self, nb::tuple const& state) + { + if (state.size() != 11) + { + throw std::runtime_error("Invalid state!"); + } + new (&self) tle::PeftCacheConfig(nb::cast(state[0]), nb::cast(state[1]), + nb::cast(state[2]), nb::cast(state[3]), nb::cast(state[4]), + nb::cast(state[5]), nb::cast(state[6]), nb::cast(state[7]), + nb::cast(state[8]), nb::cast>(state[9]), + nb::cast>(state[10])); + }; + auto peftCacheConfigGetstate = [](tle::PeftCacheConfig const& self) + { + return nb::make_tuple(self.getNumHostModuleLayer(), self.getNumDeviceModuleLayer(), + self.getOptimalAdapterSize(), self.getMaxAdapterSize(), self.getNumPutWorkers(), self.getNumEnsureWorkers(), + self.getNumCopyStreams(), self.getMaxPagesPerBlockHost(), self.getMaxPagesPerBlockDevice(), + self.getDeviceCachePercent(), self.getHostCacheSize()); + }; + nb::class_(m, "PeftCacheConfig") + .def(nb::init const&, std::optional const&, + std::optional const&>(), + nb::arg("num_host_module_layer") = 0, nb::arg("num_device_module_layer") = 0, + nb::arg("optimal_adapter_size") = 8, nb::arg("max_adapter_size") = 64, nb::arg("num_put_workers") = 1, + nb::arg("num_ensure_workers") = 1, nb::arg("num_copy_streams") = 1, + nb::arg("max_pages_per_block_host") = 24, nb::arg("max_pages_per_block_device") = 8, + nb::arg("device_cache_percent") = nb::none(), nb::arg("host_cache_size") = nb::none(), + nb::arg("lora_prefetch_dir") = nb::none()) + .def_prop_ro("num_host_module_layer", &tle::PeftCacheConfig::getNumHostModuleLayer) + .def_prop_ro("num_device_module_layer", &tle::PeftCacheConfig::getNumDeviceModuleLayer) + .def_prop_ro("optimal_adapter_size", &tle::PeftCacheConfig::getOptimalAdapterSize) + .def_prop_ro("max_adapter_size", &tle::PeftCacheConfig::getMaxAdapterSize) + .def_prop_ro("num_put_workers", &tle::PeftCacheConfig::getNumPutWorkers) + .def_prop_ro("num_ensure_workers", &tle::PeftCacheConfig::getNumEnsureWorkers) + .def_prop_ro("num_copy_streams", &tle::PeftCacheConfig::getNumCopyStreams) + .def_prop_ro("max_pages_per_block_host", &tle::PeftCacheConfig::getMaxPagesPerBlockHost) + .def_prop_ro("max_pages_per_block_device", &tle::PeftCacheConfig::getMaxPagesPerBlockDevice) + .def_prop_ro("device_cache_percent", &tle::PeftCacheConfig::getDeviceCachePercent) + .def_prop_ro("host_cache_size", &tle::PeftCacheConfig::getHostCacheSize) + .def_prop_ro("lora_prefetch_dir", &tle::PeftCacheConfig::getLoraPrefetchDir) + .def("__getstate__", peftCacheConfigGetstate) + .def("__setstate__", peftCacheConfigSetstate); + + auto decodingConfigGetstate = [](tle::DecodingConfig const& self) + { + return nb::make_tuple( + self.getDecodingMode(), self.getLookaheadDecodingConfig(), self.getMedusaChoices(), self.getEagleConfig()); + }; + auto decodingConfigSetstate = [](tle::DecodingConfig& self, nb::tuple const& state) + { + if (state.size() != 4) + { + throw std::runtime_error("Invalid state!"); + } + new (&self) tle::DecodingConfig(nb::cast>(state[0]), // DecodingMode + nb::cast>(state[1]), // LookaheadDecodingConfig + nb::cast>(state[2]), // MedusaChoices + nb::cast>(state[3]) // EagleConfig + ); + }; + nb::class_(m, "DecodingConfig") + .def(nb::init, std::optional, + std::optional, std::optional>(), + nb::arg("decoding_mode") = nb::none(), nb::arg("lookahead_decoding_config") = nb::none(), + nb::arg("medusa_choices") = nb::none(), nb::arg("eagle_config") = nb::none()) + .def_prop_rw("decoding_mode", &tle::DecodingConfig::getDecodingMode, &tle::DecodingConfig::setDecodingMode) + .def_prop_rw("lookahead_decoding_config", &tle::DecodingConfig::getLookaheadDecodingConfig, + &tle::DecodingConfig::setLookaheadDecodingConfig) + .def_prop_rw("medusa_choices", &tle::DecodingConfig::getMedusaChoices, &tle::DecodingConfig::setMedusaChoices) + .def_prop_rw("eagle_config", &tle::DecodingConfig::getEagleConfig, &tle::DecodingConfig::setEagleConfig) + .def("__getstate__", decodingConfigGetstate) + .def("__setstate__", decodingConfigSetstate); + + auto debugConfigGetstate = [](tle::DebugConfig const& self) + { + return nb::make_tuple(self.getDebugInputTensors(), self.getDebugOutputTensors(), self.getDebugTensorNames(), + self.getDebugTensorsMaxIterations()); + }; + auto debugConfigSetstate = [](tle::DebugConfig& self, nb::tuple const& state) + { + if (state.size() != 4) + { + throw std::runtime_error("Invalid state!"); + } + new (&self) tle::DebugConfig(nb::cast(state[0]), nb::cast(state[1]), + nb::cast>(state[2]), nb::cast(state[3])); + }; + nb::class_(m, "DebugConfig") + .def(nb::init, SizeType32>(), nb::arg("debug_input_tensors") = false, + nb::arg("debug_output_tensors") = false, nb::arg("debug_tensor_names") = nb::none(), + nb::arg("debug_tensors_max_iterations") = false) + .def_prop_rw( + "debug_input_tensors", &tle::DebugConfig::getDebugInputTensors, &tle::DebugConfig::setDebugInputTensors) + .def_prop_rw( + "debug_output_tensors", &tle::DebugConfig::getDebugOutputTensors, &tle::DebugConfig::setDebugOutputTensors) + .def_prop_rw( + "debug_tensor_names", &tle::DebugConfig::getDebugTensorNames, &tle::DebugConfig::setDebugTensorNames) + .def_prop_rw("debug_tensors_max_iterations", &tle::DebugConfig::getDebugTensorsMaxIterations, + &tle::DebugConfig::setDebugTensorsMaxIterations) + .def("__getstate__", debugConfigGetstate) + .def("__setstate__", debugConfigSetstate); + + auto logitsPostProcessorConfigGetstate = [](tle::LogitsPostProcessorConfig const& self) + { return nb::make_tuple(self.getProcessorMap(), self.getProcessorBatched(), self.getReplicate()); }; + + auto logitsPostProcessorConfigSetstate = [](tle::LogitsPostProcessorConfig& self, nb::tuple const& state) + { + if (state.size() != 3) + { + throw std::runtime_error("Invalid LogitsPostProcessorConfig state!"); + } + new (&self) tle::LogitsPostProcessorConfig(nb::cast>(state[0]), + nb::cast>(state[1]), nb::cast(state[2])); + }; + + nb::class_(m, "LogitsPostProcessorConfig") + .def(nb::init, std::optional, + bool>(), + nb::arg("processor_map") = nb::none(), nb::arg("processor_batched") = nb::none(), + nb::arg("replicate") = true) + .def_prop_rw("processor_map", &tle::LogitsPostProcessorConfig::getProcessorMap, + &tle::LogitsPostProcessorConfig::setProcessorMap) + .def_prop_rw("processor_batched", &tle::LogitsPostProcessorConfig::getProcessorBatched, + &tle::LogitsPostProcessorConfig::setProcessorBatched) + .def_prop_rw( + "replicate", &tle::LogitsPostProcessorConfig::getReplicate, &tle::LogitsPostProcessorConfig::setReplicate) + .def("__getstate__", logitsPostProcessorConfigGetstate) + .def("__setstate__", logitsPostProcessorConfigSetstate); + + auto extendedRuntimePerfKnobConfigSetstate = [](tle::ExtendedRuntimePerfKnobConfig& self, nb::tuple const& state) + { + if (state.size() != 4) + { + throw std::runtime_error("Invalid extendedRuntimePerfKnobConfig state!"); + } + new (&self) tle::ExtendedRuntimePerfKnobConfig(nb::cast(state[0]), nb::cast(state[1]), + nb::cast(state[2]), nb::cast(state[3])); + }; + auto extendedRuntimePerfKnobConfigGetstate = [](tle::ExtendedRuntimePerfKnobConfig const& self) + { + return nb::make_tuple(self.getMultiBlockMode(), self.getEnableContextFMHAFP32Acc(), self.getCudaGraphMode(), + self.getCudaGraphCacheSize()); + }; + nb::class_(m, "ExtendedRuntimePerfKnobConfig") + .def( + nb::init(), nb::arg("multi_block_mode") = true, nb::arg("enable_context_fmha_fp32_acc") = false) + .def_prop_rw("multi_block_mode", &tle::ExtendedRuntimePerfKnobConfig::getMultiBlockMode, + &tle::ExtendedRuntimePerfKnobConfig::setMultiBlockMode) + .def_prop_rw("enable_context_fmha_fp32_acc", &tle::ExtendedRuntimePerfKnobConfig::getEnableContextFMHAFP32Acc, + &tle::ExtendedRuntimePerfKnobConfig::setEnableContextFMHAFP32Acc) + .def_prop_rw("cuda_graph_mode", &tle::ExtendedRuntimePerfKnobConfig::getCudaGraphMode, + &tle::ExtendedRuntimePerfKnobConfig::setCudaGraphMode) + .def_prop_rw("cuda_graph_cache_size", &tle::ExtendedRuntimePerfKnobConfig::getCudaGraphCacheSize, + &tle::ExtendedRuntimePerfKnobConfig::setCudaGraphCacheSize) + .def("__getstate__", extendedRuntimePerfKnobConfigGetstate) + .def("__setstate__", extendedRuntimePerfKnobConfigSetstate); + + auto SpeculativeDecodingConfigGetState + = [](tle::SpeculativeDecodingConfig const& self) { return nb::make_tuple(self.fastLogits); }; + auto SpeculativeDecodingConfigSetState = [](tle::SpeculativeDecodingConfig& self, nb::tuple const& state) + { + if (state.size() != 1) + { + throw std::runtime_error("Invalid SpeculativeDecodingConfig state!"); + } + new (&self) tle::SpeculativeDecodingConfig(nb::cast(state[0])); + }; + nb::class_(m, "SpeculativeDecodingConfig") + .def(nb::init(), nb::arg("fast_logits") = false) + .def_rw("fast_logits", &tle::SpeculativeDecodingConfig::fastLogits) + .def("__getstate__", SpeculativeDecodingConfigGetState) + .def("__setstate__", SpeculativeDecodingConfigSetState); + + // Guided decoding config + auto pyGuidedDecodingConfig = nb::class_(m, "GuidedDecodingConfig"); + + nb::enum_(pyGuidedDecodingConfig, "GuidedDecodingBackend") + .value("XGRAMMAR", tle::GuidedDecodingConfig::GuidedDecodingBackend::kXGRAMMAR) + .value("LLGUIDANCE", tle::GuidedDecodingConfig::GuidedDecodingBackend::kLLGUIDANCE); + + auto guidedDecodingConfigGetstate = [](tle::GuidedDecodingConfig const& self) { + return nb::make_tuple( + self.getBackend(), self.getEncodedVocab(), self.getTokenizerStr(), self.getStopTokenIds()); + }; + auto guidedDecodingConfigSetstate = [](tle::GuidedDecodingConfig& self, nb::tuple state) + { + if (state.size() != 4) + { + throw std::runtime_error("Invalid GuidedDecodingConfig state!"); + } + new (&self) tle::GuidedDecodingConfig(nb::cast(state[0]), + nb::cast>>(state[1]), nb::cast>(state[2]), + nb::cast>>(state[3])); + }; + + pyGuidedDecodingConfig + .def(nb::init>, + std::optional, std::optional>>(), + nb::arg("backend"), nb::arg("encoded_vocab") = nb::none(), nb::arg("tokenizer_str") = nb::none(), + nb::arg("stop_token_ids") = nb::none()) + .def_prop_rw("backend", &tle::GuidedDecodingConfig::getBackend, &tle::GuidedDecodingConfig::setBackend) + .def_prop_rw( + "encoded_vocab", &tle::GuidedDecodingConfig::getEncodedVocab, &tle::GuidedDecodingConfig::setEncodedVocab) + .def_prop_rw( + "tokenizer_str", &tle::GuidedDecodingConfig::getTokenizerStr, &tle::GuidedDecodingConfig::setTokenizerStr) + .def_prop_rw( + "stop_token_ids", &tle::GuidedDecodingConfig::getStopTokenIds, &tle::GuidedDecodingConfig::setStopTokenIds) + .def("__getstate__", guidedDecodingConfigGetstate) + .def("__setstate__", guidedDecodingConfigSetstate); + + auto cacheTransceiverConfigGetstate = [](tle::CacheTransceiverConfig const& self) + { return nb::make_tuple(self.getBackendType(), self.getMaxTokensInBuffer()); }; + auto cacheTransceiverConfigSetstate = [](tle::CacheTransceiverConfig& self, nb::tuple const& state) + { + if (state.size() != 2) + { + throw std::runtime_error("Invalid CacheTransceiverConfig state!"); + } + new (&self) tle::CacheTransceiverConfig( + nb::cast(state[0]), nb::cast>(state[1])); + }; + + nb::enum_(m, "CacheTransceiverBackendType") + .value("DEFAULT", tle::CacheTransceiverConfig::BackendType::DEFAULT) + .value("MPI", tle::CacheTransceiverConfig::BackendType::MPI) + .value("UCX", tle::CacheTransceiverConfig::BackendType::UCX) + .value("NIXL", tle::CacheTransceiverConfig::BackendType::NIXL) + .def("from_string", + [](std::string const& str) + { + if (str == "DEFAULT" || str == "default") + return tle::CacheTransceiverConfig::BackendType::DEFAULT; + if (str == "MPI" || str == "mpi") + return tle::CacheTransceiverConfig::BackendType::MPI; + if (str == "UCX" || str == "ucx") + return tle::CacheTransceiverConfig::BackendType::UCX; + if (str == "NIXL" || str == "nixl") + return tle::CacheTransceiverConfig::BackendType::NIXL; + throw std::runtime_error("Invalid backend type: " + str); + }); + + nb::class_(m, "CacheTransceiverConfig") + .def(nb::init, std::optional>(), + nb::arg("backend") = std::nullopt, nb::arg("max_tokens_in_buffer") = std::nullopt) + .def_prop_rw( + "backend", &tle::CacheTransceiverConfig::getBackendType, &tle::CacheTransceiverConfig::setBackendType) + .def_prop_rw("max_tokens_in_buffer", &tle::CacheTransceiverConfig::getMaxTokensInBuffer, + &tle::CacheTransceiverConfig::setMaxTokensInBuffer) + .def("__getstate__", cacheTransceiverConfigGetstate) + .def("__setstate__", cacheTransceiverConfigSetstate); + + auto executorConfigGetState = [](nb::object const& self) + { + auto& c = nb::cast(self); + // Return a tuple containing C++ data and the Python __dict__ + auto cpp_states = nb::make_tuple(c.getMaxBeamWidth(), c.getSchedulerConfig(), c.getKvCacheConfig(), + c.getEnableChunkedContext(), c.getNormalizeLogProbs(), c.getIterStatsMaxIterations(), + c.getRequestStatsMaxIterations(), c.getBatchingType(), c.getMaxBatchSize(), c.getMaxNumTokens(), + c.getParallelConfig(), c.getPeftCacheConfig(), c.getLogitsPostProcessorConfig(), c.getDecodingConfig(), + c.getUseGpuDirectStorage(), c.getGpuWeightsPercent(), c.getMaxQueueSize(), + c.getExtendedRuntimePerfKnobConfig(), c.getDebugConfig(), c.getRecvPollPeriodMs(), + c.getMaxSeqIdleMicroseconds(), c.getSpecDecConfig(), c.getGuidedDecodingConfig(), + c.getAdditionalModelOutputs(), c.getCacheTransceiverConfig(), c.getGatherGenerationLogits(), + c.getPromptTableOffloading(), c.getEnableTrtOverlap()); + auto pickle_tuple = nb::make_tuple(cpp_states, nb::getattr(self, "__dict__")); + return pickle_tuple; + }; + + auto executorConfigSetState = [](nb::object self, nb::tuple const& state) + { + if (state.size() != 2) + { + throw std::runtime_error("Invalid state!"); + } + + auto cpp_states = nb::cast(state[0]); + if (cpp_states.size() != 28) + { + throw std::runtime_error("Invalid cpp_states!"); + } + + // Restore C++ data + tle::ExecutorConfig* cpp_self = nb::inst_ptr(self); + new (cpp_self) tle::ExecutorConfig( // + nb::cast(cpp_states[0]), // MaxBeamWidth + nb::cast(cpp_states[1]), // SchedulerConfig + nb::cast(cpp_states[2]), // KvCacheConfig + nb::cast(cpp_states[3]), // EnableChunkedContext + nb::cast(cpp_states[4]), // NormalizeLogProbs + nb::cast(cpp_states[5]), // IterStatsMaxIterations + nb::cast(cpp_states[6]), // RequestStatsMaxIterations + nb::cast(cpp_states[7]), // BatchingType + nb::cast>(cpp_states[8]), // MaxBatchSize + nb::cast>(cpp_states[9]), // MaxNumTokens + nb::cast>(cpp_states[10]), // ParallelConfig + nb::cast>(cpp_states[11]), // PeftCacheConfig + nb::cast>(cpp_states[12]), // LogitsPostProcessorConfig + nb::cast>(cpp_states[13]), // DecodingConfig + nb::cast(cpp_states[14]), // UseGpuDirectStorage + nb::cast(cpp_states[15]), // GpuWeightsPercent + nb::cast>(cpp_states[16]), // MaxQueueSize + nb::cast(cpp_states[17]), // ExtendedRuntimePerfKnobConfig + nb::cast>(cpp_states[18]), // DebugConfig + nb::cast(cpp_states[19]), // RecvPollPeriodMs + nb::cast(cpp_states[20]), // MaxSeqIdleMicroseconds + nb::cast>(cpp_states[21]), // SpecDecConfig + nb::cast>(cpp_states[22]), // GuidedDecodingConfig + nb::cast>>(cpp_states[23]), // AdditionalModelOutputs + nb::cast>(cpp_states[24]), // CacheTransceiverConfig + nb::cast(cpp_states[25]), // GatherGenerationLogits + nb::cast(cpp_states[26]), // PromptTableOffloading + nb::cast(cpp_states[27]) // EnableTrtOverlap + ); + + // Restore Python data + auto py_state = nb::cast(state[1]); + self.attr("__dict__").attr("update")(py_state); + + nb::inst_mark_ready(self); + }; + + nb::class_(m, "ExecutorConfig", nb::dynamic_attr()) + .def(nb::init< // + SizeType32, // MaxBeamWidth + tle::SchedulerConfig const&, // SchedulerConfig + tle::KvCacheConfig const&, // KvCacheConfig + bool, // EnableChunkedContext + bool, // NormalizeLogProbs + SizeType32, // IterStatsMaxIterations + SizeType32, // RequestStatsMaxIterations + tle::BatchingType, // BatchingType + std::optional, // MaxBatchSize + std::optional, // MaxNumTokens + std::optional, // ParallelConfig + tle::PeftCacheConfig const&, // PeftCacheConfig + std::optional, // LogitsPostProcessorConfig + std::optional, // DecodingConfig + bool, // UseGpuDirectStorage + float, // GpuWeightsPercent + std::optional, // MaxQueueSize + tle::ExtendedRuntimePerfKnobConfig const&, // ExtendedRuntimePerfKnobConfig + std::optional, // DebugConfig + SizeType32, // RecvPollPeriodMs + uint64_t, // MaxSeqIdleMicroseconds + std::optional, // SpecDecConfig + std::optional, // GuidedDecodingConfig + std::optional>, // AdditionalModelOutputs + std::optional, // CacheTransceiverConfig + bool, // GatherGenerationLogits + bool, // PromptTableOffloading + bool // EnableTrtOverlap + >(), + nb::arg("max_beam_width") = 1, nb::arg("scheduler_config") = tle::SchedulerConfig(), + nb::arg("kv_cache_config") = tle::KvCacheConfig(), nb::arg("enable_chunked_context") = false, + nb::arg("normalize_log_probs") = true, + nb::arg("iter_stats_max_iterations") = tle::ExecutorConfig::kDefaultIterStatsMaxIterations, + nb::arg("request_stats_max_iterations") = tle::ExecutorConfig::kDefaultRequestStatsMaxIterations, + nb::arg("batching_type") = tle::BatchingType::kINFLIGHT, nb::arg("max_batch_size") = nb::none(), + nb::arg("max_num_tokens") = nb::none(), nb::arg("parallel_config") = nb::none(), + nb::arg("peft_cache_config") = tle::PeftCacheConfig(), nb::arg("logits_post_processor_config") = nb::none(), + nb::arg("decoding_config") = nb::none(), nb::arg("use_gpu_direct_storage") = false, + nb::arg("gpu_weights_percent") = 1.0, nb::arg("max_queue_size") = nb::none(), + nb::arg("extended_runtime_perf_knob_config") = tle::ExtendedRuntimePerfKnobConfig(), + nb::arg("debug_config") = nb::none(), nb::arg("recv_poll_period_ms") = 0, + nb::arg("max_seq_idle_microseconds") = tle::ExecutorConfig::kDefaultMaxSeqIdleMicroseconds, + nb::arg("spec_dec_config") = nb::none(), nb::arg("guided_decoding_config") = nb::none(), + nb::arg("additional_model_outputs") = nb::none(), nb::arg("cache_transceiver_config") = nb::none(), + nb::arg("gather_generation_logits") = false, nb::arg("mm_embedding_offloading") = false, + nb::arg("enable_trt_overlap") = false) + .def_prop_rw("max_beam_width", &tle::ExecutorConfig::getMaxBeamWidth, &tle::ExecutorConfig::setMaxBeamWidth) + .def_prop_rw("max_batch_size", &tle::ExecutorConfig::getMaxBatchSize, &tle::ExecutorConfig::setMaxBatchSize) + .def_prop_rw("max_num_tokens", &tle::ExecutorConfig::getMaxNumTokens, &tle::ExecutorConfig::setMaxNumTokens) + .def_prop_rw( + "scheduler_config", &tle::ExecutorConfig::getSchedulerConfigRef, &tle::ExecutorConfig::setSchedulerConfig) + .def_prop_rw( + "kv_cache_config", &tle::ExecutorConfig::getKvCacheConfigRef, &tle::ExecutorConfig::setKvCacheConfig) + .def_prop_rw("enable_chunked_context", &tle::ExecutorConfig::getEnableChunkedContext, + &tle::ExecutorConfig::setEnableChunkedContext) + .def_prop_rw("normalize_log_probs", &tle::ExecutorConfig::getNormalizeLogProbs, + &tle::ExecutorConfig::setNormalizeLogProbs) + .def_prop_rw("iter_stats_max_iterations", &tle::ExecutorConfig::getIterStatsMaxIterations, + &tle::ExecutorConfig::setIterStatsMaxIterations) + .def_prop_rw("request_stats_max_iterations", &tle::ExecutorConfig::getRequestStatsMaxIterations, + &tle::ExecutorConfig::setRequestStatsMaxIterations) + .def_prop_rw("batching_type", &tle::ExecutorConfig::getBatchingType, &tle::ExecutorConfig::setBatchingType) + .def_prop_rw( + "parallel_config", &tle::ExecutorConfig::getParallelConfig, &tle::ExecutorConfig::setParallelConfig) + .def_prop_rw( + "peft_cache_config", &tle::ExecutorConfig::getPeftCacheConfig, &tle::ExecutorConfig::setPeftCacheConfig) + .def_prop_rw("logits_post_processor_config", &tle::ExecutorConfig::getLogitsPostProcessorConfig, + &tle::ExecutorConfig::setLogitsPostProcessorConfig) + .def_prop_rw( + "decoding_config", &tle::ExecutorConfig::getDecodingConfig, &tle::ExecutorConfig::setDecodingConfig) + .def_prop_rw("use_gpu_direct_storage", &tle::ExecutorConfig::getUseGpuDirectStorage, + &tle::ExecutorConfig::setUseGpuDirectStorage) + .def_prop_rw("gpu_weights_percent", &tle::ExecutorConfig::getGpuWeightsPercent, + &tle::ExecutorConfig::setGpuWeightsPercent) + .def_prop_rw("max_queue_size", &tle::ExecutorConfig::getMaxQueueSize, &tle::ExecutorConfig::setMaxQueueSize) + .def_prop_rw("extended_runtime_perf_knob_config", &tle::ExecutorConfig::getExtendedRuntimePerfKnobConfig, + &tle::ExecutorConfig::setExtendedRuntimePerfKnobConfig) + .def_prop_rw("debug_config", &tle::ExecutorConfig::getDebugConfig, &tle::ExecutorConfig::setDebugConfig) + .def_prop_rw( + "recv_poll_period_ms", &tle::ExecutorConfig::getRecvPollPeriodMs, &tle::ExecutorConfig::setRecvPollPeriodMs) + .def_prop_rw("max_seq_idle_microseconds", &tle::ExecutorConfig::getMaxSeqIdleMicroseconds, + &tle::ExecutorConfig::setMaxSeqIdleMicroseconds) + .def_prop_rw("spec_dec_config", &tle::ExecutorConfig::getSpecDecConfig, &tle::ExecutorConfig::setSpecDecConfig) + .def_prop_rw("guided_decoding_config", &tle::ExecutorConfig::getGuidedDecodingConfig, + &tle::ExecutorConfig::setGuidedDecodingConfig) + .def_prop_rw("additional_model_outputs", &tle::ExecutorConfig::getAdditionalModelOutputs, + &tle::ExecutorConfig::setAdditionalModelOutputs) + .def_prop_rw("cache_transceiver_config", &tle::ExecutorConfig::getCacheTransceiverConfig, + &tle::ExecutorConfig::setCacheTransceiverConfig) + .def_prop_rw("gather_generation_logits", &tle::ExecutorConfig::getGatherGenerationLogits, + &tle::ExecutorConfig::setGatherGenerationLogits) + .def_prop_rw("mm_embedding_offloading", &tle::ExecutorConfig::getPromptTableOffloading, + &tle::ExecutorConfig::setPromptTableOffloading) + .def_prop_rw( + "enable_trt_overlap", &tle::ExecutorConfig::getEnableTrtOverlap, &tle::ExecutorConfig::setEnableTrtOverlap) + .def("__getstate__", executorConfigGetState) + .def("__setstate__", executorConfigSetState); +} + +} // namespace tensorrt_llm::nanobind::executor diff --git a/cpp/tensorrt_llm/nanobind/executor/executorConfig.h b/cpp/tensorrt_llm/nanobind/executor/executorConfig.h new file mode 100644 index 000000000000..5b63e7c5a3e3 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/executor/executorConfig.h @@ -0,0 +1,30 @@ +/* + * 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::executor +{ + +// Register bindings for executor API. +void initConfigBindings(nb::module_& m); + +} // namespace tensorrt_llm::nanobind::executor diff --git a/cpp/tensorrt_llm/nanobind/executor/request.cpp b/cpp/tensorrt_llm/nanobind/executor/request.cpp new file mode 100644 index 000000000000..80b9b52bd9d4 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/executor/request.cpp @@ -0,0 +1,955 @@ +/* + * 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 "request.h" +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/executor/executor.h" +#include "tensorrt_llm/executor/serializeUtils.h" +#include "tensorrt_llm/executor/tensor.h" +#include "tensorrt_llm/executor/types.h" +#include "tensorrt_llm/nanobind/common/customCasters.h" +#include "tensorrt_llm/runtime/cudaStream.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace nb = nanobind; +namespace tle = tensorrt_llm::executor; +using Tensor = tle::Tensor; +using SizeType32 = tle::SizeType32; +using FloatType = tle::FloatType; +using VecTokens = tle::VecTokens; +using IdType = tle::IdType; +using VecTokenExtraIds = tle::VecTokenExtraIds; + +namespace tensorrt_llm::nanobind::executor +{ + +void initRequestBindings(nb::module_& m) +{ + nb::enum_(m, "RequestType") + .value("REQUEST_TYPE_CONTEXT_AND_GENERATION", tle::RequestType::REQUEST_TYPE_CONTEXT_AND_GENERATION) + .value("REQUEST_TYPE_CONTEXT_ONLY", tle::RequestType::REQUEST_TYPE_CONTEXT_ONLY) + .value("REQUEST_TYPE_GENERATION_ONLY", tle::RequestType::REQUEST_TYPE_GENERATION_ONLY); + + nb::enum_(m, "FinishReason") + .value("NOT_FINISHED", tle::FinishReason::kNOT_FINISHED) + .value("END_ID", tle::FinishReason::kEND_ID) + .value("STOP_WORDS", tle::FinishReason::kSTOP_WORDS) + .value("LENGTH", tle::FinishReason::kLENGTH) + .value("TIMED_OUT", tle::FinishReason::kTIMED_OUT) + .value("CANCELLED", tle::FinishReason::kCANCELLED); + + nb::enum_(m, "KvCacheTransferMode") + .value("DRAM", tle::KvCacheTransferMode::DRAM) + .value("GDS", tle::KvCacheTransferMode::GDS) + .value("POSIX_DEBUG_FALLBACK", tle::KvCacheTransferMode::POSIX_DEBUG_FALLBACK); + + auto samplingConfigGetstate = [](tle::SamplingConfig const& self) + { + return nb::make_tuple(self.getBeamWidth(), self.getTopK(), self.getTopP(), self.getTopPMin(), + self.getTopPResetIds(), self.getTopPDecay(), self.getSeed(), self.getTemperature(), self.getMinTokens(), + self.getBeamSearchDiversityRate(), self.getRepetitionPenalty(), self.getPresencePenalty(), + self.getFrequencyPenalty(), self.getLengthPenalty(), self.getEarlyStopping(), self.getNoRepeatNgramSize(), + self.getNumReturnSequences(), self.getMinP(), self.getBeamWidthArray()); + }; + auto samplingConfigSetstate = [](tle::SamplingConfig& samplingConfig, nb::tuple const& state) + { + if (state.size() != 19) + { + throw std::runtime_error("Invalid SamplingConfig state!"); + } + new (&samplingConfig) tle::SamplingConfig(nb::cast(state[0]), // BeamWidth + nb::cast>(state[1]), // TopK + nb::cast>(state[2]), // TopP + nb::cast>(state[3]), // TopPMin + nb::cast>(state[4]), // TopPResetIds + nb::cast>(state[5]), // TopPDecay + nb::cast>(state[6]), // Seed + nb::cast>(state[7]), // Temperature + nb::cast>(state[8]), // MinTokens + nb::cast>(state[9]), // BeamSearchDiversityRate + nb::cast>(state[10]), // RepetitionPenalty + nb::cast>(state[11]), // PresencePenalty + nb::cast>(state[12]), // FrequencyPenalty + nb::cast>(state[13]), // LengthPenalty + nb::cast>(state[14]), // EarlyStopping + nb::cast>(state[15]), // NoRepeatNgramSize + nb::cast>(state[16]), // NumReturnSequences + nb::cast>(state[17]), // MinP + nb::cast>>(state[18]) // BeamWidthArray + ); + }; + nb::class_(m, "SamplingConfig") + .def(nb::init const&, // beamWidth + std::optional const&, // topP + std::optional const&, // topPMin + std::optional const&, // topPResetIds + std::optional const&, // topPDecay + std::optional const&, // seed + std::optional const&, // temperature + std::optional const&, // minTokens + std::optional const&, // beamSearchDiversityRate + std::optional const&, // repetitionPenalty + std::optional const&, // presencePenalty + std::optional const&, // frequencyPenalty + std::optional const&, // lengthPenalty + std::optional const&, // earlyStopping + std::optional const&, // noRepeatNgramSize + std::optional const&, // numReturnSequences + std::optional const&, // minP + std::optional> const& // beamWidthArray + >(), + // clang-format off + nb::arg("beam_width") = 1, + nb::kw_only(), + nb::arg("top_k") = nb::none(), + nb::arg("top_p") = nb::none(), + nb::arg("top_p_min") = nb::none(), + nb::arg("top_p_reset_ids") = nb::none(), + nb::arg("top_p_decay") = nb::none(), + nb::arg("seed") = nb::none(), + nb::arg("temperature") = nb::none(), + nb::arg("min_tokens") = nb::none(), + nb::arg("beam_search_diversity_rate") = nb::none(), + nb::arg("repetition_penalty") = nb::none(), + nb::arg("presence_penalty") = nb::none(), + nb::arg("frequency_penalty") = nb::none(), + nb::arg("length_penalty") = nb::none(), + nb::arg("early_stopping") = nb::none(), + nb::arg("no_repeat_ngram_size") = nb::none(), + nb::arg("num_return_sequences") = nb::none(), + nb::arg("min_p") = nb::none(), + nb::arg("beam_width_array") = nb::none()) // clang-format on + .def_prop_rw("beam_width", &tle::SamplingConfig::getBeamWidth, &tle::SamplingConfig::setBeamWidth) + .def_prop_rw("top_k", &tle::SamplingConfig::getTopK, &tle::SamplingConfig::setTopK) + .def_prop_rw("top_p", &tle::SamplingConfig::getTopP, &tle::SamplingConfig::setTopP) + .def_prop_rw("top_p_min", &tle::SamplingConfig::getTopPMin, &tle::SamplingConfig::setTopPMin) + .def_prop_rw("top_p_reset_ids", &tle::SamplingConfig::getTopPResetIds, &tle::SamplingConfig::setTopPResetIds) + .def_prop_rw("top_p_decay", &tle::SamplingConfig::getTopPDecay, &tle::SamplingConfig::setTopPDecay) + .def_prop_rw("seed", &tle::SamplingConfig::getSeed, &tle::SamplingConfig::setSeed) + .def_prop_rw("temperature", &tle::SamplingConfig::getTemperature, &tle::SamplingConfig::setTemperature) + .def_prop_rw("min_tokens", &tle::SamplingConfig::getMinTokens, &tle::SamplingConfig::setMinTokens) + .def_prop_rw("beam_search_diversity_rate", &tle::SamplingConfig::getBeamSearchDiversityRate, + &tle::SamplingConfig::setBeamSearchDiversityRate) + .def_prop_rw("repetition_penalty", &tle::SamplingConfig::getRepetitionPenalty, + &tle::SamplingConfig::setRepetitionPenalty) + .def_prop_rw("presence_penalty", &tle::SamplingConfig::getPresencePenalty, + [](tle::SamplingConfig& self, std::optional v) { self.setPresencePenalty(v); }) + .def_prop_rw( + "frequency_penalty", &tle::SamplingConfig::getFrequencyPenalty, &tle::SamplingConfig::setFrequencyPenalty) + .def_prop_rw("length_penalty", &tle::SamplingConfig::getLengthPenalty, &tle::SamplingConfig::setLengthPenalty) + .def_prop_rw("early_stopping", &tle::SamplingConfig::getEarlyStopping, &tle::SamplingConfig::setEarlyStopping) + .def_prop_rw("no_repeat_ngram_size", &tle::SamplingConfig::getNoRepeatNgramSize, + &tle::SamplingConfig::setNoRepeatNgramSize) + .def_prop_rw("num_return_sequences", &tle::SamplingConfig::getNumReturnSequences, + &tle::SamplingConfig::setNumReturnSequences) + .def_prop_rw("min_p", &tle::SamplingConfig::getMinP, &tle::SamplingConfig::setMinP) + .def_prop_rw( + "beam_width_array", &tle::SamplingConfig::getBeamWidthArray, &tle::SamplingConfig::setBeamWidthArray) + .def("__getstate__", samplingConfigGetstate) + .def("__setstate__", samplingConfigSetstate); + + auto additionalModelOutputGetstate + = [](tle::AdditionalModelOutput const& self) { return nb::make_tuple(self.name, self.gatherContext); }; + auto additionalModelOutputSetstate = [](tle::AdditionalModelOutput& additionalModelOutput, nb::tuple const& state) + { + if (state.size() != 2) + { + throw std::runtime_error("Invalid AdditionalModelOutput state!"); + } + new (&additionalModelOutput) + tle::AdditionalModelOutput(nb::cast(state[0]), nb::cast(state[1])); + }; + nb::class_(m, "AdditionalModelOutput") + .def(nb::init(), nb::arg("name"), nb::arg("gather_context") = false) + .def_rw("name", &tle::AdditionalModelOutput::name) + .def_rw("gather_context", &tle::AdditionalModelOutput::gatherContext) + .def("__getstate__", additionalModelOutputGetstate) + .def("__setstate__", additionalModelOutputSetstate); + + auto outputConfigGetstate = [](tle::OutputConfig const& self) + { + return nb::make_tuple(self.returnLogProbs, self.returnContextLogits, self.returnGenerationLogits, + self.excludeInputFromOutput, self.returnEncoderOutput, self.returnPerfMetrics, self.additionalModelOutputs); + }; + auto outputConfigSetstate = [](tle::OutputConfig& outputConfig, nb::tuple const& state) + { + if (state.size() != 7) + { + throw std::runtime_error("Invalid OutputConfig state!"); + } + new (&outputConfig) tle::OutputConfig(nb::cast(state[0]), nb::cast(state[1]), + nb::cast(state[2]), nb::cast(state[3]), nb::cast(state[4]), nb::cast(state[5]), + nb::cast>>(state[6])); + }; + nb::class_(m, "OutputConfig") + .def( + "__init__", + [](tle::OutputConfig& self, std::optional return_log_probs, std::optional return_context_logits, + std::optional return_generation_logits, std::optional exclude_input_from_output, + std::optional return_encoder_output, std::optional return_perf_metrics, + std::optional> additional_model_outputs) + { + new (&self) tle::OutputConfig(return_log_probs.value_or(false), return_context_logits.value_or(false), + return_generation_logits.value_or(false), exclude_input_from_output.value_or(false), + return_encoder_output.value_or(false), return_perf_metrics.value_or(false), + additional_model_outputs); + }, + nb::arg("return_log_probs") = nb::none(), nb::arg("return_context_logits") = nb::none(), + nb::arg("return_generation_logits") = nb::none(), nb::arg("exclude_input_from_output") = nb::none(), + nb::arg("return_encoder_output") = nb::none(), nb::arg("return_perf_metrics") = nb::none(), + nb::arg("additional_model_outputs") = nb::none()) + .def_rw("return_log_probs", &tle::OutputConfig::returnLogProbs) + .def_rw("return_context_logits", &tle::OutputConfig::returnContextLogits) + .def_rw("return_generation_logits", &tle::OutputConfig::returnGenerationLogits) + .def_rw("exclude_input_from_output", &tle::OutputConfig::excludeInputFromOutput) + .def_rw("return_encoder_output", &tle::OutputConfig::returnEncoderOutput) + .def_rw("return_perf_metrics", &tle::OutputConfig::returnPerfMetrics) + .def_rw("additional_model_outputs", &tle::OutputConfig::additionalModelOutputs) + .def("__getstate__", outputConfigGetstate) + .def("__setstate__", outputConfigSetstate); + + auto externalDraftTokensConfigGetstate = [](tle::ExternalDraftTokensConfig const& self) + { return nb::make_tuple(self.getTokens(), self.getLogits(), self.getAcceptanceThreshold()); }; + auto externalDraftTokensConfigSetstate + = [](tle::ExternalDraftTokensConfig& externalDraftTokensConfig, nb::tuple const& state) + { + if (state.size() != 3) + { + throw std::runtime_error("Invalid ExternalDraftTokensConfig state!"); + } + new (&externalDraftTokensConfig) tle::ExternalDraftTokensConfig(nb::cast(state[0]), + nb::cast>(state[1]), nb::cast>(state[2])); + }; + nb::class_(m, "ExternalDraftTokensConfig") + .def(nb::init, std::optional const&, std::optional>(), + nb::arg("tokens"), nb::arg("logits") = nb::none(), nb::arg("acceptance_threshold") = nb::none(), + nb::arg("fast_logits") = nb::none()) + .def_prop_ro("tokens", &tle::ExternalDraftTokensConfig::getTokens) + .def_prop_ro("logits", &tle::ExternalDraftTokensConfig::getLogits) + .def_prop_ro("acceptance_threshold", &tle::ExternalDraftTokensConfig::getAcceptanceThreshold) + .def("__getstate__", externalDraftTokensConfigGetstate) + .def("__setstate__", externalDraftTokensConfigSetstate) + .def_prop_ro("fast_logits", &tle::ExternalDraftTokensConfig::getFastLogits); + + auto promptTuningConfigGetstate = [](tle::PromptTuningConfig const& self) + { return nb::make_tuple(self.getEmbeddingTable(), self.getInputTokenExtraIds()); }; + auto promptTuningConfigSetstate = [](tle::PromptTuningConfig& promptTuningConfig, nb::tuple const& state) + { + if (state.size() != 2) + { + throw std::runtime_error("Invalid PromptTuningConfig state!"); + } + new (&promptTuningConfig) + tle::PromptTuningConfig(nb::cast(state[0]), nb::cast>(state[1])); + }; + nb::class_(m, "PromptTuningConfig") + .def(nb::init>(), nb::arg("embedding_table"), + nb::arg("input_token_extra_ids") = nb::none()) + .def_prop_ro("embedding_table", &tle::PromptTuningConfig::getEmbeddingTable) + .def_prop_ro("input_token_extra_ids", &tle::PromptTuningConfig::getInputTokenExtraIds) + .def("__getstate__", promptTuningConfigGetstate) + .def("__setstate__", promptTuningConfigSetstate); + + auto loraConfigGetstate = [](tle::LoraConfig const& self) + { return nb::make_tuple(self.getTaskId(), self.getWeights(), self.getConfig()); }; + auto loraConfigSetstate = [](tle::LoraConfig& loraConfig, nb::tuple const& state) + { + if (state.size() != 3) + { + throw std::runtime_error("Invalid LoraConfig state!"); + } + new (&loraConfig) tle::LoraConfig(nb::cast(state[0]), nb::cast>(state[1]), + nb::cast>(state[2])); + }; + nb::class_(m, "LoraConfig") + .def(nb::init, std::optional>(), nb::arg("task_id"), + nb::arg("weights") = nb::none(), nb::arg("config") = nb::none()) + .def_prop_ro("task_id", &tle::LoraConfig::getTaskId) + .def_prop_ro("weights", &tle::LoraConfig::getWeights) + .def_prop_ro("config", &tle::LoraConfig::getConfig) + .def("__getstate__", loraConfigGetstate) + .def("__setstate__", loraConfigSetstate); + + auto multimodalInputGetstate = [](tle::MultimodalInput const& self) + { return nb::make_tuple(self.getMultimodalHashes(), self.getMultimodalPositions(), self.getMultimodalLengths()); }; + auto multimodalInputSetstate = [](tle::MultimodalInput& multimodalInput, nb::tuple const& state) + { + if (state.size() != 3) + { + throw std::runtime_error("Invalid MultimodalInput state!"); + } + new (&multimodalInput) tle::MultimodalInput(nb::cast>>(state[0]), + nb::cast>(state[1]), nb::cast>(state[2])); + }; + nb::class_(m, "MultimodalInput") + .def(nb::init>, std::vector, std::vector>(), + nb::arg("multimodal_hashes"), nb::arg("multimodal_positions"), nb::arg("multimodal_lengths")) + .def_prop_ro("multimodal_hashes", &tle::MultimodalInput::getMultimodalHashes) + .def_prop_ro("multimodal_positions", &tle::MultimodalInput::getMultimodalPositions) + .def_prop_ro("multimodal_lengths", &tle::MultimodalInput::getMultimodalLengths) + .def("__getstate__", multimodalInputGetstate) + .def("__setstate__", multimodalInputSetstate); + + auto MropeConfigGetstate = [](tle::MropeConfig const& self) + { return nb::make_tuple(self.getMRopeRotaryCosSin(), self.getMRopePositionDeltas()); }; + auto MropeConfigSetstate = [](tle::MropeConfig& mropeConfig, nb::tuple const& state) + { + if (state.size() != 2) + { + throw std::runtime_error("Invalid MropeConfig state!"); + } + new (&mropeConfig) tle::MropeConfig(nb::cast(state[0]), nb::cast(state[1])); + }; + nb::class_(m, "MropeConfig") + .def(nb::init(), nb::arg("mrope_rotary_cos_sin"), nb::arg("mrope_position_deltas")) + .def_prop_ro("mrope_rotary_cos_sin", &tle::MropeConfig::getMRopeRotaryCosSin) + .def_prop_ro("mrope_position_deltas", &tle::MropeConfig::getMRopePositionDeltas) + .def("__getstate__", MropeConfigGetstate) + .def("__setstate__", MropeConfigSetstate); + + auto lookaheadDecodingConfigGetstate = [](tle::LookaheadDecodingConfig const& self) + { return nb::make_tuple(self.getWindowSize(), self.getNgramSize(), self.getVerificationSetSize()); }; + auto lookaheadDecodingConfigSetstate + = [](tle::LookaheadDecodingConfig& lookaheadDecodingConfig, nb::tuple const& state) + { + if (state.size() != 3) + { + throw std::runtime_error("Invalid LookaheadDecodingConfig state!"); + } + new (&lookaheadDecodingConfig) tle::LookaheadDecodingConfig( + nb::cast(state[0]), nb::cast(state[1]), nb::cast(state[2])); + }; + nb::class_(m, "LookaheadDecodingConfig") + .def(nb::init(), nb::arg("max_window_size"), nb::arg("max_ngram_size"), + nb::arg("max_verification_set_size")) + .def_prop_ro("max_window_size", &tle::LookaheadDecodingConfig::getWindowSize) + .def_prop_ro("max_ngram_size", &tle::LookaheadDecodingConfig::getNgramSize) + .def_prop_ro("max_verification_set_size", &tle::LookaheadDecodingConfig::getVerificationSetSize) + .def("calculate_speculative_resource", &tle::LookaheadDecodingConfig::calculateSpeculativeResource) + .def_static( + "calculate_speculative_resource_tuple", &tle::LookaheadDecodingConfig::calculateSpeculativeResourceTuple) + .def("__getstate__", lookaheadDecodingConfigGetstate) + .def("__setstate__", lookaheadDecodingConfigSetstate) + .def_static("get_default_lookahead_decoding_window", + []() { return tle::LookaheadDecodingConfig::kDefaultLookaheadDecodingWindow; }) + .def_static("get_default_lookahead_decoding_ngram", + []() { return tle::LookaheadDecodingConfig::kDefaultLookaheadDecodingNgram; }) + .def_static("get_default_lookahead_decoding_verification_set", + []() { return tle::LookaheadDecodingConfig::kDefaultLookaheadDecodingVerificationSet; }); + + auto TokenRangeRetentionConfigGetstate = [](tle::KvCacheRetentionConfig::TokenRangeRetentionConfig const& self) + { return nb::make_tuple(self.tokenStart, self.tokenEnd, self.priority, self.durationMs); }; + auto TokenRangeRetentionConfigSetstate + = [](tle::KvCacheRetentionConfig::TokenRangeRetentionConfig& tokenRangeRetentionConfig, nb::tuple const& state) + { + if (state.size() != 4) + { + throw std::runtime_error("Invalid state!"); + } + new (&tokenRangeRetentionConfig) tle::KvCacheRetentionConfig::TokenRangeRetentionConfig( + nb::cast(state[0]), nb::cast>(state[1]), + nb::cast(state[2]), nb::cast>(state[3])); + }; + auto kvCacheRetentionConfigGetstate = [](tle::KvCacheRetentionConfig const& self) + { + return nb::make_tuple(self.getTokenRangeRetentionConfigs(), self.getDecodeRetentionPriority(), + self.getDecodeDurationMs(), self.getTransferMode(), self.getDirectory()); + }; + auto kvCacheRetentionConfigSetstate + = [](tle::KvCacheRetentionConfig& kvCacheRetentionConfig, nb::tuple const& state) + { + if (state.size() != 5) + { + throw std::runtime_error("Invalid state!"); + } + new (&kvCacheRetentionConfig) tle::KvCacheRetentionConfig( + nb::cast>(state[0]), + nb::cast(state[1]), nb::cast>(state[2]), + nb::cast(state[3]), nb::cast>(state[4])); + }; + + auto kvCacheRetentionConfig = nb::class_(m, "KvCacheRetentionConfig"); + + nb::class_( + kvCacheRetentionConfig, "TokenRangeRetentionConfig") + .def(nb::init, tle::RetentionPriority, + std::optional>(), + nb::arg("token_start"), nb::arg("token_end"), nb::arg("priority"), nb::arg("duration_ms") = nb::none()) + .def_rw("token_start", &tle::KvCacheRetentionConfig::TokenRangeRetentionConfig::tokenStart) + .def_rw("token_end", &tle::KvCacheRetentionConfig::TokenRangeRetentionConfig::tokenEnd) + .def_rw("priority", &tle::KvCacheRetentionConfig::TokenRangeRetentionConfig::priority) + .def_rw("duration_ms", &tle::KvCacheRetentionConfig::TokenRangeRetentionConfig::durationMs) + .def("__getstate__", TokenRangeRetentionConfigGetstate) + .def("__setstate__", TokenRangeRetentionConfigSetstate) + .def("__eq__", &tle::KvCacheRetentionConfig::TokenRangeRetentionConfig::operator==); + + // There's a circular dependency between the declaration of the TokenRangeRetentionPriority and + // KvCacheRetentionConfig bindings. Defer definition of the KvCacheRetentionConfig bindings until the + // TokenRangeRetentionPriority bindings have been defined. + kvCacheRetentionConfig + .def(nb::init, tle::RetentionPriority, + std::optional, tle::KvCacheTransferMode, std::optional>(), + nb::arg("token_range_retention_configs"), + nb::arg("decode_retention_priority") = tle::KvCacheRetentionConfig::kDefaultRetentionPriority, + nb::arg("decode_duration_ms") = nb::none(), nb::arg("transfer_mode") = tle::KvCacheTransferMode::DRAM, + nb::arg("directory") = nb::none()) + .def_prop_ro("token_range_retention_configs", &tle::KvCacheRetentionConfig::getTokenRangeRetentionConfigs) + .def_prop_ro("decode_retention_priority", &tle::KvCacheRetentionConfig::getDecodeRetentionPriority) + .def_prop_ro("decode_duration_ms", &tle::KvCacheRetentionConfig::getDecodeDurationMs) + .def_prop_ro("transfer_mode", &tle::KvCacheRetentionConfig::getTransferMode) + .def_prop_ro("directory", &tle::KvCacheRetentionConfig::getDirectory) + .def("__getstate__", kvCacheRetentionConfigGetstate) + .def("__setstate__", kvCacheRetentionConfigSetstate) + .def("__eq__", &tle::KvCacheRetentionConfig::operator==); + + auto ContextPhaseParamsGetState = [](tle::ContextPhaseParams const& self) + { + if (self.getState() != nullptr) + { + auto serializedState = self.getSerializedState(); + return nb::make_tuple(self.getFirstGenTokens(), self.getReqId(), + nb::bytes(serializedState.data(), serializedState.size()), self.getDraftTokens()); + } + return nb::make_tuple(self.getFirstGenTokens(), self.getReqId(), nb::none(), self.getDraftTokens()); + }; + + auto ContextPhaseParamsSetState = [](tle::ContextPhaseParams& contextPhaseParams, nb::tuple const& state) + { + if (state.size() != 4) + { + throw std::runtime_error("Invalid ContextPhaseParams state!"); + } + if (!state[2].is_none()) + { + auto opaque_state = nb::cast(state[2]); + auto opaque_state_str_view = std::string_view(opaque_state.c_str(), opaque_state.size()); + new (&contextPhaseParams) tle::ContextPhaseParams(nb::cast(state[0]), + nb::cast(state[1]), + std::vector(opaque_state_str_view.begin(), opaque_state_str_view.end()), + nb::cast>(state[3])); + } + else + { + new (&contextPhaseParams) tle::ContextPhaseParams(nb::cast(state[0]), + nb::cast(state[1]), + nb::cast>(state[3])); + } + }; + + nb::class_(m, "ContextPhaseParams") + .def( + "__init__", + [](tle::ContextPhaseParams& self, VecTokens const& first_gen_tokens, + tle::ContextPhaseParams::RequestIdType req_id, std::optional const& opaque_state, + std::optional const& draft_tokens) + { + if (opaque_state) + { + auto opaque_state_str_view + = std::string_view(opaque_state.value().c_str(), opaque_state.value().size()); + new (&self) tle::ContextPhaseParams(first_gen_tokens, req_id, + std::vector(opaque_state_str_view.begin(), opaque_state_str_view.end()), draft_tokens); + } + else + { + new (&self) tle::ContextPhaseParams(first_gen_tokens, req_id, draft_tokens); + } + }, + nb::arg("first_gen_tokens"), nb::arg("req_id"), nb::arg("opaque_state").none(), + nb::arg("draft_tokens").none()) + .def_prop_ro("first_gen_tokens", [](tle::ContextPhaseParams const& self) { return self.getFirstGenTokens(); }) + .def_prop_ro("draft_tokens", [](tle::ContextPhaseParams const& self) { return self.getDraftTokens(); }) + .def_prop_ro("req_id", &tle::ContextPhaseParams::getReqId) + .def_prop_ro("opaque_state", + [](tle::ContextPhaseParams const& self) + { + std::optional opaque_state{std::nullopt}; + if (self.getState() != nullptr) + { + auto serializedState = self.getSerializedState(); + opaque_state = nb::bytes(serializedState.data(), serializedState.size()); + } + return opaque_state; + }) + .def("__getstate__", ContextPhaseParamsGetState) + .def("__setstate__", ContextPhaseParamsSetState); + + auto EagleDecodingConfigGetstate = [](tle::EagleConfig const& self) + { + return nb::make_tuple(self.getEagleChoices(), self.isGreedySampling(), self.getPosteriorThreshold(), + self.useDynamicTree(), self.getDynamicTreeMaxTopK()); + }; + auto EagleDecodingConfigSetstate = [](tle::EagleConfig& self, nb::tuple const& state) + { + if (state.size() != 5) + { + throw std::runtime_error("Invalid EagleConfig state!"); + } + new (&self) tle::EagleConfig(nb::cast>(state[0]), nb::cast(state[1]), + nb::cast>(state[2]), nb::cast(state[3]), + nb::cast>(state[4])); + }; + nb::class_(m, "EagleConfig") + .def(nb::init, bool, std::optional, bool, std::optional>(), + nb::arg("eagle_choices") = nb::none(), nb::arg("greedy_sampling") = true, + nb::arg("posterior_threshold") = nb::none(), nb::arg("use_dynamic_tree") = false, + nb::arg("dynamic_tree_max_topK") = nb::none()) + .def_prop_ro("eagle_choices", &tle::EagleConfig::getEagleChoices) + .def_prop_ro("greedy_sampling", &tle::EagleConfig::isGreedySampling) + .def_prop_ro("posterior_threshold", &tle::EagleConfig::getPosteriorThreshold) + .def_prop_ro("use_dynamic_tree", &tle::EagleConfig::useDynamicTree) + .def_prop_ro("dynamic_tree_max_topK", &tle::EagleConfig::getDynamicTreeMaxTopK) + .def("__getstate__", EagleDecodingConfigGetstate) + .def("__setstate__", EagleDecodingConfigSetstate); + + // Guided decoding params + auto pyGuidedDecodingParams = nb::class_(m, "GuidedDecodingParams"); + + nb::enum_(pyGuidedDecodingParams, "GuideType") + .value("JSON", tle::GuidedDecodingParams::GuideType::kJSON) + .value("JSON_SCHEMA", tle::GuidedDecodingParams::GuideType::kJSON_SCHEMA) + .value("REGEX", tle::GuidedDecodingParams::GuideType::kREGEX) + .value("EBNF_GRAMMAR", tle::GuidedDecodingParams::GuideType::kEBNF_GRAMMAR) + .value("STRUCTURAL_TAG", tle::GuidedDecodingParams::GuideType::kSTRUCTURAL_TAG); + + auto guidedDecodingParamsGetstate + = [](tle::GuidedDecodingParams const& self) { return nb::make_tuple(self.getGuideType(), self.getGuide()); }; + + auto guidedDecodingParamsSetstate = [](tle::GuidedDecodingParams& self, nb::tuple const& state) + { + if (state.size() != 2) + { + throw std::runtime_error("Invalid GuidedDecodingParams state!"); + } + new (&self) tle::GuidedDecodingParams( + nb::cast(state[0]), nb::cast>(state[1])); + }; + + pyGuidedDecodingParams + .def(nb::init>(), nb::arg("guide_type"), + nb::arg("guide") = nb::none()) + .def_prop_ro("guide_type", &tle::GuidedDecodingParams::getGuideType) + .def_prop_ro("guide", &tle::GuidedDecodingParams::getGuide) + .def("__getstate__", guidedDecodingParamsGetstate) + .def("__setstate__", guidedDecodingParamsSetstate); + + auto requestGetstate = [](tle::Request const& self) + { + return nb::make_tuple(self.getInputTokenIds(), self.getMaxTokens(), self.getStreaming(), + self.getSamplingConfig(), self.getOutputConfig(), self.getEndId(), self.getPadId(), self.getPositionIds(), + self.getBadWords(), self.getStopWords(), self.getEmbeddingBias(), self.getExternalDraftTokensConfig(), + self.getPromptTuningConfig(), self.getMultimodalInput(), self.getMultimodalEmbedding(), + self.getMropeConfig(), self.getLoraConfig(), self.getLookaheadConfig(), self.getKvCacheRetentionConfig(), + self.getLogitsPostProcessorName(), self.getLogitsPostProcessor(), self.getEncoderInputTokenIds(), + self.getClientId(), self.getReturnAllGeneratedTokens(), self.getPriority(), self.getRequestType(), + self.getContextPhaseParams(), self.getEncoderInputFeatures(), self.getEncoderOutputLength(), + self.getCrossAttentionMask(), self.getEagleConfig(), self.getSkipCrossAttnBlocks(), + self.getGuidedDecodingParams()); + }; + auto requestSetstate = [](tle::Request& self, nb::tuple const& state) + { + if (state.size() != 33) + { + throw std::runtime_error("Invalid Request state!"); + } + new (&self) tle::Request(nb::cast(state[0]), nb::cast(state[1]), + nb::cast(state[2]), nb::cast(state[3]), nb::cast(state[4]), + nb::cast>(state[5]), nb::cast>(state[6]), + nb::cast>>(state[7]), + nb::cast>>(state[8]), + nb::cast>>(state[9]), nb::cast>(state[10]), + nb::cast>(state[11]), + nb::cast>(state[12]), + nb::cast>(state[13]), nb::cast>(state[14]), + nb::cast>(state[15]), nb::cast>(state[16]), + nb::cast>(state[17]), + nb::cast>(state[18]), + nb::cast>(state[19]), + nb::cast>(state[20]), nb::cast>(state[21]), + nb::cast>(state[22]), nb::cast(state[23]), + nb::cast(state[24]), nb::cast(state[25]), + nb::cast>(state[26]), + nb::cast>(state[27]), nb::cast>(state[28]), + nb::cast>(state[29]), 1, nb::cast>(state[30]), + nb::cast>(state[31]), + nb::cast>(state[32])); + }; + + nb::class_ request(m, "Request", nb::dynamic_attr()); + request + .def(nb::init const&, // endId + std::optional const&, // padId + std::optional>, // positionIds + std::optional>, // badWords + std::optional>, // stopWords + std::optional, // embeddingBias + std::optional, // externalDraftTokensConfig + std::optional, // pTuningConfig + std::optional, // multimodalInput + std::optional, // multimodalEmbedding + std::optional, // mRopeConfig + std::optional, // loraConfig + std::optional, // lookaheadConfig + std::optional, // kvCacheRetentionConfig + std::optional, // logitsPostProcessorName + std::optional, // logitsPostProcessor + std::optional, // encoderInputTokenIds + std::optional, // clientId + bool, // returnAllGeneratedTokens + tle::PriorityType, // priority + tle::RequestType, // type + std::optional, // contextPhaseParams + std::optional, // encoderInputFeatures + std::optional, // encoderOutputLength + std::optional, // crossAttentionMask + SizeType32, // numReturnSequences + std::optional, // eagleConfig + std::optional, // skipCrossAttnBlocks + std::optional, // guidedDecodingParams + std::optional, // languageAdapterUid + std::optional // allottedTimeMs + >(), + // clang-format off + nb::arg("input_token_ids"), + nb::arg("max_tokens"), + nb::kw_only(), + nb::arg("streaming") = false, + nb::arg("sampling_config") = tle::SamplingConfig(), + nb::arg("output_config") = tle::OutputConfig(), + nb::arg("end_id") = nb::none(), + nb::arg("pad_id") = nb::none(), + nb::arg("position_ids") = nb::none(), + nb::arg("bad_words") = nb::none(), + nb::arg("stop_words") = nb::none(), + nb::arg("embedding_bias") = nb::none(), + nb::arg("external_draft_tokens_config") = nb::none(), + nb::arg("prompt_tuning_config") = nb::none(), + nb::arg("multimodal_input") = nb::none(), + nb::arg("multimodal_embedding") = nb::none(), + nb::arg("mrope_config") = nb::none(), + nb::arg("lora_config") = nb::none(), + nb::arg("lookahead_config") = nb::none(), + nb::arg("kv_cache_retention_config") = nb::none(), + nb::arg("logits_post_processor_name") = nb::none(), + nb::arg("logits_post_processor") = nb::none(), + nb::arg("encoder_input_token_ids") = nb::none(), + nb::arg("client_id") = nb::none(), + nb::arg("return_all_generated_tokens") = false, + nb::arg("priority") = tle::Request::kDefaultPriority, + nb::arg("type") = tle::RequestType::REQUEST_TYPE_CONTEXT_AND_GENERATION, + nb::arg("context_phase_params") = nb::none(), + nb::arg("encoder_input_features") = nb::none(), + nb::arg("encoder_output_length") = nb::none(), + nb::arg("cross_attention_mask") = nb::none(), + nb::arg("num_return_sequences") = 1, + nb::arg("eagle_config") = nb::none(), + nb::arg("skip_cross_attn_blocks") = nb::none(), + nb::arg("guided_decoding_params") = nb::none(), + nb::arg("language_adapter_uid") = nb::none(), + nb::arg("allotted_time_ms") = nb::none() + ) // clang-format on + .def_prop_ro("input_token_ids", &tle::Request::getInputTokenIds) + .def_prop_ro("max_tokens", &tle::Request::getMaxTokens) + .def_prop_rw("streaming", &tle::Request::getStreaming, &tle::Request::setStreaming) + .def_prop_rw("sampling_config", &tle::Request::getSamplingConfig, &tle::Request::setSamplingConfig) + .def_prop_rw("output_config", &tle::Request::getOutputConfig, &tle::Request::setOutputConfig) + .def_prop_rw("end_id", &tle::Request::getEndId, &tle::Request::setEndId) + .def_prop_rw("pad_id", &tle::Request::getPadId, &tle::Request::setPadId) + .def_prop_rw("position_ids", &tle::Request::getPositionIds, &tle::Request::setPositionIds) + .def_prop_rw("bad_words", &tle::Request::getBadWords, &tle::Request::setBadWords) + .def_prop_rw("stop_words", &tle::Request::getStopWords, &tle::Request::setStopWords) + .def_prop_rw("embedding_bias", &tle::Request::getEmbeddingBias, &tle::Request::setEmbeddingBias) + .def_prop_rw("external_draft_tokens_config", &tle::Request::getExternalDraftTokensConfig, + &tle::Request::setExternalDraftTokensConfig) + .def_prop_rw("prompt_tuning_config", &tle::Request::getPromptTuningConfig, &tle::Request::setPromptTuningConfig) + .def_prop_rw("multimodal_input", &tle::Request::getMultimodalInput, &tle::Request::setMultimodalInput) + .def_prop_rw( + "multimodal_embedding", &tle::Request::getMultimodalEmbedding, &tle::Request::setMultimodalEmbedding) + .def_prop_rw("mrope_config", &tle::Request::getMropeConfig, &tle::Request::setMropeConfig) + .def_prop_rw("lora_config", &tle::Request::getLoraConfig, &tle::Request::setLoraConfig) + .def_prop_rw("lookahead_config", &tle::Request::getLookaheadConfig, &tle::Request::setLookaheadConfig) + .def_prop_rw("kv_cache_retention_config", &tle::Request::getKvCacheRetentionConfig, + &tle::Request::setKvCacheRetentionConfig) + .def_prop_rw("logits_post_processor_name", &tle::Request::getLogitsPostProcessorName, + &tle::Request::setLogitsPostProcessorName) + .def_prop_rw( + "logits_post_processor", &tle::Request::getLogitsPostProcessor, &tle::Request::setLogitsPostProcessor) + .def_prop_rw( + "encoder_input_token_ids", &tle::Request::getEncoderInputTokenIds, &tle::Request::setEncoderInputTokenIds) + .def_prop_rw("client_id", &tle::Request::getClientId, &tle::Request::setClientId) + .def_prop_rw("return_all_generated_tokens", &tle::Request::getReturnAllGeneratedTokens, + &tle::Request::setReturnAllGeneratedTokens) + .def_prop_rw("request_type", &tle::Request::getRequestType, &tle::Request::setRequestType) + .def_prop_rw( + "encoder_input_features", &tle::Request::getEncoderInputFeatures, &tle::Request::setEncoderInputFeatures) + .def_prop_rw("cross_attention_mask", &tle::Request::getCrossAttentionMask, &tle::Request::setCrossAttentionMask) + .def_prop_rw("eagle_config", &tle::Request::getEagleConfig, &tle::Request::setEagleConfig) + .def_prop_rw( + "skip_cross_attn_blocks", &tle::Request::getSkipCrossAttnBlocks, &tle::Request::setSkipCrossAttnBlocks) + .def_prop_rw( + "guided_decoding_params", &tle::Request::getGuidedDecodingParams, &tle::Request::setGuidedDecodingParams) + .def_prop_rw("allotted_time_ms", &tle::Request::getAllottedTimeMs, &tle::Request::setAllottedTimeMs) + .def_prop_rw("context_phase_params", &tle::Request::getContextPhaseParams, &tle::Request::setContextPhaseParams) + .def("__getstate__", requestGetstate) + .def("__setstate__", requestSetstate); + request.attr("BATCHED_POST_PROCESSOR_NAME") = tle::Request::kBatchedPostProcessorName; + + nb::class_(m, "SpeculativeDecodingFastLogitsInfo") + .def(nb::init<>()) + .def_rw("draft_request_id", &tle::SpeculativeDecodingFastLogitsInfo::draftRequestId) + .def_rw("draft_participant_id", &tle::SpeculativeDecodingFastLogitsInfo::draftParticipantId) + .def("to_tensor", &tle::SpeculativeDecodingFastLogitsInfo::toTensor); + + auto requestPerfMetrics = nb::class_(m, "RequestPerfMetrics"); + + auto timingMetricsGetstate = [](tle::RequestPerfMetrics::TimingMetrics const& self) + { + return nb::make_tuple(self.arrivalTime, self.firstScheduledTime, self.firstTokenTime, self.lastTokenTime, + self.kvCacheTransferStart, self.kvCacheTransferEnd, self.kvCacheSize); + }; + auto timingMetricsSetstate = [](tle::RequestPerfMetrics::TimingMetrics& timingMetrics, nb::tuple const& state) + { + if (state.size() != 7) + { + throw std::runtime_error("Invalid TimingMetrics state!"); + } + new (&timingMetrics) + tle::RequestPerfMetrics::TimingMetrics{nb::cast(state[0]), + nb::cast(state[1]), + nb::cast(state[2]), + nb::cast(state[3]), + nb::cast(state[4]), + nb::cast(state[5]), nb::cast(state[6])}; + }; + nb::class_(m, "TimingMetrics") + .def(nb::init<>()) + .def_rw("arrival_time", &tle::RequestPerfMetrics::TimingMetrics::arrivalTime) + .def_rw("first_scheduled_time", &tle::RequestPerfMetrics::TimingMetrics::firstScheduledTime) + .def_rw("first_token_time", &tle::RequestPerfMetrics::TimingMetrics::firstTokenTime) + .def_rw("last_token_time", &tle::RequestPerfMetrics::TimingMetrics::lastTokenTime) + .def_rw("kv_cache_transfer_start", &tle::RequestPerfMetrics::TimingMetrics::kvCacheTransferStart) + .def_rw("kv_cache_transfer_end", &tle::RequestPerfMetrics::TimingMetrics::kvCacheTransferEnd) + .def_rw("kv_cache_size", &tle::RequestPerfMetrics::TimingMetrics::kvCacheSize) + .def("__getstate__", timingMetricsGetstate) + .def("__setstate__", timingMetricsSetstate); + + auto kvCacheMetricsGetstate = [](tle::RequestPerfMetrics::KvCacheMetrics const& self) + { + return nb::make_tuple(self.numTotalAllocatedBlocks, self.numNewAllocatedBlocks, self.numReusedBlocks, + self.numMissedBlocks, self.kvCacheHitRate); + }; + auto kvCacheMetricsSetstate = [](tle::RequestPerfMetrics::KvCacheMetrics& kvCacheMetrics, nb::tuple const& state) + { + if (state.size() != 5) + { + throw std::runtime_error("Invalid KvCacheMetrics state!"); + } + new (&kvCacheMetrics) + tle::RequestPerfMetrics::KvCacheMetrics{nb::cast(state[0]), nb::cast(state[1]), + nb::cast(state[2]), nb::cast(state[3]), nb::cast(state[4])}; + }; + nb::class_(m, "KvCacheMetrics") + .def(nb::init<>()) + .def_rw("num_total_allocated_blocks", &tle::RequestPerfMetrics::KvCacheMetrics::numTotalAllocatedBlocks) + .def_rw("num_new_allocated_blocks", &tle::RequestPerfMetrics::KvCacheMetrics::numNewAllocatedBlocks) + .def_rw("num_reused_blocks", &tle::RequestPerfMetrics::KvCacheMetrics::numReusedBlocks) + .def_rw("num_missed_blocks", &tle::RequestPerfMetrics::KvCacheMetrics::numMissedBlocks) + .def_rw("kv_cache_hit_rate", &tle::RequestPerfMetrics::KvCacheMetrics::kvCacheHitRate) + .def("__getstate__", kvCacheMetricsGetstate) + .def("__setstate__", kvCacheMetricsSetstate); + + auto speculativeDecodingMetricsGetstate = [](tle::RequestPerfMetrics::SpeculativeDecodingMetrics const& self) + { return nb::make_tuple(self.acceptanceRate, self.totalAcceptedDraftTokens, self.totalDraftTokens); }; + auto speculativeDecodingMetricsSetstate + = [](tle::RequestPerfMetrics::SpeculativeDecodingMetrics& speculativeDecodingMetrics, nb::tuple const& state) + { + if (state.size() != 3) + { + throw std::runtime_error("Invalid SpeculativeDecodingMetrics state!"); + } + new (&speculativeDecodingMetrics) tle::RequestPerfMetrics::SpeculativeDecodingMetrics{ + nb::cast(state[0]), nb::cast(state[1]), nb::cast(state[2])}; + }; + + nb::class_(m, "SpeculativeDecodingMetrics") + .def(nb::init<>()) + .def_rw("acceptance_rate", &tle::RequestPerfMetrics::SpeculativeDecodingMetrics::acceptanceRate) + .def_rw("total_accepted_draft_tokens", + &tle::RequestPerfMetrics::SpeculativeDecodingMetrics::totalAcceptedDraftTokens) + .def_rw("total_draft_tokens", &tle::RequestPerfMetrics::SpeculativeDecodingMetrics::totalDraftTokens) + .def("__getstate__", speculativeDecodingMetricsGetstate) + .def("__setstate__", speculativeDecodingMetricsSetstate); + + auto requestPerfMetricsGetstate = [](tle::RequestPerfMetrics const& self) + { + return nb::make_tuple(self.timingMetrics, self.kvCacheMetrics, self.speculativeDecoding, self.firstIter, + self.lastIter, self.iter); + }; + auto requestPerfMetricsSetstate = [](tle::RequestPerfMetrics& self, nb::tuple const& state) + { + if (state.size() != 6) + { + throw std::runtime_error("Invalid RequestPerfMetrics state!"); + } + new (&self) tle::RequestPerfMetrics{nb::cast(state[0]), + nb::cast(state[1]), + nb::cast(state[2]), + nb::cast>(state[3]), + nb::cast>(state[4]), + nb::cast>(state[5])}; + }; + + // There's a circular dependency between the declaration of the TimingMetrics and RequestPerfMetrics bindings. + // Defer definition of the RequestPerfMetrics bindings until the TimingMetrics have been defined. + requestPerfMetrics.def(nb::init<>()) + .def_rw("timing_metrics", &tle::RequestPerfMetrics::timingMetrics) + .def_rw("kv_cache_metrics", &tle::RequestPerfMetrics::kvCacheMetrics) + .def_rw("speculative_decoding", &tle::RequestPerfMetrics::speculativeDecoding) + .def_rw("first_iter", &tle::RequestPerfMetrics::firstIter) + .def_rw("last_iter", &tle::RequestPerfMetrics::lastIter) + .def_rw("iter", &tle::RequestPerfMetrics::iter) + .def("__getstate__", requestPerfMetricsGetstate) + .def("__setstate__", requestPerfMetricsSetstate); + + nb::class_(m, "AdditionalOutput") + .def(nb::init(), nb::arg("name"), nb::arg("output")) + .def_rw("name", &tle::AdditionalOutput::name) + .def_rw("output", &tle::AdditionalOutput::output); + + auto resultSetstate = [](tle::Result& self, nb::tuple const& state) + { + if (state.size() != 13) + { + throw std::runtime_error("Invalid Request state!"); + } + tle::Result result; + result.isFinal = nb::cast(state[0]); + result.outputTokenIds = nb::cast>(state[1]); + result.cumLogProbs = nb::cast>>(state[2]); + result.logProbs = nb::cast>>>(state[3]); + result.contextLogits = nb::cast>(state[4]); + result.generationLogits = nb::cast>(state[5]); + result.encoderOutput = nb::cast>(state[6]); + result.finishReasons = nb::cast>(state[7]); + result.sequenceIndex = nb::cast(state[8]); + result.isSequenceFinal = nb::cast(state[9]); + result.decodingIter = nb::cast(state[10]); + result.contextPhaseParams = nb::cast>(state[11]); + result.requestPerfMetrics = nb::cast>(state[12]); + new (&self) tle::Result(result); + }; + + auto resultGetstate = [](tle::Result const& self) + { + return nb::make_tuple(self.isFinal, self.outputTokenIds, self.cumLogProbs, self.logProbs, self.contextLogits, + self.generationLogits, self.encoderOutput, self.finishReasons, self.sequenceIndex, self.isSequenceFinal, + self.decodingIter, self.contextPhaseParams, self.requestPerfMetrics); + }; + + nb::class_(m, "Result") + .def(nb::init<>()) + .def_rw("is_final", &tle::Result::isFinal) + .def_rw("output_token_ids", &tle::Result::outputTokenIds) + .def_rw("cum_log_probs", &tle::Result::cumLogProbs) + .def_rw("log_probs", &tle::Result::logProbs) + .def_rw("context_logits", &tle::Result::contextLogits) + .def_rw("generation_logits", &tle::Result::generationLogits) + .def_rw("spec_dec_fast_logits_info", &tle::Result::specDecFastLogitsInfo) + .def_rw("encoder_output", &tle::Result::encoderOutput) + .def_rw("finish_reasons", &tle::Result::finishReasons) + .def_rw("sequence_index", &tle::Result::sequenceIndex) + .def_rw("is_sequence_final", &tle::Result::isSequenceFinal) + .def_rw("decoding_iter", &tle::Result::decodingIter) + .def_rw("context_phase_params", &tle::Result::contextPhaseParams) + .def_rw("request_perf_metrics", &tle::Result::requestPerfMetrics) + .def_rw("additional_outputs", &tle::Result::additionalOutputs) + .def("__getstate__", resultGetstate) + .def("__setstate__", resultSetstate); + + m.def("deserialize_result", + [](nb::bytes& x) + { + std::string str(x.c_str(), x.size()); + std::istringstream is(str); + return tle::serialize_utils::deserialize(is); + }); + + auto responseGetstate = [](tle::Response const& self) + { return nb::make_tuple(self.getRequestId(), self.getResult(), self.getClientId()); }; + + auto responseSetstate = [](tle::Response& response, nb::tuple const& state) + { + if (state.size() != 3) + { + throw std::runtime_error("Invalid Request state!"); + } + new (&response) tle::Response( + nb::cast(state[0]), nb::cast(state[1]), nb::cast(state[2])); + }; + + nb::class_(m, "Response") + .def(nb::init>(), nb::arg("request_id"), nb::arg("error_msg"), + nb::arg("client_id") = std::nullopt) + .def(nb::init>(), nb::arg("request_id"), nb::arg("result"), + nb::arg("client_id") = std::nullopt) + .def_prop_ro("request_id", &tle::Response::getRequestId) + .def_prop_ro("client_id", &tle::Response::getClientId) + .def("has_error", &tle::Response::hasError) + .def_prop_ro("error_msg", &tle::Response::getErrorMsg) + .def_prop_ro("result", &tle::Response::getResult) + .def("clear_context_logits", + [](tle::Response& self) + { + if (!self.hasError()) + { + auto& result = const_cast(self.getResult()); + result.contextLogits.reset(); + } + }) + .def("clear_generation_logits", + [](tle::Response& self) + { + if (!self.hasError()) + { + auto& result = const_cast(self.getResult()); + result.generationLogits.reset(); + } + }) + .def("__getstate__", responseGetstate) + .def("__setstate__", responseSetstate); +} + +} // namespace tensorrt_llm::nanobind::executor diff --git a/cpp/tensorrt_llm/nanobind/executor/request.h b/cpp/tensorrt_llm/nanobind/executor/request.h new file mode 100644 index 000000000000..5a5cf9acbee6 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/executor/request.h @@ -0,0 +1,29 @@ +/* + * 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::executor +{ + +// Register bindings for executor API. +void initRequestBindings(nb::module_& m); + +} // namespace tensorrt_llm::nanobind::executor diff --git a/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp b/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp new file mode 100644 index 000000000000..f3be85bbbf24 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp @@ -0,0 +1,388 @@ +/* + * 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 "bindings.h" +#include "moeBindings.h" +#include "tensorrt_llm/kernels/communicationKernels/allReduceWorkspace.h" +#include "tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.h" +#include "tensorrt_llm/kernels/customAllReduceKernels.h" +#include "tensorrt_llm/kernels/delayStream.h" +#include "tensorrt_llm/nanobind/common/customCasters.h" +#include "tensorrt_llm/runtime/cudaEvent.h" +#include "tensorrt_llm/runtime/cudaStream.h" +#include "tensorrt_llm/runtime/decoderState.h" +#include "tensorrt_llm/runtime/decodingInput.h" +#include "tensorrt_llm/runtime/decodingOutput.h" +#include "tensorrt_llm/runtime/gptDecoder.h" +#include "tensorrt_llm/runtime/gptDecoderBatched.h" +#include "tensorrt_llm/runtime/iBuffer.h" +#include "tensorrt_llm/runtime/iGptDecoderBatched.h" +#include "tensorrt_llm/runtime/iTensor.h" +#include "tensorrt_llm/runtime/ipcUtils.h" +#include "tensorrt_llm/runtime/lookaheadBuffers.h" +#include "tensorrt_llm/runtime/loraCache.h" +#include "tensorrt_llm/runtime/mcastGPUBuffer.h" +#include "tensorrt_llm/runtime/request.h" +#include "tensorrt_llm/runtime/speculativeDecodingMode.h" +#include "tensorrt_llm/runtime/tllmRuntime.h" +#include "tensorrt_llm/runtime/torchView.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace tr = tensorrt_llm::runtime; +namespace te = tensorrt_llm::executor; + +class PyIGptDecoder : public tr::IGptDecoder +{ +public: + NB_TRAMPOLINE(tr::IGptDecoder, 5); + + 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> const& lookaheadPrompt = std::nullopt, + std::optional> const& lookaheadAlgoConfigs = std::nullopt) override + { + NB_OVERRIDE_PURE(setup, samplingConfig, batchSize, batchSlots, output, explicitDraftTokensDType, + lookaheadPrompt, lookaheadAlgoConfigs); + } + + void forwardAsync(tr::DecodingOutput& output, tr::DecodingInput const& input) override + { + NB_OVERRIDE_PURE(forwardAsync, output, input); + } + + void forwardSync(tr::DecodingOutput& output, tr::DecodingInput const& input) override + { + NB_OVERRIDE_PURE(forwardSync, output, input); + } + + tr::SamplingConfig const& getSamplingConfig() override + { + NB_OVERRIDE_PURE(getSamplingConfig); + } + + void disableLookahead(std::optional const& samplingConfig, tr::SizeType32 batchSize, + tr::DecodingInput::TensorConstPtr batchSlots) override + { + NB_OVERRIDE_PURE(disableLookahead, samplingConfig, batchSize, batchSlots); + } +}; + +namespace tensorrt_llm::nanobind::runtime +{ + +void initBindings(nb::module_& m) +{ + + nb::class_(m, "TaskLayerModuleConfig") + .def(nb::init<>()) + .def_rw("page_id", &tr::LoraCache::TaskLayerModuleConfig::pageId) + .def_rw("slot_idx", &tr::LoraCache::TaskLayerModuleConfig::slotIdx) + .def_rw("in_size", &tr::LoraCache::TaskLayerModuleConfig::inSize) + .def_rw("out_size", &tr::LoraCache::TaskLayerModuleConfig::outSize) + .def_rw("module_id", &tr::LoraCache::TaskLayerModuleConfig::moduleId) + .def_rw("layer_id", &tr::LoraCache::TaskLayerModuleConfig::layerId) + .def_rw("adapter_size", &tr::LoraCache::TaskLayerModuleConfig::adapterSize) + .def_rw("num_slots", &tr::LoraCache::TaskLayerModuleConfig::numSlots) + .def_rw("weights_in_pointer", &tr::LoraCache::TaskLayerModuleConfig::weightsInPointer) + .def_rw("weights_out_pointer", &tr::LoraCache::TaskLayerModuleConfig::weightsOutPointer) + .def_rw("scaling_vec_pointer", &tr::LoraCache::TaskLayerModuleConfig::scalingVecPointer) + .def(nb::self == nb::self); + + nb::class_(m, "BufferManager") + .def(nb::init(), nb::arg("stream"), nb::arg("trim_pool") = false) + .def_prop_ro("stream", &tr::BufferManager::getStream); + + 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")) + .def("clear_contexts", &tr::TllmRuntime::clearContexts) + .def("execute_context", &tr::TllmRuntime::executeContext, nb::arg("context_id")) + .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) + .def("has_layer_profiler", &tr::TllmRuntime::hasLayerProfiler, nb::arg("context_id")) + .def_prop_ro("layer_profiler_info", &tr::TllmRuntime::getLayerProfileInfo) + .def("report_to_profiler", &tr::TllmRuntime::reportToProfiler, nb::arg("context_id")) + .def_prop_ro("logits_dtype_from_engine", + [](tr::TllmRuntime& self) { return self.getEngine().getTensorDataType("logits"); }); + + nb::class_(m, "Request") + .def(nb::init, + std::optional>(), + nb::arg("ids"), nb::arg("input_len"), nb::arg("max_new_tokens") = std::nullopt, + nb::arg("end_id") = std::nullopt) + .def_rw("ids", &tr::decoder_batch::Request::ids) + .def_rw("input_len", &tr::decoder_batch::Request::inputLen) + .def_rw("max_new_tokens", &tr::decoder_batch::Request::maxNewTokens) + .def_rw("end_id", &tr::decoder_batch::Request::endId) + .def_rw("draft_logits", &tr::decoder_batch::Request::draftLogits) + .def_rw("embedding_bias", &tr::decoder_batch::Request::embeddingBias) + .def_rw("bad_words_list", &tr::decoder_batch::Request::badWordsList) + .def_rw("stop_words_list", &tr::decoder_batch::Request::stopWordsList) + .def_rw("generated_tokens_per_engine_step", &tr::decoder_batch::Request::generatedTokensPerEngineStep) + .def_rw("medusa_paths", &tr::decoder_batch::Request::medusaPaths) + .def_rw("medusa_tree_ids", &tr::decoder_batch::Request::medusaTreeIds) + .def_rw("lookahead_runtime_config", &tr::decoder_batch::Request::lookaheadRuntimeConfig); + nb::bind_vector>(m, "RequestVector"); + + nb::class_(m, "DecoderBatchInput") + .def(nb::init>, tr::SizeType32>(), nb::arg("logits"), + nb::arg("max_decoding_engine_tokens")) + .def(nb::init>(), nb::arg("logits")) + .def_rw("logits", &tr::decoder_batch::Input::logits) + .def_rw("max_decoder_steps", &tr::decoder_batch::Input::maxDecoderSteps) + .def_rw("batch_slots", &tr::decoder_batch::Input::batchSlots); + + nb::class_(m, "LookaheadDecodingBuffers") + .def(nb::init(), nb::arg("max_num_sequences"), + nb::arg("max_tokens_per_step"), nb::arg("buffer_manager")) + .def_rw("generation_lengths", &tr::LookaheadDecodingBuffers::generationLengths) + .def_rw("position_offsets", &tr::LookaheadDecodingBuffers::positionOffsets) + .def_rw("packed_masks", &tr::LookaheadDecodingBuffers::packedMasks) + .def_rw("position_ids", &tr::LookaheadDecodingBuffers::positionIds); + + nb::class_(m, "ExplicitDraftTokensBuffersInputs") + .def("create", &tr::ExplicitDraftTokensBuffers::Inputs::create, nb::arg("max_num_sequences"), + nb::arg("runtime"), nb::arg("model_config"), nb::arg("world_config")) + .def_rw("temperatures", &tr::ExplicitDraftTokensBuffers::Inputs::temperatures) + .def_rw("position_ids_base", &tr::ExplicitDraftTokensBuffers::Inputs::positionIdsBase) + .def_rw("generation_lengths", &tr::ExplicitDraftTokensBuffers::Inputs::generationLengths) + .def_rw("random_data_sample", &tr::ExplicitDraftTokensBuffers::Inputs::randomDataSample) + .def_rw("random_data_validation", &tr::ExplicitDraftTokensBuffers::Inputs::randomDataValidation) + .def_rw("draft_tokens", &tr::ExplicitDraftTokensBuffers::Inputs::draftTokens) + .def_rw("draft_indices", &tr::ExplicitDraftTokensBuffers::Inputs::draftIndices) + .def_rw("draft_probs", &tr::ExplicitDraftTokensBuffers::Inputs::draftProbs) + .def_rw("packed_masks", &tr::ExplicitDraftTokensBuffers::Inputs::packedMasks) + .def_rw("position_ids", &tr::ExplicitDraftTokensBuffers::Inputs::positionIds) + .def_rw("max_gen_length_host", &tr::ExplicitDraftTokensBuffers::Inputs::maxGenLengthHost) + .def_rw("generation_lengths_host", &tr::ExplicitDraftTokensBuffers::Inputs::generationLengthsHost); + + nb::class_(m, "DecodingInput"); + nb::class_(m, "DecodingOutput"); + + nb::class_(m, "CudaEvent") + .def(nb::init(), nb::arg("flags") = cudaEventDisableTiming) + .def("synchronize", &tr::CudaEvent::synchronize); + + nb::class_(m, "IGptDecoder") + .def( + "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> const& lookaheadPrompt = std::nullopt, + std::optional> const& lookaheadAlgoConfigs = std::nullopt) + { + auto tensorPtrBatchSlots = tr::TorchView::of(batchSlots); + self.setup(samplingConfig, batchSize, std::move(tensorPtrBatchSlots), output, explicitDraftTokensDType, + lookaheadPrompt, lookaheadAlgoConfigs); + }, + nb::arg("sampling_config"), nb::arg("batch_size"), nb::arg("batch_slots"), nb::arg("output") = std::nullopt, + nb::arg("explicit_draft_tokens_d_type") = std::nullopt, nb::arg("lookahead_prompt") = std::nullopt, + nb::arg("lookahead_algo_configs") = std::nullopt); + + nb::class_(m, "DecoderState") + .def(nb::init<>()) + .def("setup", &tr::decoder::DecoderState::setup, nb::arg("max_batch_size"), nb::arg("max_beam_width"), + nb::arg("max_attention_window"), nb::arg("sink_token_length"), nb::arg("max_sequence_length"), + nb::arg("dtype"), nb::arg("model_config"), nb::arg("world_config"), nb::arg("buffer_manager")) + .def("setup_cache_indirection", &tr::decoder::DecoderState::setupCacheIndirection, nb::arg("max_batch_size"), + nb::arg("max_beam_width"), nb::arg("max_attention_window"), nb::arg("buffer_manager")) + .def("setup_speculative_decoding", &tr::decoder::DecoderState::setupSpeculativeDecoding, + nb::arg("speculative_decoding_mode"), nb::arg("max_tokens_per_engine_step"), nb::arg("dtype"), + nb::arg("model_config"), nb::arg("world_config"), nb::arg("buffer_manager")) + .def_prop_ro("joint_decoding_input", &tr::decoder::DecoderState::getJointDecodingInput) + .def_prop_ro("joint_decoding_output", &tr::decoder::DecoderState::getJointDecodingOutput) + .def_prop_ro("cache_indirection_input", &tr::decoder::DecoderState::getCacheIndirectionInput) + .def_prop_ro("cache_indirection_output", &tr::decoder::DecoderState::getCacheIndirectionOutput) + .def_prop_ro( + "sequence_lengths", nb::overload_cast<>(&tr::decoder::DecoderState::getSequenceLengths, nb::const_)) + .def("get_sequence_lengths", + nb::overload_cast(&tr::decoder::DecoderState::getSequenceLengths, nb::const_), + nb::arg("batch_idx")) + .def_prop_ro("all_new_tokens", &tr::decoder::DecoderState::getAllNewTokens) + .def_prop_ro("finished_sum", &tr::decoder::DecoderState::getFinishedSum) + .def_prop_ro("finish_reasons", &tr::decoder::DecoderState::getFinishReasons) + .def_prop_ro("ids", nb::overload_cast<>(&tr::decoder::DecoderState::getIds, nb::const_)) + .def("get_ids", nb::overload_cast(&tr::decoder::DecoderState::getIds, nb::const_), + nb::arg("batch_idx")) + .def_prop_ro("gathered_ids", nb::overload_cast<>(&tr::decoder::DecoderState::getGatheredIds, nb::const_)) + .def("get_gathered_ids", + nb::overload_cast(&tr::decoder::DecoderState::getGatheredIds, nb::const_), + nb::arg("batch_idx")) + .def_prop_ro("parent_ids", &tr::decoder::DecoderState::getParentIds) + .def_prop_ro("cum_log_probs", nb::overload_cast<>(&tr::decoder::DecoderState::getCumLogProbs, nb::const_)) + .def("get_cum_log_probs", + nb::overload_cast(&tr::decoder::DecoderState::getCumLogProbs, nb::const_), + nb::arg("batch_idx")) + .def_prop_ro("log_probs", nb::overload_cast<>(&tr::decoder::DecoderState::getLogProbs, nb::const_)) + .def("get_log_probs", nb::overload_cast(&tr::decoder::DecoderState::getLogProbs, nb::const_), + nb::arg("batch_idx")) + .def_prop_ro("next_draft_tokens", &tr::decoder::DecoderState::getNextDraftTokens) + .def_prop_ro("prev_draft_tokens_lengths", &tr::decoder::DecoderState::getPrevDraftTokensLengths) + .def_prop_ro("next_draft_tokens_lengths", &tr::decoder::DecoderState::getNextDraftTokensLengths) + .def_prop_ro("accepted_lengths_cum_sum", &tr::decoder::DecoderState::getAcceptedLengthsCumSum) + .def_prop_ro("accepted_packed_paths", &tr::decoder::DecoderState::getAcceptedPackedPaths) + .def_prop_ro("finished_steps", &tr::decoder::DecoderState::getFinishedSteps) + .def_prop_ro("max_beam_width", &tr::decoder::DecoderState::getMaxBeamWidth) + .def_prop_ro("max_sequence_length", &tr::decoder::DecoderState::getMaxSequenceLength) + .def_prop_ro("max_decoding_decoder_tokens", &tr::decoder::DecoderState::getMaxDecodingDecoderTokens) + .def_prop_ro("max_decoding_engine_tokens", &tr::decoder::DecoderState::getMaxDecodingEngineTokens) + .def_prop_ro("num_decoding_engine_tokens", + nb::overload_cast<>(&tr::decoder::DecoderState::getNumDecodingEngineTokens, nb::const_)) + .def("get_num_decoding_engine_tokens", + nb::overload_cast(&tr::decoder::DecoderState::getNumDecodingEngineTokens, nb::const_), + nb::arg("batch_idx")) + .def("set_num_decoding_engine_tokens", &tr::decoder::DecoderState::setNumDecodingEngineTokens, + nb::arg("batch_idx"), nb::arg("num_tokens")) + .def_prop_ro("speculative_decoding_mode", &tr::decoder::DecoderState::getSpeculativeDecodingMode) + .def_prop_rw("generation_steps", &tr::decoder::DecoderState::getGenerationSteps, + &tr::decoder::DecoderState::setGenerationSteps); + + nb::class_(m, "GptDecoderBatched") + .def(nb::init(), nb::arg("stream")) + .def("setup", &tr::GptDecoderBatched::setup, nb::arg("mode"), nb::arg("max_batch_size"), + nb::arg("max_beam_width"), nb::arg("dtype"), nb::arg("model_config"), nb::arg("world_config")) + .def("forward_async", &tr::GptDecoderBatched::forwardAsync, nb::arg("output"), nb::arg("input")) + .def("underlying_decoder", &tr::GptDecoderBatched::getUnderlyingDecoder, nb::rv_policy::reference) + .def("finalize", &tr::GptDecoderBatched::finalize, nb::arg("decoder_state"), nb::arg("batch_idx"), + nb::arg("sampling_config"), nb::arg("streaming")) + .def_prop_ro( + "decoder_stream", + [](tr::GptDecoderBatched& self) -> tr::CudaStream const& { return *self.getDecoderStream(); }, + nb::rv_policy::reference); + + m.def( + "lamport_initialize_all", + [](intptr_t buffer_0, intptr_t buffer_1, intptr_t buffer_2, size_t size) + { + tr::lamportInitializeAll(reinterpret_cast(buffer_0), reinterpret_cast(buffer_1), + reinterpret_cast(buffer_2), size); + }, + "Lamport initialize all buffers"); + m.def( + "lamport_initialize", + [](intptr_t buffer, size_t size) + { tensorrt_llm::kernels::ar_fusion::lamport_initialize(reinterpret_cast(buffer), size, 0); }, + "Lmaport initialize buffer"); + m.def( + "delay_kernel", + [](int64_t delay_micro_secs, nb::object py_stream) + { + // Get the raw stream handle from PyTorch stream object + auto stream_ptr = nb::cast(py_stream.attr("cuda_stream")); + cudaStream_t stream = reinterpret_cast(stream_ptr); + tensorrt_llm::kernels::invokeDelayStreamKernel(delay_micro_secs, stream); + }, + "Delay kernel launch on the default stream"); + m.def( + "max_workspace_size_lowprecision", + [](int32_t tp_size) { return tensorrt_llm::kernels::max_workspace_size_lowprecision(tp_size); }, + "Calculate the maximum workspace size needed for low precision all-reduce operations"); + + nb::class_(m, "McastGPUBuffer") + .def(nb::init()) + .def("get_uc_buffer", &tensorrt_llm::runtime::McastGPUBuffer::getUCBuffer) + .def("get_mc_buffer", &tensorrt_llm::runtime::McastGPUBuffer::getMCBuffer); + + nb::enum_(m, "AllReduceFusionOp") + .value("NONE", tensorrt_llm::kernels::AllReduceFusionOp::NONE) + .value("RESIDUAL_RMS_NORM", tensorrt_llm::kernels::AllReduceFusionOp::RESIDUAL_RMS_NORM) + .value("LAST_PROCESS_FOR_UB", tensorrt_llm::kernels::AllReduceFusionOp::LAST_PROCESS_FOR_UB) + .value("RESIDUAL_RMS_PREPOST_NORM", tensorrt_llm::kernels::AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM) + .value("RESIDUAL_RMS_NORM_QUANT_FP8", tensorrt_llm::kernels::AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_FP8) + .value("RESIDUAL_RMS_NORM_QUANT_NVFP4", tensorrt_llm::kernels::AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4) + .value("RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4", + tensorrt_llm::kernels::AllReduceFusionOp::RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4) + .value("RESIDUAL_RMS_NORM_OUT_QUANT_FP8", + tensorrt_llm::kernels::AllReduceFusionOp::RESIDUAL_RMS_NORM_OUT_QUANT_FP8); + + nb::enum_(m, "AllReduceStrategy") + .value("NCCL", tensorrt_llm::kernels::AllReduceStrategyType::NCCL) + .value("MIN_LATENCY", tensorrt_llm::kernels::AllReduceStrategyType::MIN_LATENCY) + .value("AUTO", tensorrt_llm::kernels::AllReduceStrategyType::AUTO) + .value("UB", tensorrt_llm::kernels::AllReduceStrategyType::UB) + .value("ONESHOT", tensorrt_llm::kernels::AllReduceStrategyType::ONESHOT) + .value("TWOSHOT", tensorrt_llm::kernels::AllReduceStrategyType::TWOSHOT); + + // Initialize MoeLoadBalancer bindings + initMoeBindings(m); +} + +void initBindingsEarly(nb::module_& m) +{ + nb::class_(m, "SpeculativeDecodingMode") + .def(nb::init(), nb::arg("state")) + .def_static("NoneType", &tr::SpeculativeDecodingMode::None) + .def_static("DraftTokensExternal", &tr::SpeculativeDecodingMode::DraftTokensExternal) + .def_static("Medusa", &tr::SpeculativeDecodingMode::Medusa) + .def_static("Eagle", &tr::SpeculativeDecodingMode::Eagle) + .def_static("LookaheadDecoding", &tr::SpeculativeDecodingMode::LookaheadDecoding) + .def_static("ExplicitDraftTokens", &tr::SpeculativeDecodingMode::ExplicitDraftTokens) + .def_prop_ro("is_none", &tr::SpeculativeDecodingMode::isNone) + .def_prop_ro("is_draft_tokens_external", &tr::SpeculativeDecodingMode::isDraftTokensExternal) + .def_prop_ro("is_medusa", &tr::SpeculativeDecodingMode::isMedusa) + .def_prop_ro("is_eagle", &tr::SpeculativeDecodingMode::isEagle) + .def_prop_ro("is_lookahead_decoding", &tr::SpeculativeDecodingMode::isLookaheadDecoding) + .def_prop_ro("is_explicit_draft_tokens", &tr::SpeculativeDecodingMode::isExplicitDraftTokens) + .def_prop_ro("updates_position_ids", &tr::SpeculativeDecodingMode::updatesPositionIds) + .def_prop_ro("requires_attention_mask", &tr::SpeculativeDecodingMode::requiresAttentionMask) + .def_prop_ro("predicts_draft_tokens", &tr::SpeculativeDecodingMode::predictsDraftTokens) + .def_prop_ro("needs_kv_cache_rewind", &tr::SpeculativeDecodingMode::needsKVCacheRewind) + .def_prop_ro("variable_draft_length", &tr::SpeculativeDecodingMode::variableDraftLength) + .def_prop_ro("has_draft_logits", &tr::SpeculativeDecodingMode::hasDraftLogits) + .def_prop_ro("needs_decoder_prologue", &tr::SpeculativeDecodingMode::needsDecoderPrologue); +} +} // namespace tensorrt_llm::nanobind::runtime diff --git a/cpp/tensorrt_llm/nanobind/runtime/bindings.h b/cpp/tensorrt_llm/nanobind/runtime/bindings.h new file mode 100644 index 000000000000..410dac80b05e --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/runtime/bindings.h @@ -0,0 +1,30 @@ +/* + * 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::runtime +{ + +void initBindings(nb::module_& m); +void initBindingsEarly(nb::module_& m); + +} // namespace tensorrt_llm::nanobind::runtime diff --git a/cpp/tensorrt_llm/nanobind/runtime/moeBindings.cpp b/cpp/tensorrt_llm/nanobind/runtime/moeBindings.cpp new file mode 100644 index 000000000000..c26fa84b661f --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/runtime/moeBindings.cpp @@ -0,0 +1,124 @@ +/* + * 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 "moeBindings.h" +#include "tensorrt_llm/nanobind/common/customCasters.h" +#include "tensorrt_llm/runtime/moeLoadBalancer/hostAccessibleDeviceAllocator.h" +#include "tensorrt_llm/runtime/moeLoadBalancer/moeLoadBalancer.h" +#include +#include +#include + +namespace nb = nanobind; +namespace tr = tensorrt_llm::runtime; +namespace tk = tensorrt_llm::kernels; + +namespace tensorrt_llm::nanobind::runtime +{ + +void pyDoReplication(tk::MoeLoadBalanceMetaInfo const& metaInfo, std::vector& expertLoadFactor, + tr::MoePlacementCpuInfo* cpuPlacement) +{ + TLLM_CHECK_WITH_INFO( + metaInfo.expertCount == expertLoadFactor.size(), "expert_count and expert_load_factor size mismatch"); + tr::doReplication(metaInfo, expertLoadFactor.data(), cpuPlacement); +}; + +void pyDoPlacement(tk::MoeLoadBalanceMetaInfo const& metaInfo, std::vector& expertLoadFactor, + tr::MoePlacementCpuInfo* cpuPlacement) +{ + TLLM_CHECK_WITH_INFO( + metaInfo.expertCount == expertLoadFactor.size(), "expert_count and expert_load_factor size mismatch"); + tr::doPlacement(metaInfo, expertLoadFactor.data(), cpuPlacement); +}; + +void initMoeBindings(nb::module_& m) +{ + // Bind MoeWeight struct + nb::class_(m, "MoeWeight") + .def(nb::init<>()) + .def_prop_rw("weight_ptr", &tr::MoeWeight::getWeightPtr, &tr::MoeWeight::setWeightPtr) + .def_rw("height", &tr::MoeWeight::mHeight) + .def_rw("width", &tr::MoeWeight::mWidth) + .def_rw("pitch", &tr::MoeWeight::mPitch) + .def("__repr__", + [](tr::MoeWeight const& self) + { + return ""; + }); + + // Bind MoeLoadBalanceMetaInfo struct + nb::class_(m, "MoeLoadBalanceMetaInfo") + .def(nb::init(), nb::arg("expert_count"), nb::arg("top_k"), nb::arg("ep_rank"), + nb::arg("ep_size"), nb::arg("slot_count_per_rank")) + .def_rw("expert_count", &tk::MoeLoadBalanceMetaInfo::expertCount) + .def_rw("top_k", &tk::MoeLoadBalanceMetaInfo::topK) + .def_rw("ep_rank", &tk::MoeLoadBalanceMetaInfo::epRank) + .def_rw("ep_size", &tk::MoeLoadBalanceMetaInfo::epSize) + .def_rw("slot_count_per_rank", &tk::MoeLoadBalanceMetaInfo::slotCountPerRank); + + // Bind MoePlacementCpuInfo struct + nb::class_(m, "MoePlacementCpuInfo") + .def(nb::init<>()) + .def_rw("expert_replica_count", &tr::MoePlacementCpuInfo::expertReplicaCount) + .def_rw("rank_expert_ids", &tr::MoePlacementCpuInfo::rankExpertIds); + + // Bind SingleLayerMoeLoadBalancer class + nb::class_(m, "SingleLayerMoeLoadBalancer") + .def("add_single_weight_slot", &tr::SingleLayerMoeLoadBalancer::addSingleWeightSlot, nb::arg("slot_id"), + nb::arg("name"), nb::arg("weight_slot"), "Add a single weight slot for a specific slot ID") + .def("add_single_host_weight", &tr::SingleLayerMoeLoadBalancer::addSingleHostWeight, nb::arg("expert_id"), + nb::arg("name"), nb::arg("host_weight"), "Add a single host weight for a specific expert ID") + .def("set_initial_weight_assignments", &tr::SingleLayerMoeLoadBalancer::setInitialWeightAssignments, + nb::arg("initial_weight_assignments"), "Set initial weight assignments for each slot") + .def("get_pointer", &tr::SingleLayerMoeLoadBalancer::getSelfPtr, + "Get the pointer of the SingleLayerMoeLoadBalancer") + .def("get_layer_id", &tr::SingleLayerMoeLoadBalancer::getLayerId, + "Get the layer id of the SingleLayerMoeLoadBalancer"); + + // Bind MoeLoadBalancer class + nb::class_(m, "MoeLoadBalancer") + .def(nb::init(), nb::arg("ep_rank"), nb::arg("ep_size"), nb::arg("layer_updates_per_iter"), + "Initialize the MoeLoadBalancer with the specified expert parallel rank, size, and update frequency") + .def("set_use_gpu_memcpy", &tr::MoeLoadBalancer::setUseGpuMemcpy, nb::arg("use_gpu_memcpy"), + "Set whether to use GPU memcpy for weight updates") + .def("add_layer", &tr::MoeLoadBalancer::AddLayer, nb::arg("expert_count"), nb::arg("top_k"), + nb::arg("slot_count_per_rank"), "Add a new MOE layer to the load balancer") + .def("finalize_model", &tr::MoeLoadBalancer::finalizeModel, + "Finalize the model structure, must be called after all layers are added") + .def("set_warm_up_iter_count", &tr::MoeLoadBalancer::setWarmUpIterCount, nb::arg("iter_count"), + "Set the number of warm-up iterations") + .def("start_iter", &tr::MoeLoadBalancer::startIter, nb::arg("iter_id"), nb::arg("enable_statistic"), + nb::arg("enable_update_weights"), "Start a new iteration with the given ID and settings") + .def("end_iter", &tr::MoeLoadBalancer::endIter, nb::arg("iter_id"), "End the iteration with the given ID") + .def("shutdown", &tr::MoeLoadBalancer::shutdown, "Shutdown the load balancer and clean up resources"); + + m.def("is_host_accessible_device_memory_supported", &tr::HostAccessibleDeviceAllocator::isSupported, + "If current system support host accessible device memory"); + + // Bind do_replication function for testing + m.def("do_replication", &pyDoReplication, nb::arg("meta_info"), nb::arg("expert_load_factor"), + nb::arg("cpu_placement"), "Do replication"); + + // Bind do_placement function for testing + m.def("do_placement", &pyDoPlacement, nb::arg("meta_info"), nb::arg("expert_load_factor"), nb::arg("cpu_placement"), + "Do placement"); +} + +} // namespace tensorrt_llm::nanobind::runtime diff --git a/cpp/tensorrt_llm/nanobind/runtime/moeBindings.h b/cpp/tensorrt_llm/nanobind/runtime/moeBindings.h new file mode 100644 index 000000000000..73b9a3ceec8f --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/runtime/moeBindings.h @@ -0,0 +1,29 @@ +/* + * 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::runtime +{ + +void initMoeBindings(nb::module_& m); + +} // namespace tensorrt_llm::nanobind::runtime diff --git a/cpp/tensorrt_llm/nanobind/testing/modelSpecBinding.cpp b/cpp/tensorrt_llm/nanobind/testing/modelSpecBinding.cpp new file mode 100644 index 000000000000..caef94c5defd --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/testing/modelSpecBinding.cpp @@ -0,0 +1,87 @@ +/* + * 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 new file mode 100644 index 000000000000..1aababc6ff89 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/testing/modelSpecBinding.h @@ -0,0 +1,29 @@ +/* + * 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/nanobind/userbuffers/bindings.cpp b/cpp/tensorrt_llm/nanobind/userbuffers/bindings.cpp new file mode 100644 index 000000000000..82e0d0a1f0c7 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/userbuffers/bindings.cpp @@ -0,0 +1,47 @@ +/* + * 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 "bindings.h" +#include "tensorrt_llm/kernels/userbuffers/ub_interface.h" +#include "tensorrt_llm/kernels/userbuffers/userbuffersManager.h" +#include "tensorrt_llm/nanobind/common/customCasters.h" +#include + +namespace nb = nanobind; +namespace tub = tensorrt_llm::runtime::ub; + +namespace tensorrt_llm::kernels::userbuffers +{ + +void UserBufferBindings::initBindings(nb::module_& m) +{ + nb::class_(m, "UBBuffer") + .def_ro("size", &tub::UBBuffer::size) + .def_prop_ro("addr", [](tub::UBBuffer& self) { return reinterpret_cast(self.addr); }) + .def_ro("handle", &tub::UBBuffer::handle) + .def("invalid", &tub::UBBuffer::invalid); + + m.def("ub_initialize", [](int tp_size) { tub::ub_initialize(tp_size); }); + m.def("ub_is_initialized", &tub::ub_is_initialized); + m.def("ub_allocate", [](size_t bytes) { return tub::ub_allocate(bytes); }); + m.def("ub_deallocate", [](intptr_t addr) { return tub::ub_deallocate(reinterpret_cast(addr)); }); + m.def("ub_get", &tub::ub_get); + m.def("ub_supported", &tub::ub_supported); + + m.def("initialize_userbuffers_manager", &tub::initialize_userbuffers_manager); +} +} // namespace tensorrt_llm::kernels::userbuffers diff --git a/cpp/tensorrt_llm/nanobind/userbuffers/bindings.h b/cpp/tensorrt_llm/nanobind/userbuffers/bindings.h new file mode 100644 index 000000000000..15728bf6c1d0 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/userbuffers/bindings.h @@ -0,0 +1,30 @@ +/* + * 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 + +#include +namespace nb = nanobind; + +namespace tensorrt_llm::kernels::userbuffers +{ +class UserBufferBindings +{ +public: + static void initBindings(nb::module_& m); +}; +} // namespace tensorrt_llm::kernels::userbuffers diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.cpp b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.cpp index 8d80827b9008..4cec38b046a6 100644 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.cpp +++ b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.cpp @@ -108,6 +108,8 @@ 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; @@ -123,7 +125,7 @@ void GemmAllReducePlugin::allocatePersistentWorkspace() // Register and allocate workspace mWorkspace = static_cast( - getPluginRegistry()->acquirePluginResource(mWorkspaceKey, &unallocated_resource)); + getPluginRegistry()->acquirePluginResource(mWorkspaceKey.c_str(), &unallocated_resource)); TLLM_CHECK(mWorkspace != nullptr); } @@ -395,6 +397,7 @@ int GemmAllReducePlugin::enqueue(PluginTensorDesc const* inputDesc, PluginTensor 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."); @@ -513,7 +516,7 @@ void GemmAllReducePlugin::terminate() noexcept // free mWorkspace if (mWorkspace) { - getPluginRegistry()->releasePluginResource(mWorkspaceKey); + getPluginRegistry()->releasePluginResource(mWorkspaceKey.c_str()); mWorkspace = nullptr; } } diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.h b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.h index 4cd2a77a5c46..457926246002 100644 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.h +++ b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.h @@ -154,7 +154,7 @@ class GemmAllReducePlugin : public BasePlugin int mNbOutputs = 0; std::map mTypedInstantiators; - char const* mWorkspaceKey = "gemm_allreduce_workspace"; + std::string mWorkspaceKey; std::shared_ptr mGemm; // Params that are initialized during configurePlugin() GemmAllReducePersistentWorkspace* mWorkspace = nullptr; diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.cpp b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.cpp index d6e0f3b8ac69..a6f7ca2615df 100644 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.cpp +++ b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.cpp @@ -60,8 +60,12 @@ void GemmAllReducePluginProfiler::deserializeFromOwnFile(GemmIdCore gemmId, Gemm bool GemmAllReducePluginProfiler::useProfiler() { - char const* envDir = getenv("GEMM_AR_PLUGIN_PROFILE_DIR"); - return envDir != nullptr; + // 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) diff --git a/cpp/tensorrt_llm/pybind/batch_manager/algorithms.cpp b/cpp/tensorrt_llm/pybind/batch_manager/algorithms.cpp index 0f391d166508..f6bd8f02491d 100644 --- a/cpp/tensorrt_llm/pybind/batch_manager/algorithms.cpp +++ b/cpp/tensorrt_llm/pybind/batch_manager/algorithms.cpp @@ -133,16 +133,16 @@ void tensorrt_llm::pybind::batch_manager::algorithms::initBindings(pybind11::mod py::class_(m, MakeDecodingBatchInputOutput::name) .def(py::init()) - .def("__call__", &MakeDecodingBatchInputOutput::operator(), py::arg("context_requests"), - py::arg("generation_requests"), py::arg("decoder_input_buffers"), py::arg("decoder_state"), - py::arg("model_config"), py::arg("max_num_sequences"), py::arg("fused_runtime_buffers") = std::nullopt) + .def("__call__", &MakeDecodingBatchInputOutput::operator(), py::arg("decoder_input_buffers"), + py::arg("decoder_state"), py::arg("model_config"), py::arg("max_num_sequences"), + py::arg("fused_runtime_buffers") = std::nullopt) .def("name", [](MakeDecodingBatchInputOutput const&) { return MakeDecodingBatchInputOutput::name; }); py::class_(m, LogitsPostProcessor::name) .def(py::init()) - .def("__call__", &LogitsPostProcessor::operator(), py::arg("context_requests"), py::arg("generation_requests"), - py::arg("replicate_logits_post_processor"), py::arg("decoder_buffers"), py::arg("world_config"), - py::arg("runtime"), py::arg("logits_post_processor_batched") = std::nullopt) + .def("__call__", &LogitsPostProcessor::operator(), py::arg("decoder_input_buffers"), + py::arg("replicate_logits_post_processor"), py::arg("world_config"), py::arg("stream"), + py::arg("logits_post_processor_batched") = std::nullopt) .def("name", [](LogitsPostProcessor const&) { return LogitsPostProcessor::name; }); py::class_(m, CreateNewDecoderRequests::name) diff --git a/cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp index f7ba20920c9a..63d91ddab3d9 100644 --- a/cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp @@ -393,16 +393,16 @@ void initBindings(pybind11::module_& m) py::arg("max_num_sequences"), py::arg("model_config"), py::arg("world_config"), py::arg("buffer_manager")); py::class_(m, "DecoderInputBuffers") - .def(py::init(), - py::arg("max_num_sequences"), py::arg("max_batch_size"), py::arg("max_tokens_per_engine_step"), - py::arg("manager")) + .def(py::init(), py::arg("max_batch_size"), + py::arg("max_tokens_per_engine_step"), py::arg("manager")) .def_readwrite("setup_batch_slots", &tb::DecoderInputBuffers::setupBatchSlots) .def_readwrite("setup_batch_slots_device", &tb::DecoderInputBuffers::setupBatchSlotsDevice) .def_readwrite("fill_values", &tb::DecoderInputBuffers::fillValues) .def_readwrite("fill_values_device", &tb::DecoderInputBuffers::fillValuesDevice) .def_readwrite("inputs_ids", &tb::DecoderInputBuffers::inputsIds) .def_readwrite("forward_batch_slots", &tb::DecoderInputBuffers::forwardBatchSlots) - .def_readwrite("logits", &tb::DecoderInputBuffers::logits); + .def_readwrite("logits", &tb::DecoderInputBuffers::logits) + .def_readwrite("decoder_requests", &tb::DecoderInputBuffers::decoderRequests); py::class_(m, "DecoderOutputBuffers") .def_readwrite("sequence_lengths_host", &tb::DecoderOutputBuffers::sequenceLengthsHost) diff --git a/cpp/tensorrt_llm/pybind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/pybind/batch_manager/cacheTransceiver.cpp index 87b0a26a79e7..d92336e6bdf7 100644 --- a/cpp/tensorrt_llm/pybind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/pybind/batch_manager/cacheTransceiver.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -80,21 +81,15 @@ void tb::CacheTransceiverBindings::initBindings(py::module_& m) .def("check_gen_transfer_status", &BaseCacheTransceiver::checkGenTransferStatus) .def("check_gen_transfer_complete", &BaseCacheTransceiver::checkGenTransferComplete); - py::enum_(m, "CommType") - .value("UNKNOWN", tb::CacheTransceiver::CommType::UNKNOWN) - .value("MPI", tb::CacheTransceiver::CommType::MPI) - .value("UCX", tb::CacheTransceiver::CommType::UCX) - .value("NIXL", tb::CacheTransceiver::CommType::NIXL); - py::enum_(m, "AttentionType") .value("DEFAULT", executor::kv_cache::CacheState::AttentionType::kDEFAULT) .value("MLA", executor::kv_cache::CacheState::AttentionType::kMLA); py::classh(m, "CacheTransceiver") - .def(py::init, SizeType32, SizeType32, runtime::WorldConfig, nvinfer1::DataType, - executor::kv_cache::CacheState::AttentionType, std::optional>(), - py::arg("cache_manager"), py::arg("comm_type"), py::arg("num_kv_heads_per_layer"), py::arg("size_per_head"), + .def(py::init, SizeType32, SizeType32, + runtime::WorldConfig, nvinfer1::DataType, executor::kv_cache::CacheState::AttentionType, + std::optional>(), + py::arg("cache_manager"), py::arg("num_kv_heads_per_layer"), py::arg("size_per_head"), py::arg("tokens_per_block"), py::arg("world_config"), py::arg("dtype"), py::arg("attention_type"), py::arg("cache_transceiver_config") = std::nullopt); @@ -102,5 +97,5 @@ void tb::CacheTransceiverBindings::initBindings(py::module_& m) .def(py::init>(), py::arg("cache_manager"), py::arg("max_num_tokens") = std::nullopt) .def_static("pre_alloc_buffer_size", &tb::kv_cache_manager::CacheTransBufferManager::preAllocBufferSize, - py::arg("max_num_tokens") = std::nullopt); + py::arg("cache_size_bytes_per_token_per_window"), py::arg("cache_transceiver_config") = py::none()); } diff --git a/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp index a75db66eaee1..255b0f8efa33 100644 --- a/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp @@ -298,7 +298,8 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(py::module_& m) .def_readwrite("alloc_new_blocks", &tbk::KvCacheStats::allocNewBlocks) .def_readwrite("reused_blocks", &tbk::KvCacheStats::reusedBlocks) .def_readwrite("missed_blocks", &tbk::KvCacheStats::missedBlocks) - .def_readwrite("cache_hit_rate", &tbk::KvCacheStats::cacheHitRate); + .def_readwrite("cache_hit_rate", &tbk::KvCacheStats::cacheHitRate) + .def_readwrite("num_free_blocks_per_window_size", &tbk::KvCacheStats::numFreeBlocksPerWindowSize); py::class_(m, "TempAttentionWindowInputs") .def(py::init<>()) @@ -468,7 +469,8 @@ void tb::BasePeftCacheManagerBindings::initBindings(py::module_& m) py::classh(m, "PeftCacheManager") .def(py::init(), - py::arg("config"), py::arg("model_config"), py::arg("world_config"), py::arg("buffer_manager")); + py::arg("config"), py::arg("model_config"), py::arg("world_config"), py::arg("buffer_manager")) + .def("is_task_cached", &tb::PeftCacheManager::isTaskCached, py::arg("taskId")); py::classh(m, "NoOpPeftCacheManager").def(py::init()); } diff --git a/cpp/tensorrt_llm/pybind/bindings.cpp b/cpp/tensorrt_llm/pybind/bindings.cpp index 1a5841d4b7aa..a004c872a7fc 100644 --- a/cpp/tensorrt_llm/pybind/bindings.cpp +++ b/cpp/tensorrt_llm/pybind/bindings.cpp @@ -170,7 +170,7 @@ PYBIND11_MODULE(TRTLLM_PYBIND_MODULE, m) .value("CONTINUOUS", tr::ModelConfig::KVCacheType::kCONTINUOUS) .value("PAGED", tr::ModelConfig::KVCacheType::kPAGED) .value("DISABLED", tr::ModelConfig::KVCacheType::kDISABLED) - .def(py::init(&tr::ModelConfig::KVCacheTypeFromString)); + .def("from_string", &tr::ModelConfig::KVCacheTypeFromString); py::enum_(m, "LayerType") .value("ATTENTION", tr::ModelConfig::LayerType::kATTENTION) @@ -355,7 +355,10 @@ PYBIND11_MODULE(TRTLLM_PYBIND_MODULE, m) }; auto SamplingConfigSetState = [](py::tuple t) -> tr::SamplingConfig { - assert(t.size() == 19); + if (t.size() != 19) + { + throw std::runtime_error("Invalid SamplingConfig state!"); + } tr::SamplingConfig config; config.beamWidth = t[0].cast(); diff --git a/cpp/tensorrt_llm/pybind/executor/bindings.cpp b/cpp/tensorrt_llm/pybind/executor/bindings.cpp index d09157e1a8bf..a8f6aaef73d7 100644 --- a/cpp/tensorrt_llm/pybind/executor/bindings.cpp +++ b/cpp/tensorrt_llm/pybind/executor/bindings.cpp @@ -244,7 +244,17 @@ void initBindings(pybind11::module_& m) py::class_>( executor_kv_cache, "KVCacheEventManager") - .def("get_latest_events", &tle::KVCacheEventManager::getLatestEvents, py::arg("timeout") = std::nullopt); + .def( + "get_latest_events", + [](tle::KVCacheEventManager& self, std::optional timeout_ms = std::nullopt) + { + if (timeout_ms) + { + return self.getLatestEvents(std::chrono::milliseconds(static_cast(*timeout_ms))); + } + return self.getLatestEvents(std::nullopt); + }, + py::arg("timeout_ms") = std::nullopt); tensorrt_llm::pybind::executor::initRequestBindings(m); tensorrt_llm::pybind::executor::initConfigBindings(m); diff --git a/cpp/tensorrt_llm/pybind/executor/executorConfig.cpp b/cpp/tensorrt_llm/pybind/executor/executorConfig.cpp index 71a0b4af7241..87f326358666 100644 --- a/cpp/tensorrt_llm/pybind/executor/executorConfig.cpp +++ b/cpp/tensorrt_llm/pybind/executor/executorConfig.cpp @@ -336,7 +336,7 @@ void initConfigBindings(pybind11::module_& m) throw std::runtime_error("Invalid extendedRuntimePerfKnobConfig state!"); } return tle::ExtendedRuntimePerfKnobConfig( - state[0].cast(), state[1].cast(), state[2].cast(), state[2].cast()); + state[0].cast(), state[1].cast(), state[2].cast(), state[3].cast()); }; auto extendedRuntimePerfKnobConfigGetstate = [](tle::ExtendedRuntimePerfKnobConfig const& self) { @@ -407,21 +407,44 @@ void initConfigBindings(pybind11::module_& m) "stop_token_ids", &tle::GuidedDecodingConfig::getStopTokenIds, &tle::GuidedDecodingConfig::setStopTokenIds) .def(py::pickle(guidedDecodingConfigGetstate, guidedDecodingConfigSetstate)); - auto cacheTransceiverConfigGetstate - = [](tle::CacheTransceiverConfig const& self) { return py::make_tuple(self.getMaxNumTokens()); }; + auto cacheTransceiverConfigGetstate = [](tle::CacheTransceiverConfig const& self) + { return py::make_tuple(self.getBackendType(), self.getMaxTokensInBuffer()); }; auto cacheTransceiverConfigSetstate = [](py::tuple const& state) { - if (state.size() != 1) + if (state.size() != 2) { throw std::runtime_error("Invalid CacheTransceiverConfig state!"); } - return tle::CacheTransceiverConfig(state[0].cast>()); + return tle::CacheTransceiverConfig( + state[0].cast(), state[1].cast>()); }; + py::enum_(m, "CacheTransceiverBackendType") + .value("DEFAULT", tle::CacheTransceiverConfig::BackendType::DEFAULT) + .value("MPI", tle::CacheTransceiverConfig::BackendType::MPI) + .value("UCX", tle::CacheTransceiverConfig::BackendType::UCX) + .value("NIXL", tle::CacheTransceiverConfig::BackendType::NIXL) + .def("from_string", + [](std::string const& str) + { + if (str == "DEFAULT" || str == "default") + return tle::CacheTransceiverConfig::BackendType::DEFAULT; + if (str == "MPI" || str == "mpi") + return tle::CacheTransceiverConfig::BackendType::MPI; + if (str == "UCX" || str == "ucx") + return tle::CacheTransceiverConfig::BackendType::UCX; + if (str == "NIXL" || str == "nixl") + return tle::CacheTransceiverConfig::BackendType::NIXL; + throw std::runtime_error("Invalid backend type: " + str); + }); + py::class_(m, "CacheTransceiverConfig") - .def(py::init>(), py::arg("max_num_tokens") = py::none()) - .def_property("max_num_tokens", &tle::CacheTransceiverConfig::getMaxNumTokens, - &tle::CacheTransceiverConfig::setMaxNumTokens) + .def(py::init, std::optional>(), + py::arg("backend") = std::nullopt, py::arg("max_tokens_in_buffer") = std::nullopt) + .def_property( + "backend", &tle::CacheTransceiverConfig::getBackendType, &tle::CacheTransceiverConfig::setBackendType) + .def_property("max_tokens_in_buffer", &tle::CacheTransceiverConfig::getMaxTokensInBuffer, + &tle::CacheTransceiverConfig::setMaxTokensInBuffer) .def(py::pickle(cacheTransceiverConfigGetstate, cacheTransceiverConfigSetstate)); auto executorConfigGetState = [](py::object const& self) diff --git a/cpp/tensorrt_llm/runtime/ipcNvlsMemory.cu b/cpp/tensorrt_llm/runtime/ipcNvlsMemory.cu index c685966148f5..031ac92168a2 100644 --- a/cpp/tensorrt_llm/runtime/ipcNvlsMemory.cu +++ b/cpp/tensorrt_llm/runtime/ipcNvlsMemory.cu @@ -295,6 +295,7 @@ public: // Clean up MPI_Group_free(&new_group); MPI_Group_free(&world_group); + MPI_Comm_free(&new_comm); return nvls_handle; } @@ -401,14 +402,14 @@ void MPI_group_barrier(std::set group) MPI_Comm new_comm; // Get the group of the world communicator - MPI_Comm_group(MPI_COMM_WORLD, &world_group); + MPI_Comm_group(COMM_SESSION, &world_group); // Create a new group containing only the ranks we want std::vector ranks(group.begin(), group.end()); MPI_Group_incl(world_group, ranks.size(), ranks.data(), &new_group); // Create a new communicator from the group - MPI_Comm_create_group(MPI_COMM_WORLD, new_group, 0, &new_comm); + MPI_Comm_create_group(COMM_SESSION, new_group, 0, &new_comm); // Use the new communicator for the barrier MPI_Barrier(new_comm); @@ -510,6 +511,8 @@ IpcNvlsHandle* ipcNvlsAllocate(size_t size, std::set group) MPI_Barrier(new_comm); + MPI_Comm_free(&new_comm); + return handle; #else TLLM_THROW("ipcNvlsAllocate needs to be compiled with ENABLE_MULTI_DEVICE"); diff --git a/cpp/tensorrt_llm/runtime/mcastDeviceMemory.cpp b/cpp/tensorrt_llm/runtime/mcastDeviceMemory.cpp index b45d903367e1..950215e7542c 100644 --- a/cpp/tensorrt_llm/runtime/mcastDeviceMemory.cpp +++ b/cpp/tensorrt_llm/runtime/mcastDeviceMemory.cpp @@ -50,7 +50,7 @@ McastDeviceMemory::McastDeviceMemory( , mMcHandle(0) { - cudaSetDevice(mDeviceIdx); + TLLM_CUDA_CHECK(cudaSetDevice(mDeviceIdx)); // Check if the device support multicasting int multicast_supported{0}; TLLM_CU_CHECK(cuDeviceGetAttribute(&multicast_supported, CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, mDeviceIdx)); @@ -82,34 +82,41 @@ McastDeviceMemory::McastDeviceMemory( { allocNvlsMcastMem(mSignalPadOffset + kSIGNAL_PAD_SIZE); } - mSignalPadsDev.resize(mGroupSize); + // Initialize signal pads + mSignalPads.resize(mGroupSize); for (size_t i = 0; i < mGroupSize; i++) { - mSignalPadsDev[i] = mUcPtrs[i] + mSignalPadOffset; + mSignalPads[i] = mUcPtrs[i] + mSignalPadOffset; if (i == mGroupRank) { - cuMemsetD8(mSignalPadsDev[i], 0, kSIGNAL_PAD_SIZE); + cuMemsetD8(mSignalPads[i], 0, kSIGNAL_PAD_SIZE); } } + // Copy host array of pointers to device array + TLLM_CUDA_CHECK(cudaMalloc(&mSignalPadsDev, mGroupSize * sizeof(CUdeviceptr))); + TLLM_CUDA_CHECK(cudaMalloc(&mUcPtrsDev, mGroupSize * sizeof(CUdeviceptr))); + TLLM_CUDA_CHECK( + cudaMemcpy(mSignalPadsDev, mSignalPads.data(), mGroupSize * sizeof(CUdeviceptr), cudaMemcpyHostToDevice)); + TLLM_CUDA_CHECK(cudaMemcpy(mUcPtrsDev, mUcPtrs.data(), mGroupSize * sizeof(CUdeviceptr), cudaMemcpyHostToDevice)); } McastDeviceMemory::~McastDeviceMemory() { tensorrt_llm::common::unregisterMcastDevMemBuffer(this); + TLLM_CUDA_CHECK(cudaFree(mSignalPadsDev)); + TLLM_CUDA_CHECK(cudaFree(mUcPtrsDev)); + if (mIsMNNvlink) { for (uint32_t rank = 0; rank < mGroupSize; rank++) { - if (rank == mGroupRank) - { - cuMemRelease(mUcHandles[rank]); - } - else - { - mUcHandles[rank] = 0; - } + TLLM_CU_CHECK(cuMemUnmap(mUcPtrs[rank], mAllocationSize)); + // We need to release the handle on each rank + TLLM_CU_CHECK(cuMemRelease(mUcHandles[rank])); } - cuMemRelease(mMcHandle); + TLLM_CU_CHECK(cuMemUnmap(mMcPtr, mAllocationSize)); + TLLM_CU_CHECK(cuMemAddressFree(mMcPtr, mAllocationSize)); + TLLM_CU_CHECK(cuMemRelease(mMcHandle)); } else { diff --git a/cpp/tensorrt_llm/runtime/mcastDeviceMemory.h b/cpp/tensorrt_llm/runtime/mcastDeviceMemory.h index 339c62c310c4..4afcc05223d9 100644 --- a/cpp/tensorrt_llm/runtime/mcastDeviceMemory.h +++ b/cpp/tensorrt_llm/runtime/mcastDeviceMemory.h @@ -44,16 +44,18 @@ class McastDeviceMemory McastDeviceMemory(size_t bufSize, uint32_t groupSize, uint32_t groupRank, int deviceIdx, bool mnNvlink); + // We don't register the pointer in these two functions since we don't expect any python-level code would call + // to obtain the raw pointers. //! Get the raw array of signal pad pointers to all ranks (including self) void** getSignalPadPtrsDev() { - return reinterpret_cast(mSignalPadsDev.data()); + return mSignalPadsDev; } //! Get the raw array of unicast pointers to all ranks (including self) void** getBufferPtrsDev() { - return reinterpret_cast(mUcPtrs.data()); + return mUcPtrsDev; } //! Get the raw unicast pointer to a given rank @@ -93,11 +95,17 @@ class McastDeviceMemory size_t mAllocationSize; CUdeviceptr mMcPtr; - std::vector mUcPtrs; - std::vector mSignalPadsDev; CUmemGenericAllocationHandle mMcHandle; std::vector mUcHandles; + // Host array of pointers + std::vector mUcPtrs; + std::vector mSignalPads; + + // Device array of pointers + void** mUcPtrsDev; + void** mSignalPadsDev; + // For intra-node mcast tensorrt_llm::runtime::IpcNvlsHandle* mNvlsHandle; diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index b593147b5847..8e41e2a2886f 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -85,6 +85,7 @@ add_library( selectiveScanOp.cpp userbuffersFinalizeOp.cpp userbuffersTensor.cpp + weightOnlyQuantGemm.cpp weightOnlyQuantOp.cpp mtpOp.cpp loraOp.cpp diff --git a/cpp/tensorrt_llm/thop/allgatherOp.cpp b/cpp/tensorrt_llm/thop/allgatherOp.cpp index 5a1a759d6a40..81310003de75 100644 --- a/cpp/tensorrt_llm/thop/allgatherOp.cpp +++ b/cpp/tensorrt_llm/thop/allgatherOp.cpp @@ -55,70 +55,61 @@ class AllgatherOp return 0; } - torch::Tensor run(torch::Tensor input, torch::optional> sizes) + std::vector run_list(torch::TensorList input_list, torch::optional> sizes) { TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used"); - auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); - auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); - std::vector outputShape = input.sizes().vec(); - if (sizes.has_value()) - { - outputShape[0] = std::accumulate(sizes.value().begin(), sizes.value().end(), 0, std::plus<>{}); - } - else - { - outputShape[0] *= mGroup.size(); - } - auto output = torch::empty(outputShape, input.options()); bool use_nccl_allgather = !sizes.has_value() || std::all_of(sizes.value().begin(), sizes.value().end(), [&sizes](int64_t size) { return size == sizes.value()[0]; }); - if (use_nccl_allgather) - { - NCCLCHECK_THROW(ncclAllGather(input.data_ptr(), output.mutable_data_ptr(), input.numel(), - (*getDtypeMap())[type], *mNcclComm, stream)); - } - else - { - size_t numel_base = std::accumulate(outputShape.cbegin() + 1, outputShape.cend(), 1, std::multiplies<>{}); - int64_t split_offset = 0; - ncclGroupStart(); - for (int root = 0; root < static_cast(mGroup.size()); ++root) - { - auto split_size = sizes.value()[root]; - NCCLCHECK_THROW(ncclBroadcast(input.data_ptr(), - output.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).mutable_data_ptr(), - numel_base * split_size, (*getDtypeMap())[type], root, *mNcclComm, stream)); - split_offset += split_size; - } - ncclGroupEnd(); - } - return output; - } - - std::vector run_list(torch::TensorList input_list, torch::optional> sizes) - { + int64_t sum_sizes + = sizes.has_value() ? std::accumulate(sizes.value().begin(), sizes.value().end(), 0, std::plus<>{}) : 0; std::vector output_list; output_list.reserve(input_list.size()); - bool use_nccl_allgather = !sizes.has_value() - || std::all_of(sizes.value().begin(), sizes.value().end(), - [&sizes](int64_t size) { return size == sizes.value()[0]; }); - if (use_nccl_allgather) - { - ncclGroupStart(); - } + ncclGroupStart(); for (auto const& input : input_list) { - auto output = run(input, sizes); + auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); + auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); + std::vector outputShape = input.sizes().vec(); + if (sizes.has_value()) + { + outputShape[0] = sum_sizes; + } + else + { + outputShape[0] *= mGroup.size(); + } + auto output = torch::empty(outputShape, input.options()); + if (use_nccl_allgather) + { + ncclAllGather(input.data_ptr(), output.mutable_data_ptr(), input.numel(), (*getDtypeMap())[type], + *mNcclComm, stream); + } + else + { + size_t numel_base + = std::accumulate(outputShape.cbegin() + 1, outputShape.cend(), 1, std::multiplies<>{}); + int64_t split_offset = 0; + for (int root = 0; root < static_cast(mGroup.size()); ++root) + { + auto split_size = sizes.value()[root]; + ncclBroadcast(input.data_ptr(), + output.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).mutable_data_ptr(), + numel_base * split_size, (*getDtypeMap())[type], root, *mNcclComm, stream); + split_offset += split_size; + } + } output_list.push_back(output); } - if (use_nccl_allgather) - { - ncclGroupEnd(); - } + NCCLCHECK_THROW(ncclGroupEnd()); return output_list; } + torch::Tensor run(torch::Tensor input, torch::optional> sizes) + { + return run_list({input}, sizes)[0]; + } + private: std::set mGroup; std::shared_ptr mNcclComm; diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index f377220be886..7a77fc49bbf3 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -101,7 +101,9 @@ class Runner : public RunnerBase // Always reserve SemaphoreArray (for multi-block mode) as MMHA may enable multi-block mode when shared memory // is not enough. - op.reserveSemaphoreArray(op.mNumHeads * max_num_requests); + // The attention kernel might split the heads into multiple blocks, so we might need to reserve more semaphores. + // Use mMultiProcessorCount as the lower-bound to make sure we reserve enough semaphores. + op.reserveSemaphoreArray(std::max(op.mNumHeads * max_num_requests, op.getMultiProcessorCount())); } int64_t getWorkspaceSize(AttentionOp const& op, int const num_tokens, int const max_attention_window_size, @@ -671,7 +673,8 @@ bool attention_supports_nvfp4_output(int64_t const num_heads, int64_t const num_ bool const use_paged_context_fmha, bool is_mla_enable) { // Only Blackwell supports NVFP4 output. - if (tensorrt_llm::common::getSMVersion() < 100) + // SM 120 does not support NVFP4 output. + if (tensorrt_llm::common::getSMVersion() < 100 || tensorrt_llm::common::getSMVersion() == 120) { return false; } diff --git a/cpp/tensorrt_llm/thop/cublasScaledMM.cpp b/cpp/tensorrt_llm/thop/cublasScaledMM.cpp index ed90c31cf5d2..d39b7b693fe2 100644 --- a/cpp/tensorrt_llm/thop/cublasScaledMM.cpp +++ b/cpp/tensorrt_llm/thop/cublasScaledMM.cpp @@ -66,6 +66,9 @@ AlgoListType fp8_algo_list = { {{8, 8192, 8192}, {393, 36, 1, 0, 0, 5, 2}}, // [-algo66 -m_tile10 -m_stages36 -m_numsK1 -m_reduction0 -m_swizzle0 -m_custom1 -m_mma0 -m_cga2 -m_scheduling1] {{8, 8192, 57344}, {10, 36, 1, 0, 0, 1, 2}}, + // Llama-3.3-70B TP4 (this is the default algo on B200. Here we aim to use the same algo on GB200.) + // [-algo66 -m_tile393 -m_stages36 -m_numsK1 -m_reduction0 -m_swizzle0 -m_custom1 -m_mma0 -m_cga4 -m_scheduling1] + {{8, 8192, 14336}, {393, 36, 1, 0, 1, 1, 4}}, }; void set_algo_attr(cublasLtMatmulAlgo_t& algo, std::array const& attr_list) diff --git a/cpp/tensorrt_llm/thop/finegrained_mixed_dtype_gemm_thop.cpp b/cpp/tensorrt_llm/thop/finegrained_mixed_dtype_gemm_thop.cpp index 9fa36d16b8e4..f2255604e214 100644 --- a/cpp/tensorrt_llm/thop/finegrained_mixed_dtype_gemm_thop.cpp +++ b/cpp/tensorrt_llm/thop/finegrained_mixed_dtype_gemm_thop.cpp @@ -44,51 +44,107 @@ namespace torch_ext { -W4A16GemmRunner::W4A16GemmRunner(at::ScalarType activationDtype, int64_t quant_mode) +finegrainedMixedDtypeGemmRunner::finegrainedMixedDtypeGemmRunner( + at::ScalarType activationDtype, at::ScalarType outputDtype, int64_t quant_mode) : mActivationDtype(activationDtype) + , mOutputDtype(outputDtype) { if (quant_mode == 0) { if (activationDtype == at::ScalarType::Half) { + TORCH_CHECK( + outputDtype == activationDtype, "Activation dtype needs to match Output stype", activationDtype); mGemmRunner = std::make_shared>(); } else if (activationDtype == at::ScalarType::BFloat16) { + TORCH_CHECK( + outputDtype == activationDtype, "Activation dtype needs to match Output stype", activationDtype); mGemmRunner = std::make_shared< tensorrt_llm::kernels::cutlass_kernels::CutlassFpAIntBGemmRunner<__nv_bfloat16, cutlass::uint4b_t, cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY, __nv_bfloat16, __nv_bfloat16, __nv_bfloat16>>(); } + + else if (activationDtype == at::ScalarType::Float8_e4m3fn) + { + if (outputDtype == at::ScalarType::BFloat16) + { + mGemmRunner = std::make_shared< + tensorrt_llm::kernels::cutlass_kernels::CutlassFpAIntBGemmRunner<__nv_fp8_e4m3, cutlass::uint4b_t, + cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY, half, __nv_bfloat16, __nv_bfloat16>>(); + } + else if (outputDtype == at::ScalarType::Half) + { + mGemmRunner + = std::make_shared>(); + } + else + { + TORCH_CHECK(false, "Unsupported output dtype for Float8_e4m3fn activation", outputDtype); + } + } + else + { + TORCH_CHECK(false, "Unsupported activation dtype", activationDtype); + } } + else if (quant_mode == 1) { if (activationDtype == at::ScalarType::Half) { + TORCH_CHECK( + outputDtype == activationDtype, "Activation dtype needs to match Output stype", activationDtype); mGemmRunner = std::make_shared>(); } else if (activationDtype == at::ScalarType::BFloat16) { + TORCH_CHECK( + outputDtype == activationDtype, "Activation dtype needs to match Output stype", activationDtype); mGemmRunner = std::make_shared>(); } + else if (activationDtype == at::ScalarType::Float8_e4m3fn) + { + if (outputDtype == at::ScalarType::BFloat16) + { + mGemmRunner = std::make_shared< + tensorrt_llm::kernels::cutlass_kernels::CutlassFpAIntBGemmRunner<__nv_fp8_e4m3, cutlass::uint4b_t, + cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS, half, __nv_bfloat16, __nv_bfloat16>>(); + } + else if (outputDtype == at::ScalarType::Half) + { + mGemmRunner = std::make_shared< + tensorrt_llm::kernels::cutlass_kernels::CutlassFpAIntBGemmRunner<__nv_fp8_e4m3, cutlass::uint4b_t, + cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS, half, half, half>>(); + } + else + { + TORCH_CHECK(false, "Unsupported output dtype for Float8_e4m3fn activation", outputDtype); + } + } } else { - TORCH_CHECK(false, "Unsupported quant mode for W4A16GemmRunner: ", quant_mode); + TORCH_CHECK(false, "Unsupported quant mode for finegrainedMixedDtypeGemmRunner: ", quant_mode); } - TORCH_CHECK(mGemmRunner, "Failed to create W4A16 GEMM runner for activation type ", c10::toString(activationDtype)); + TORCH_CHECK(mGemmRunner, "Failed to create finegrained Mixed Dtype GEMM runner for activation type ", + c10::toString(activationDtype)); mConfigs = mGemmRunner->getConfigs(); // Get configs via the interface - TORCH_CHECK(!mConfigs.empty(), "Failed to get CUTLASS configs for W4A16 GEMM with activation type ", + TORCH_CHECK(!mConfigs.empty(), "Failed to get CUTLASS configs for finegrainedMixedDtype GEMM with activation type ", c10::toString(activationDtype)); } -at::Tensor W4A16GemmRunner::runGemm(at::Tensor const& A, at::Tensor const& B_packed, at::Tensor const& scales, - int64_t group_size_long, int64_t configIdx, std::optional bias, std::optional zeros) const +at::Tensor finegrainedMixedDtypeGemmRunner::runGemm(at::Tensor const& A, at::Tensor const& B_packed, + at::Tensor const& scales, int64_t group_size_long, int64_t configIdx, std::optional bias, + std::optional zeros, double alpha) const { TORCH_CHECK(A.is_cuda() && B_packed.is_cuda() && scales.is_cuda(), "All input tensors must be on CUDA"); TORCH_CHECK(A.scalar_type() == mActivationDtype, "Activation tensor A's dtype ", c10::toString(A.scalar_type()), @@ -96,6 +152,7 @@ at::Tensor W4A16GemmRunner::runGemm(at::Tensor const& A, at::Tensor const& B_pac TORCH_CHECK(B_packed.scalar_type() == torch::kQUInt4x2 || B_packed.scalar_type() == torch::kInt8 || B_packed.scalar_type() == torch::kUInt8, "B_packed must be quint4x2, int8, or uint8 (view of quantized data)"); + TORCH_CHECK(A.is_contiguous() && B_packed.is_contiguous() && scales.is_contiguous(), "All input tensors (A, B_packed, scales) must be contiguous"); @@ -156,19 +213,18 @@ at::Tensor W4A16GemmRunner::runGemm(at::Tensor const& A, at::Tensor const& B_pac output_shape_vec.push_back(N_orig); } - // Set output dtype based on activation dtype torch::ScalarType output_dtype; - if (mActivationDtype == at::ScalarType::Half) + if (mOutputDtype == at::ScalarType::Half) { output_dtype = torch::kFloat16; } - else if (mActivationDtype == at::ScalarType::BFloat16) + else if (mOutputDtype == at::ScalarType::BFloat16) { output_dtype = torch::kBFloat16; } else { - TORCH_CHECK(false, "Unsupported activation type for output dtype determination"); + TORCH_CHECK(false, "Unsupported output dtype"); } torch::Tensor C_tensor = torch::empty(output_shape_vec, A.options().dtype(output_dtype)); @@ -201,16 +257,15 @@ at::Tensor W4A16GemmRunner::runGemm(at::Tensor const& A, at::Tensor const& B_pac cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.device().index()); - mGemmRunner->gemm(A_ptr, B_ptr, scales_ptr, zeros_ptr, bias_ptr, - 1.0f, // alpha - C_ptr, M, N_orig, K, group_size, gemm_config_to_use, workspace_ptr, workspace_bytes, stream); + mGemmRunner->gemm(A_ptr, B_ptr, scales_ptr, zeros_ptr, bias_ptr, static_cast(alpha), C_ptr, M, N_orig, K, + group_size, gemm_config_to_use, workspace_ptr, workspace_bytes, stream); return C_tensor; } -int64_t W4A16GemmRunner::getNumConfigs() const +int64_t finegrainedMixedDtypeGemmRunner::getNumConfigs() const { - TORCH_CHECK(mGemmRunner, "W4A16GemmRunner not initialized properly."); + TORCH_CHECK(mGemmRunner, "finegrainedMixedDtypeGemmRunner not initialized properly."); return static_cast(mConfigs.size()); } @@ -218,8 +273,8 @@ int64_t W4A16GemmRunner::getNumConfigs() const TORCH_LIBRARY_FRAGMENT(trtllm, m) { - m.class_("W4A16GemmRunner") - .def(torch::init()) - .def("run_gemm", &torch_ext::W4A16GemmRunner::runGemm) - .def("get_num_configs", &torch_ext::W4A16GemmRunner::getNumConfigs); + m.class_("finegrainedMixedDtypeGemmRunner") + .def(torch::init()) + .def("run_gemm", &torch_ext::finegrainedMixedDtypeGemmRunner::runGemm) + .def("get_num_configs", &torch_ext::finegrainedMixedDtypeGemmRunner::getNumConfigs); } diff --git a/cpp/tensorrt_llm/thop/finegrained_mixed_dtype_gemm_thop.h b/cpp/tensorrt_llm/thop/finegrained_mixed_dtype_gemm_thop.h index 1b2083de5a01..5bda7be3eb6d 100644 --- a/cpp/tensorrt_llm/thop/finegrained_mixed_dtype_gemm_thop.h +++ b/cpp/tensorrt_llm/thop/finegrained_mixed_dtype_gemm_thop.h @@ -24,14 +24,15 @@ namespace torch_ext { -class W4A16GemmRunner : public torch::CustomClassHolder +class finegrainedMixedDtypeGemmRunner : public torch::CustomClassHolder { public: - explicit W4A16GemmRunner(at::ScalarType activationDtype, int64_t quant_mode = 0); + explicit finegrainedMixedDtypeGemmRunner( + at::ScalarType activationDtype, at::ScalarType outputDtype, int64_t quant_mode = 0); at::Tensor runGemm(at::Tensor const& A, at::Tensor const& B_packed, at::Tensor const& scales, int64_t group_size_long, int64_t configIdx = -1, std::optional bias = std::nullopt, - std::optional zeros = std::nullopt) const; + std::optional zeros = std::nullopt, double alpha = 1.0f) const; int64_t getNumConfigs() const; @@ -39,6 +40,7 @@ class W4A16GemmRunner : public torch::CustomClassHolder std::shared_ptr mGemmRunner; std::vector mConfigs; at::ScalarType mActivationDtype; + at::ScalarType mOutputDtype; }; } // namespace torch_ext diff --git a/cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp b/cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp index 11b24e7a9897..ad4588a6ce58 100644 --- a/cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp +++ b/cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp @@ -54,8 +54,8 @@ void logitsBitmask(std::vector const& logits, std::vector(bitmask[i].data_ptr()); } - auto logitsPtrs = logitsPtrsHost.to(torch::kCUDA); - auto bitmaskPtrs = bitmaskPtrsHost.to(torch::kCUDA); + auto logitsPtrs = logitsPtrsHost.to(torch::kCUDA, /*non_blocking=*/true); + auto bitmaskPtrs = bitmaskPtrsHost.to(torch::kCUDA, /*non_blocking=*/true); auto stream = at::cuda::getCurrentCUDAStream(logits[0].get_device()).stream(); diff --git a/cpp/tensorrt_llm/thop/reducescatterOp.cpp b/cpp/tensorrt_llm/thop/reducescatterOp.cpp index 05535089d26a..4f7157e381c1 100644 --- a/cpp/tensorrt_llm/thop/reducescatterOp.cpp +++ b/cpp/tensorrt_llm/thop/reducescatterOp.cpp @@ -55,16 +55,16 @@ class ReducescatterOp return 0; } - torch::Tensor run(torch::Tensor const& input, torch::optional> sizes) + std::vector run_list(torch::TensorList input_list, torch::optional> sizes) { TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used"); - auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); - auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); - std::vector outputShape = input.sizes().vec(); + bool use_nccl_reducescatter = !sizes.has_value() + || std::all_of(sizes.value().begin(), sizes.value().end(), + [&sizes](int64_t size) { return size == sizes.value()[0]; }); + int groupRank = 0; if (sizes.has_value()) { auto rank = COMM_SESSION.getRank(); - int groupRank = 0; for (auto const& currentRank : mGroup) { if (rank == currentRank) @@ -72,62 +72,52 @@ class ReducescatterOp ++groupRank; } TLLM_CHECK(static_cast(groupRank) < mGroup.size()); - outputShape[0] = sizes.value()[groupRank]; - } - else - { - outputShape[0] = outputShape[0] / mGroup.size(); } - auto output = torch::empty(outputShape, input.options()); - bool use_nccl_reducescatter = !sizes.has_value() - || std::all_of(sizes.value().begin(), sizes.value().end(), - [&sizes](int64_t size) { return size == sizes.value()[0]; }); - if (use_nccl_reducescatter) - { - NCCLCHECK_THROW(ncclReduceScatter(input.data_ptr(), output.mutable_data_ptr(), output.numel(), - (*getDtypeMap())[type], ncclSum, *mNcclComm, stream)); - } - else + std::vector output_list; + output_list.reserve(input_list.size()); + ncclGroupStart(); + for (auto const& input : input_list) { - size_t numel_base = std::accumulate(outputShape.cbegin() + 1, outputShape.cend(), 1, std::multiplies<>{}); - int64_t split_offset = 0; - ncclGroupStart(); - for (int root = 0; root < static_cast(mGroup.size()); ++root) + auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); + auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); + std::vector outputShape = input.sizes().vec(); + if (sizes.has_value()) { - auto split_size = sizes.value()[root]; - NCCLCHECK_THROW( + outputShape[0] = sizes.value()[groupRank]; + } + else + { + outputShape[0] = outputShape[0] / mGroup.size(); + } + auto output = torch::empty(outputShape, input.options()); + if (use_nccl_reducescatter) + { + ncclReduceScatter(input.data_ptr(), output.mutable_data_ptr(), output.numel(), (*getDtypeMap())[type], + ncclSum, *mNcclComm, stream); + } + else + { + size_t numel_base + = std::accumulate(outputShape.cbegin() + 1, outputShape.cend(), 1, std::multiplies<>{}); + int64_t split_offset = 0; + for (int root = 0; root < static_cast(mGroup.size()); ++root) + { + auto split_size = sizes.value()[root]; ncclReduce(input.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).data_ptr(), output.mutable_data_ptr(), numel_base * split_size, (*getDtypeMap())[type], ncclSum, root, - *mNcclComm, stream)); - split_offset += split_size; + *mNcclComm, stream); + split_offset += split_size; + } } - ncclGroupEnd(); + output_list.push_back(output); } - return output; + NCCLCHECK_THROW(ncclGroupEnd()); + return output_list; } - std::vector run_list( - torch::TensorList input_list, torch::optional> sizes) noexcept + torch::Tensor run(torch::Tensor const& input, torch::optional> sizes) { - std::vector output_list; - output_list.reserve(input_list.size()); - bool use_nccl_reducescatter = !sizes.has_value() - || std::all_of(sizes.value().begin(), sizes.value().end(), - [&sizes](int64_t size) { return size == sizes.value()[0]; }); - if (use_nccl_reducescatter) - { - ncclGroupStart(); - } - for (auto const& input : input_list) - { - auto output = run(input, sizes); - output_list.push_back(output); - } - if (use_nccl_reducescatter) - { - ncclGroupEnd(); - } - return output_list; + return run_list({input}, sizes)[0]; } private: diff --git a/cpp/tensorrt_llm/thop/weightOnlyQuantGemm.cpp b/cpp/tensorrt_llm/thop/weightOnlyQuantGemm.cpp new file mode 100644 index 000000000000..a00b51e16e41 --- /dev/null +++ b/cpp/tensorrt_llm/thop/weightOnlyQuantGemm.cpp @@ -0,0 +1,165 @@ +/* + * Copyright (c) 2025, 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 "weightOnlyQuantGemm.h" +#include "cutlass/numeric_types.h" + +#include +#include + +using namespace tensorrt_llm::kernels::cutlass_kernels; +using namespace tensorrt_llm::kernels; + +namespace torch_ext +{ + +namespace +{ +void check_input_dtypes(at::Tensor const& mat_a, at::Tensor const& mat_b) +{ + TORCH_CHECK(mat_a.scalar_type() == at::ScalarType::BFloat16 || mat_a.scalar_type() == at::ScalarType::Half, + "Activation matrix dtype must be BF16 or FP16"); + + TORCH_CHECK(mat_b.scalar_type() == at::ScalarType::Char, "Weight matrix dtype must be INT8"); +} + +#define DISPATCH_ACTIVATION_TYPE(scalar_type, ...) \ + if (scalar_type == at::ScalarType::Half) \ + { \ + using ActivationType = half; \ + __VA_ARGS__(); \ + } \ + else if (scalar_type == at::ScalarType::BFloat16) \ + { \ + using ActivationType = __nv_bfloat16; \ + __VA_ARGS__(); \ + } \ + else \ + { \ + TORCH_CHECK(false, "Unsupported activation type"); \ + } + +#define DISPATCH_WEIGHT_TYPE(scalar_type, ...) \ + if (scalar_type == at::ScalarType::Char) \ + { \ + using WeightType = uint8_t; \ + __VA_ARGS__(); \ + } \ + else if (scalar_type == at::ScalarType::QUInt4x2) \ + { \ + using WeightType = cutlass::uint4b_t; \ + __VA_ARGS__(); \ + } \ + else \ + { \ + TORCH_CHECK(false, "Unsupported weight type"); \ + } + +} // namespace + +WeightOnlyQuantGemmRunner::WeightOnlyQuantGemmRunner(at::ScalarType activation_dtype, at::ScalarType weight_dtype) + : mActivationDtype(activation_dtype) + , mWeightDtype(weight_dtype) +{ + DISPATCH_ACTIVATION_TYPE(activation_dtype, + [&] + { + using ADtypeStatic = ActivationType; + DISPATCH_WEIGHT_TYPE(weight_dtype, + [&] + { + using BDtypeStatic = WeightType; + mGemmRunner = std::make_shared>(); + }) + }) + mConfigs = mGemmRunner->getConfigs(); + TORCH_CHECK(!mConfigs.empty(), "Failed to get CUTLASS configs for WeightOnlyQuantGemmRunner with activation type ", + c10::toString(mActivationDtype), ", weight type ", c10::toString(mWeightDtype)); +} + +at::Tensor WeightOnlyQuantGemmRunner::runGemm(at::Tensor const& mat_a, at::Tensor const& mat_b, + at::Tensor const& weight_scales, int64_t config_idx, bool to_userbuffers, std::optional out_dtype) +{ + check_input_dtypes(mat_a, mat_b); + + TORCH_CHECK(mat_a.dim() == 2, "mat_a must be a matrix"); + TORCH_CHECK(mat_b.dim() == 2, "mat_b must be a matrix"); + TORCH_CHECK(mat_a.sizes()[1] == mat_b.sizes()[0], "mat_a and mat_b shapes cannot be multiplied"); + TORCH_CHECK(mat_a.is_cuda() && mat_b.is_cuda() && weight_scales.is_cuda(), "All input tensors must be on CUDA"); + + auto const m = mat_a.sizes()[0]; + auto const k = mat_a.sizes()[1]; + auto const n = mat_b.sizes()[1]; + auto real_n = n; + if (mWeightDtype == at::ScalarType::QUInt4x2) + { + real_n = n * 2; + } + + auto const dtype = out_dtype.value_or(mActivationDtype); + at::Tensor out; + if (to_userbuffers) + { + out = torch_ext::create_userbuffers_tensor({m, real_n}, dtype).first; + } + else + { + out = at::detail::empty_cuda({m, real_n}, dtype, mat_a.device(), std::nullopt); + } + + auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device()); + + auto workspace_size = mGemmRunner->getWorkspaceSize(m, real_n, k); + at::Tensor workspace; + char* workspace_ptr = nullptr; + if (workspace_size > 0) + { + workspace = at::detail::empty_cuda( + {static_cast(workspace_size)}, at::ScalarType::Byte, mat_a.device(), std::nullopt); + workspace_ptr = static_cast(workspace.data_ptr()); + } + + tensorrt_llm::cutlass_extensions::CutlassGemmConfig gemm_config_to_use; + if (config_idx >= 0 && config_idx < getNumConfigs()) + { + gemm_config_to_use = mConfigs.at(config_idx); + } + else + { + gemm_config_to_use = mConfigs.at(0); + } + + mGemmRunner->gemm(mat_a.data_ptr(), mat_b.data_ptr(), weight_scales.data_ptr(), out.data_ptr(), m, real_n, k, + gemm_config_to_use, workspace_ptr, workspace_size, stream); + + return out; +} + +int64_t WeightOnlyQuantGemmRunner::getNumConfigs() const +{ + TORCH_CHECK(mGemmRunner, "WeightOnlyQuantGemmRunner not initialized properly."); + return static_cast(mConfigs.size()); +} + +} // namespace torch_ext + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.class_("WeightOnlyQuantGemmRunner") + .def(torch::init()) + .def("run_gemm", &torch_ext::WeightOnlyQuantGemmRunner::runGemm) + .def("get_num_configs", &torch_ext::WeightOnlyQuantGemmRunner::getNumConfigs); +} diff --git a/cpp/tensorrt_llm/thop/weightOnlyQuantGemm.h b/cpp/tensorrt_llm/thop/weightOnlyQuantGemm.h new file mode 100644 index 000000000000..df062d79a52b --- /dev/null +++ b/cpp/tensorrt_llm/thop/weightOnlyQuantGemm.h @@ -0,0 +1,53 @@ +/* + * Copyright (c) 2025, 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 "cutlass_extensions/gemm_configs.h" +#include "cutlass_extensions/weight_only_quant_op.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.h" +#include "tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm.h" +#include "tensorrt_llm/thop/thUtils.h" +#include "tensorrt_llm/thop/userbuffersTensor.h" + +#include + +using namespace tensorrt_llm::kernels::cutlass_kernels; +using namespace tensorrt_llm::kernels; + +namespace torch_ext +{ +using WeightOnlyQuantGemmRunnerPtr = std::shared_ptr; + +class WeightOnlyQuantGemmRunner : public torch::CustomClassHolder +{ +public: + explicit WeightOnlyQuantGemmRunner(at::ScalarType activation_dtype, at::ScalarType weight_dtype); + + at::Tensor runGemm(at::Tensor const& mat_a, at::Tensor const& mat_b, at::Tensor const& weight_scales, + int64_t config_idx, bool to_userbuffers, std::optional out_dtype); + + int64_t getNumConfigs() const; + +private: + WeightOnlyQuantGemmRunnerPtr mGemmRunner; + at::ScalarType mActivationDtype; + at::ScalarType mWeightDtype; + std::vector mConfigs; +}; + +} // namespace torch_ext diff --git a/cpp/tests/batch_manager/cacheTransceiverTest.cpp b/cpp/tests/batch_manager/cacheTransceiverTest.cpp index 5c3bbcb1278d..99c40f810f67 100644 --- a/cpp/tests/batch_manager/cacheTransceiverTest.cpp +++ b/cpp/tests/batch_manager/cacheTransceiverTest.cpp @@ -186,8 +186,8 @@ texec::kv_cache::CommState MockDataSender::mState; class MockDataReceiver : public DataReceiver { public: - MOCK_METHOD(void, sendRequestInfo, (LlmRequest const&), (override)); - MOCK_METHOD(void, receiveSync, (LlmRequest const&), (override)); + MOCK_METHOD(TransferSession, sendRequestInfo, (LlmRequest const&), (override)); + MOCK_METHOD(void, receiveSync, (TransferSession&), (override)); }; class MockTransceiverTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) @@ -237,12 +237,14 @@ TEST_F(MockTransceiverTest, MpiRequesterBasic) GTEST_SKIP() << "mpirun with procs<=2 is required to run this test."; } auto receiver = std::make_unique(); - EXPECT_CALL(*receiver, sendRequestInfo).WillOnce(Return()); + auto state = std::make_unique(); + state->setCommState(texec::kv_cache::CommState{std::vector{0}}); + EXPECT_CALL(*receiver, sendRequestInfo) + .WillOnce(Return(TransferSession({nullptr}, DataContext{0}, *state, *state, + tensorrt_llm::runtime::BufferManager{std::make_shared()}, nullptr))); EXPECT_CALL(*receiver, receiveSync).WillOnce(Return()); DataRequester requester{std::move(receiver)}; auto request = makeLlmRequest(0); - auto state = std::make_unique(); - state->setCommState(texec::kv_cache::CommState{std::vector{0}}); auto stats = texec::ContextPhaseParams({}, 0, state.release(), std::nullopt); request->setContextPhaseParams(std::move(stats)); auto future = requester.requestAndReceiveAsync(*request); diff --git a/cpp/tests/batch_manager/guidedDecoderTest.cpp b/cpp/tests/batch_manager/guidedDecoderTest.cpp index 4b193ba3498f..8358e9873343 100644 --- a/cpp/tests/batch_manager/guidedDecoderTest.cpp +++ b/cpp/tests/batch_manager/guidedDecoderTest.cpp @@ -17,9 +17,9 @@ #include #include #include -#include #include "tensorrt_llm/batch_manager/common.h" +#include "tensorrt_llm/batch_manager/decoderBuffers.h" #include "tensorrt_llm/batch_manager/guidedDecoder.h" #include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/executor/executor.h" @@ -128,11 +128,21 @@ class GuidedDecoderTest : public ::testing::Test RequestVector contextRequests{llmReq1, llmReq2}; RequestVector generationRequests{}; ScheduledRequests scheduledRequests{contextRequests, generationRequests}; + DecoderInputBuffers decoderInputBuffers(mMaxNumRequests, 1, *mRuntimeBufferManager); + + for (auto const& requests : {scheduledRequests.contextRequests, scheduledRequests.generationRequests}) + { + for (auto const& llmReq : requests) + { + decoderInputBuffers.decoderRequests.push_back(llmReq); + } + } + decoderInputBuffers.logits = mLogits; // Context phase resetLogits(); mGuidedDecoder->build(scheduledRequests); - mGuidedDecoder->execute(scheduledRequests, *mRuntimeBufferManager, mLogits); + mGuidedDecoder->execute(decoderInputBuffers, *mRuntimeBufferManager); syncLogitsToHost(); mRuntimeBufferManager->getStream().synchronize(); @@ -143,8 +153,18 @@ class GuidedDecoderTest : public ::testing::Test generationRequests.push_back(llmReq1); llmReq2->setState(LlmRequestState::kGENERATION_IN_PROGRESS); generationRequests.push_back(llmReq2); - EXPECT_EQ(countRejected(1), mExpectedNumRejected[0]); - EXPECT_EQ(countRejected(2), 0); + + decoderInputBuffers.decoderRequests.clear(); + for (auto const& requests : {scheduledRequests.contextRequests, scheduledRequests.generationRequests}) + { + for (auto const& llmReq : requests) + { + decoderInputBuffers.decoderRequests.push_back(llmReq); + } + } + + EXPECT_EQ(countRejected(0), mExpectedNumRejected[0]); + EXPECT_EQ(countRejected(1), 0); // Generation phase for (int i = 0; i < mOutputIds.size(); i++) @@ -154,12 +174,12 @@ class GuidedDecoderTest : public ::testing::Test resetLogits(); mGuidedDecoder->build(scheduledRequests); - mGuidedDecoder->execute(scheduledRequests, *mRuntimeBufferManager, mLogits); + mGuidedDecoder->execute(decoderInputBuffers, *mRuntimeBufferManager); syncLogitsToHost(); mRuntimeBufferManager->getStream().synchronize(); - EXPECT_EQ(countRejected(1), mExpectedNumRejected[i + 1]); - EXPECT_EQ(countRejected(2), 0); + EXPECT_EQ(countRejected(0), mExpectedNumRejected[i + 1]); + EXPECT_EQ(countRejected(1), 0); } } diff --git a/cpp/tests/executor/disaggExecutorTest.cpp b/cpp/tests/executor/disaggExecutorTest.cpp index 49c8c00f0489..75ab6dccb444 100644 --- a/cpp/tests/executor/disaggExecutorTest.cpp +++ b/cpp/tests/executor/disaggExecutorTest.cpp @@ -662,6 +662,8 @@ TEST_P(DisaggParamsTest, DisaggTokenComparison) 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); @@ -894,6 +896,8 @@ TEST_P(DisaggOrchestratorParamsTest, DisaggTokenComparison) 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); @@ -994,6 +998,8 @@ TEST_P(ConditionalDisaggParamsTest, DisaggTokenComparison) 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); diff --git a/cpp/tests/runtime/gptDecoderBatchedTest.cpp b/cpp/tests/runtime/gptDecoderBatchedTest.cpp index e1a86e4479a6..7c152f48a9e8 100644 --- a/cpp/tests/runtime/gptDecoderBatchedTest.cpp +++ b/cpp/tests/runtime/gptDecoderBatchedTest.cpp @@ -322,7 +322,7 @@ void testDecoder(nvinfer1::DataType const dtype, std::vector& sa modelConfig, worldConfig, manager); // set up inputs and outputs - tb::DecoderInputBuffers inputBuffers(batchSize, batchSize, maxGeneratedTokensPerStep, manager); + tb::DecoderInputBuffers inputBuffers(batchSize, maxGeneratedTokensPerStep, manager); auto batchSlotsRange = BufferRange(*inputBuffers.setupBatchSlots); std::iota(batchSlotsRange.begin(), batchSlotsRange.end(), 0); @@ -456,7 +456,7 @@ void testDecoderWavefront(nvinfer1::DataType const dtype, std::vector @@ -110,8 +111,13 @@ TEST_F(CacheTransBufferTest, TestPreAllocBufferSize) size_t sendBufferCount = tensorrt_llm::common::getEnvParallelCacheSend() ? tensorrt_llm::common::getEnvKVCacheSendMaxConcurrenceNum() : 1; - size_t bufferSizeBytes = CacheTransBufferManager::preAllocBufferSize(maxNumTokens) - * kvCacheSizePerToken(4, 2, 64, CacheType::kSELFKONLY); + size_t cacheSizeBytesPerToken = kvCacheSizePerToken(4, 2, 64, CacheType::kSELFKONLY); + std::map cacheSizeBytesPerTokenPerWindow{ + {maxBlocksPerSeq * tokensPerBlock, cacheSizeBytesPerToken}}; + tensorrt_llm::executor::CacheTransceiverConfig cacheTransceiverConfig{ + tensorrt_llm::executor::CacheTransceiverConfig::BackendType::UCX, maxNumTokens}; + size_t bufferSizeBytes + = CacheTransBufferManager::preAllocBufferSize(cacheSizeBytesPerTokenPerWindow, cacheTransceiverConfig); auto bufferId = mTransBufferManager->assignBufferIndexForSend(); EXPECT_TRUE(bufferId.has_value()); EXPECT_EQ(bufferId.value(), 0); @@ -149,15 +155,18 @@ TEST_F(CacheTransBufferTest, TestPreAllocBufferSize2) size_t sendBufferCount = tensorrt_llm::common::getEnvParallelCacheSend() ? tensorrt_llm::common::getEnvKVCacheSendMaxConcurrenceNum() : 1; - size_t bufferSizeBytes = CacheTransBufferManager::preAllocBufferSize(maxNumTokens) - * kvCacheSizePerToken(4, 2, 64, CacheType::kSELF); + size_t cacheSizeBytesPerToken = kvCacheSizePerToken(4, 2, 64, CacheType::kSELF); + tensorrt_llm::executor::CacheTransceiverConfig cacheTransceiverConfig{ + tensorrt_llm::executor::CacheTransceiverConfig::BackendType::UCX, maxNumTokens}; + std::map cacheSizeBytesPerTokenPerWindow{ + {maxBlocksPerSeq * tokensPerBlock, cacheSizeBytesPerToken}}; + size_t bufferSizeBytes + = CacheTransBufferManager::preAllocBufferSize(cacheSizeBytesPerTokenPerWindow, cacheTransceiverConfig); auto bufferId = mTransBufferManager->assignBufferIndexForSend(); EXPECT_TRUE(bufferId.has_value()); EXPECT_EQ(bufferId.value(), 0); EXPECT_EQ(bufferSizeBytes, mTransBufferManager->getSendBuffer(bufferId)->getSizeInBytes() * (recvbufferCount + sendBufferCount)); - TLLM_LOG_INFO("bufferSizeBytes: %ld , getSizeINBytes: %ld", bufferSizeBytes, - mTransBufferManager->getSendBuffer(bufferId)->getSizeInBytes() * (recvbufferCount + sendBufferCount)); mTransBufferManager->freeBufferIndexForSend(bufferId); exit(testing::Test::HasFailure() ? 1 : 0); } diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index db502614141b..ba10a17b26db 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -1034,6 +1034,182 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); } +TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest) +{ + using VecTokenExtraIds = LlmRequest::VecTokenExtraIds; + + auto constexpr numLayers = 12; + auto constexpr numKvHeads = 6; + auto constexpr sizePerHead = 16; + auto constexpr tokensPerBlock = 4; + auto constexpr maxBlocksPerSeq = 4; + auto constexpr blocksInPrimaryPool = 16; + auto constexpr blocksInSecondaryPool = 0; + auto constexpr maxNumSequences = 8; + auto const stream = std::make_shared(); + auto constexpr onboardBlocks = true; + auto constexpr numReturnSequences = 1; + auto constexpr maxAttentionWindow = tokensPerBlock * maxBlocksPerSeq; + auto constexpr beamWidth = 1; + + auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; + + BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, + maxNumSequences, stream, maxAttentionWindow, beamWidth, + std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0, + onboardBlocks); + blockManager.allocatePools(false); + + EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); + EXPECT_EQ(blockManager.getMaxNumBlocks(), blocksInPrimaryPool); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); + + SizeType32 constexpr maxNewTokens{0}; + tr::SamplingConfig const samplingConfig{beamWidth}; + bool constexpr isStreaming{false}; + + // Create multimodal hash data (256-bit hash = 8 int32 values) + auto multimodalHashes = std::make_shared>>(std::vector>{ + {0x12345678, -0x6F543211, 0x11111111, 0x22222222, 0x33333333, 0x44444444, 0x55555555, 0x66666666} // Hash 1 + }); + auto multimodalPositions + = std::make_shared>(std::vector{2}); // Start at token 2 + auto multimodalLengths = std::make_shared>(std::vector{4}); // Length 4 tokens + // assume prompt id starts from 100 + auto inputTokens = std::make_shared(VecTokens{100, 101, 102, 103, 104, 105, 0, 1, 2}); + auto const inputLength = static_cast(inputTokens->size()); + LlmRequest::RequestIdType requestId{0}; + auto llmRequest0 = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming, + std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + multimodalHashes, multimodalPositions, multimodalLengths, 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, 0.5, std::nullopt, std::nullopt, + std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences); + + GenerationRequest seq0{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()}; + + /////////////////////////////////////////////////////////////////////////// + // add request and then remove it + auto constexpr beamIdx = 0; + auto promptLen0 = llmRequest0->getNumTokens(beamIdx); + auto numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock()); + blockManager.addSequence(seq0, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow); + EXPECT_EQ(llmRequest0->getContextCurrentPosition(), 0); + EXPECT_THAT(seq0.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({0, 1, 2})); + llmRequest0->addNewToken(3, beamIdx); + llmRequest0->addNewToken(4, beamIdx); + auto numTokens = llmRequest0->getNumTokens(beamIdx); + auto numBlocks = tc::ceilDiv(numTokens, tokensPerBlock); + EXPECT_EQ(numBlocks, 3); + EXPECT_EQ(blockManager.getNumAllocatedBlocks(), numBlocks); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks); + + // Input: [100, 101, 102, 103, 104, 105, 0, 1, 2] (9 tokens) + // Multimodal: starts at token 2, length 4 → [102, 103, 104, 105] + + // Block 0: [100, 101, 102, 103] ← Contains multimodal (102, 103) + // Block 1: [104, 105, 0, 1] ← Contains multimodal (104, 105) + // Block 2: [2, 3, 4] ← No multimodal + blockManager.releaseBlocks(seq0, llmRequest0); + EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); + + /////////////////////////////////////////////////////////////////////////// + // new request with same tokens and same multimodal hash - should reuse + requestId = 1; + auto llmRequest1 = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming, + std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + multimodalHashes, multimodalPositions, multimodalLengths, 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, 0.5, std::nullopt, std::nullopt, + std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences); + GenerationRequest seq1{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()}; + + // should reuse blocks 0, 1 and get new block 3 + auto promptLen1 = llmRequest1->getNumTokens(beamIdx); + auto numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock()); + blockManager.addSequence(seq1, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow); + EXPECT_EQ(llmRequest1->getContextCurrentPosition(), 2 * tokensPerBlock); + EXPECT_THAT(seq1.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({0, 1, 3})); + llmRequest1->addNewToken(3, beamIdx); + llmRequest1->addNewToken(4, beamIdx); + EXPECT_EQ(blockManager.getNumAllocatedBlocks(), numBlocks); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks); + // block 3 matches block 2 and will be freed + blockManager.releaseBlocks(seq1, llmRequest1); + EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); + + /////////////////////////////////////////////////////////////////////////// + // Test Case 2: Different multimodal hash + requestId = 2; + auto multimodalHashes2 + = std::make_shared>>(std::vector>{ + {0x45678123, 0x23456789, 0x34567890, 0x12121212, 0x56565656, 0x78787878, 0x54545454, 0x67676767} // Hash 2 + }); + auto multimodalPositions2 + = std::make_shared>(std::vector{2}); // Start at token 2 + auto multimodalLengths2 = std::make_shared>(std::vector{4}); // Length 4 tokens + auto llmRequest2 = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming, + std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + multimodalHashes2, multimodalPositions2, multimodalLengths2, 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, 0.5, std::nullopt, std::nullopt, + std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences); + + GenerationRequest seq2{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()}; + // no reuse, get new blocks 4, 5, 6 + auto promptLen2 = llmRequest2->getNumTokens(beamIdx); + auto numContextBlocks2 = tc::ceilDiv(promptLen2, blockManager.getTokensPerBlock()); + blockManager.addSequence(seq2, promptLen2, numContextBlocks2, *llmRequest2, maxAttentionWindow); + EXPECT_EQ(llmRequest2->getContextCurrentPosition(), 0); + EXPECT_THAT(seq2.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({4, 5, 6})); + llmRequest2->addNewToken(9, beamIdx); + numTokens = llmRequest2->getNumTokens(beamIdx); + numBlocks = tc::ceilDiv(numTokens, tokensPerBlock); + EXPECT_EQ(blockManager.getNumAllocatedBlocks(), numBlocks); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks); + + /////////////////////////////////////////////////////////////////////////// + // Test Case 3: Multiple multimodal hashes and partial reuse + requestId = 3; + auto multimodalHashes3 + = std::make_shared>>(std::vector>{ + {0x12345678, -0x6F543211, 0x11111111, 0x22222222, 0x33333333, 0x44444444, 0x55555555, 0x66666666}, // Hash 1 + {0x45678123, 0x23456789, 0x34567890, 0x12121212, 0x56565656, 0x78787878, 0x54545454, 0x67676767} // Hash 2 + }); + auto multimodalPositions3 + = std::make_shared>(std::vector{2, 4}); // Start at token 2 and 4 + auto multimodalLengths3 + = std::make_shared>(std::vector{2, 2}); // Length 2 tokens + + auto llmRequest3 = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming, + std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + multimodalHashes3, multimodalPositions3, multimodalLengths3, 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, 0.5, std::nullopt, std::nullopt, + std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences); + GenerationRequest seq3{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()}; + // reuse block 0, get new blocks 7, 8 + auto promptLen3 = llmRequest3->getNumTokens(beamIdx); + auto numContextBlocks3 = tc::ceilDiv(promptLen3, blockManager.getTokensPerBlock()); + blockManager.addSequence(seq3, promptLen3, numContextBlocks3, *llmRequest3, maxAttentionWindow); + EXPECT_EQ(llmRequest3->getContextCurrentPosition(), + tokensPerBlock); // only reuse block 0 [100, 101, 102, 103] with same hash/offset + EXPECT_THAT(seq3.getCacheBlockIds(maxAttentionWindow).at(beamIdx), ::testing::ElementsAreArray({0, 7, 8})); + llmRequest3->addNewToken(11, beamIdx); + numTokens = llmRequest3->getNumTokens(beamIdx); + numBlocks = tc::ceilDiv(numTokens, tokensPerBlock); + EXPECT_EQ(blockManager.getNumAllocatedBlocks(), numBlocks * 2); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks * 2); + + // clean up + blockManager.releaseBlocks(seq2, llmRequest2); + blockManager.releaseBlocks(seq3, llmRequest3); + EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); +} + TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) { // tc::Logger::getLogger()->setLevel(tc::Logger::Level::DEBUG); @@ -2115,7 +2291,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerLeafBlockWithDependentTest) EXPECT_THAT(cacheBlockIds0, ::testing::ElementsAreArray({0, 1, 2})); // Lower priority of middle block to prevent offloading - auto const blockManager = kvCacheManager.getBlockManager(); + auto const& blockManager = kvCacheManager.getBlockManager(); auto middleBlock = blockManager.getBlockById(cacheBlockIds0[1], maxAttentionWindow); middleBlock->setPriority(0); @@ -2215,7 +2391,7 @@ TEST_P(KVCacheManagerTest, DISABLED_KVCacheManagerAllocationTest) std::nullopt, nvinfer1::DataType::kHALF, sinkTokenLength, stream, std::nullopt, enableBlockReuse, onboardBlocks); - auto const blockManager = kvCacheManager.getBlockManager(); + auto const& blockManager = kvCacheManager.getBlockManager(); auto const& bufferManager = blockManager.getBufferManager(theOnlyWindowSize(kvCacheManager)); auto const memoryPoolUsedBefore = bufferManager.memoryPoolUsed(); kvCacheManager.allocatePools(useUvm); diff --git a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp index d29cf0350caf..18f7e6f5379e 100644 --- a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp +++ b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp @@ -785,8 +785,8 @@ TEST(SerializeUtilsTest, ExecutorConfig) texec::SpeculativeDecodingConfig(true), texec::GuidedDecodingConfig( texec::GuidedDecodingConfig::GuidedDecodingBackend::kXGRAMMAR, std::initializer_list{"eos"}), - std::vector{tensorrt_llm::executor::AdditionalModelOutput{"output_name"}}, texec::CacheTransceiverConfig(1024), - true, true, true); + std::vector{tensorrt_llm::executor::AdditionalModelOutput{"output_name"}}, + texec::CacheTransceiverConfig(std::nullopt, 1024), true, true, true); auto executorConfig2 = serializeDeserialize(executorConfig); EXPECT_EQ(executorConfig.getMaxBeamWidth(), executorConfig2.getMaxBeamWidth()); @@ -862,7 +862,9 @@ TEST(SerializeUtilsTest, MethodReturnType) TEST(SerializeUtilsTest, CacheTransceiverConfig) { - texec::CacheTransceiverConfig cacheTransceiverConfig(1024); + texec::CacheTransceiverConfig cacheTransceiverConfig( + tensorrt_llm::executor::CacheTransceiverConfig::BackendType::UCX, 1024); auto cacheTransceiverConfig2 = serializeDeserialize(cacheTransceiverConfig); - EXPECT_EQ(cacheTransceiverConfig.getMaxNumTokens(), cacheTransceiverConfig2.getMaxNumTokens()); + EXPECT_EQ(cacheTransceiverConfig.getBackendType(), cacheTransceiverConfig2.getBackendType()); + EXPECT_EQ(cacheTransceiverConfig.getMaxTokensInBuffer(), cacheTransceiverConfig2.getMaxTokensInBuffer()); } diff --git a/cpp/tests/unit_tests/executor/transferAgentTest.cpp b/cpp/tests/unit_tests/executor/transferAgentTest.cpp index 38f57204534c..c73d9a2140bd 100644 --- a/cpp/tests/unit_tests/executor/transferAgentTest.cpp +++ b/cpp/tests/unit_tests/executor/transferAgentTest.cpp @@ -228,7 +228,7 @@ TEST_F(TransferAgentTest, Connect) TEST_F(TransferAgentTest, SyncMessage) { - + constexpr std::size_t MAX_QUERY_TIMES = std::numeric_limits::max(); std::string const agent0{"agent0"}, agent1{"agent1"}; BaseAgentConfig config0{agent0, true}, config1{agent1, true}; auto nixlAgent0 = makeTransferAgent(config0); @@ -255,47 +255,46 @@ TEST_F(TransferAgentTest, SyncMessage) checked = nixlAgent0->checkRemoteDescs(agent1, regMem3.getDescs()); } while (!checked); auto syncMessage = std::string("agent_sync_message"); - nixlAgent0->notifySyncMessage(agent1, syncMessage); - TransferRequest writeReq{TransferOp::kWRITE, regMem0.getDescs(), regMem3.getDescs(), agent1}; + TransferRequest writeReq{TransferOp::kWRITE, regMem0.getDescs(), regMem3.getDescs(), agent1, syncMessage}; auto status = nixlAgent0->submitTransferRequests(writeReq); - status->wait(); auto notif = nixlAgent1->getNotifiedSyncMessages(); + for (std::size_t i = 0; i < MAX_QUERY_TIMES && notif.size() == 0; i++) + { + notif = nixlAgent1->getNotifiedSyncMessages(); + } + TLLM_CHECK(status->isCompleted()); TLLM_CHECK(notif.size() == 1); TLLM_CHECK(notif[agent0].size() == 1); TLLM_CHECK(notif[agent0][0] == syncMessage); + TLLM_CHECK(memory0 == memory1); std::string syncMessage2 = "two_agent_sync_message"; nixlAgent0->notifySyncMessage(agent1, syncMessage2); - while (true) + auto notif2 = nixlAgent1->getNotifiedSyncMessages(); + for (std::size_t i = 0; i < MAX_QUERY_TIMES && notif2.size() == 0; i++) { - auto notif2 = nixlAgent1->getNotifiedSyncMessages(); - if (notif2.size() > 0) - { - TLLM_CHECK(notif2.size() == 1); - TLLM_CHECK(notif2[agent0].size() == 1); - TLLM_CHECK(notif2[agent0][0] == syncMessage2); - break; - } + notif2 = nixlAgent1->getNotifiedSyncMessages(); } + TLLM_CHECK(notif2.size() == 1); + TLLM_CHECK(notif2[agent0].size() == 1); + TLLM_CHECK(notif2[agent0][0] == syncMessage2); // nixlAgent1->loadRemoteAgent(agent0); auto connectionInfo2 = nixlAgent0->getConnectionInfo(); nixlAgent1->connectRemoteAgent(agent0, connectionInfo2); std::string syncMessage3 = "three_agent_sync_message"; nixlAgent1->notifySyncMessage(agent0, syncMessage3); - while (true) + auto notif3 = nixlAgent0->getNotifiedSyncMessages(); + for (std::size_t i = 0; i < MAX_QUERY_TIMES && notif3.size() == 0; i++) { - auto notif3 = nixlAgent0->getNotifiedSyncMessages(); - if (notif3.size() > 0) - { - TLLM_CHECK(notif3.size() == 1); - TLLM_CHECK(notif3[agent1].size() == 1); - TLLM_CHECK(notif3[agent1][0] == syncMessage3); - break; - } + notif3 = nixlAgent0->getNotifiedSyncMessages(); } + TLLM_CHECK(notif3.size() == 1); + TLLM_CHECK(notif3[agent1].size() == 1); + TLLM_CHECK(notif3[agent1][0] == syncMessage3); + bool checked2 = false; do { @@ -303,11 +302,14 @@ TEST_F(TransferAgentTest, SyncMessage) } while (!checked2); std::string syncMessage4 = "four_agent_sync_message"; - nixlAgent1->notifySyncMessage(agent0, syncMessage4); - TransferRequest writeReq1{TransferOp::kWRITE, regMem2.getDescs(), regMem1.getDescs(), agent0}; + TransferRequest writeReq1{TransferOp::kWRITE, regMem2.getDescs(), regMem1.getDescs(), agent0, syncMessage4}; auto status1 = nixlAgent1->submitTransferRequests(writeReq1); - status1->wait(); auto notif4 = nixlAgent0->getNotifiedSyncMessages(); + for (std::size_t i = 0; i < MAX_QUERY_TIMES && notif4.size() == 0; i++) + { + notif4 = nixlAgent0->getNotifiedSyncMessages(); + } + TLLM_CHECK(status1->isCompleted()); TLLM_CHECK(notif4.size() == 1); TLLM_CHECK(notif4[agent1].size() == 1); TLLM_CHECK(notif4[agent1][0] == syncMessage4); @@ -322,20 +324,17 @@ TEST_F(TransferAgentTest, SyncMessage) Serialization::serialize(state, ss); std::string serializedState = ss.str(); nixlAgent0->notifySyncMessage(agent1, serializedState); - while (true) + auto notif5 = nixlAgent1->getNotifiedSyncMessages(); + for (size_t i = 0; i < MAX_QUERY_TIMES && notif5.size() == 0; i++) { - auto notif5 = nixlAgent1->getNotifiedSyncMessages(); - if (notif5.size() > 0) - { - TLLM_CHECK(notif5.size() == 1); - TLLM_CHECK(notif5[agent0].size() == 1); - TLLM_CHECK(notif5[agent0][0] == serializedState); - std::stringstream ss2(notif5[agent0][0]); - auto state2 = Serialization::deserializeCommState(ss2); - TLLM_CHECK(state2 == state); - break; - } + notif5 = nixlAgent1->getNotifiedSyncMessages(); } + TLLM_CHECK(notif5.size() == 1); + TLLM_CHECK(notif5[agent0].size() == 1); + TLLM_CHECK(notif5[agent0][0] == serializedState); + std::stringstream ss2(notif5[agent0][0]); + auto state2 = Serialization::deserializeCommState(ss2); + TLLM_CHECK(state2 == state); nixlAgent0->invalidateRemoteAgent(agent1); nixlAgent1->invalidateRemoteAgent(agent0); diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi index da67d0a49948..95aa670a090e 100644 --- a/docker/Dockerfile.multi +++ b/docker/Dockerfile.multi @@ -99,6 +99,10 @@ COPY --from=triton /opt/tritonserver/caches /opt/tritonserver/caches COPY docker/common/install_triton.sh install_triton.sh RUN bash ./install_triton.sh && rm install_triton.sh +# Install UCX first +COPY docker/common/install_ucx.sh install_ucx.sh +RUN bash ./install_ucx.sh && rm install_ucx.sh + # Install NIXL COPY docker/common/install_nixl.sh install_nixl.sh RUN bash ./install_nixl.sh && rm install_nixl.sh @@ -123,7 +127,7 @@ RUN mkdir -p /root/.cache/pip /root/.cache/ccache ENV CCACHE_DIR=/root/.cache/ccache # Build the TRT-LLM wheel ARG GITHUB_MIRROR="" -ARG BUILD_WHEEL_ARGS="--clean --python_bindings --benchmarks" +ARG BUILD_WHEEL_ARGS="--clean --benchmarks" RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=${CCACHE_DIR} \ GITHUB_MIRROR=$GITHUB_MIRROR python3 scripts/build_wheel.py ${BUILD_WHEEL_ARGS} diff --git a/docker/Dockerfile.user b/docker/Dockerfile.user index f6f912ddd3a7..3d1ff9e52fd0 100644 --- a/docker/Dockerfile.user +++ b/docker/Dockerfile.user @@ -21,3 +21,5 @@ RUN (getent group ${GROUP_ID} || groupadd --gid ${GROUP_ID} ${GROUP_NAME}) && \ rm -rf /var/lib/apt/lists/* USER ${USER_NAME} + +ENV PATH="/home/${USER_NAME}/.local/bin:$PATH" diff --git a/docker/Makefile b/docker/Makefile index 71d3b6511d5f..2b5022b1ee8e 100644 --- a/docker/Makefile +++ b/docker/Makefile @@ -152,11 +152,11 @@ endif $(if $(and $(filter 1,$(LOCAL_USER)),$(shell [ -w "$(USER_CACHE_DIR)" ] && echo 1)),--volume $(USER_CACHE_DIR):/home/$(USER_NAME)/.cache:rw) \ --env "CCACHE_DIR=$(CCACHE_DIR)" \ --env "CCACHE_BASEDIR=$(CODE_DIR)" \ - --env "CONAN_HOME=$(CONAN_DIR)" \ + --env "CONAN_HOME=$(CONAN_DIR)" \ --workdir $(WORK_DIR) \ --hostname $(shell hostname)-$* \ --name $(CONTAINER_NAME)-$*-$(USER_NAME) \ - --tmpfs /tmp:exec \ + --tmpfs /tmp:exec \ $(IMAGE_WITH_TAG)$(IMAGE_TAG_SUFFIX) $(RUN_CMD) devel_%: STAGE = devel @@ -180,7 +180,8 @@ jenkins-aarch64_%: IMAGE_WITH_TAG = $(shell . ../jenkins/current_image_tags.prop jenkins-aarch64_%: STAGE = tritondevel # For x86_64 -jenkins-rockylinux8_%: IMAGE_WITH_TAG = $(shell . ../jenkins/current_image_tags.properties && echo $$LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE) +jenkins-rockylinux8_%: PYTHON_VERSION_TAG_ID = $(if $(findstring 3.12,${PYTHON_VERSION}),PY312,$(if $(findstring 3.10,${PYTHON_VERSION}),PY310,$(error Unknown PYTHON_VERSION specified))) +jenkins-rockylinux8_%: IMAGE_WITH_TAG = $(shell . ../jenkins/current_image_tags.properties && echo $$LLM_ROCKYLINUX8_${PYTHON_VERSION_TAG_ID}_DOCKER_IMAGE) jenkins-rockylinux8_%: STAGE = tritondevel jenkins-rockylinux8_%: BASE_IMAGE = nvidia/cuda jenkins-rockylinux8_%: BASE_TAG = 12.9.0-devel-rockylinux8 diff --git a/docker/README.md b/docker/README.md index 3bfac62a2c41..fa1b80a9fd72 100644 --- a/docker/README.md +++ b/docker/README.md @@ -89,13 +89,10 @@ equivalent containers as [described above](#building-docker-images-with-gnu-make ### Jenkins Integration [`Makefile`](Makefile) has special targets for building, pushing and running the Docker build image used on Jenkins. -The full image name and tag is defined in [`L0_MergeRequest.groovy`](../jenkins/L0_MergeRequest.groovy). The `make` -system will parse this name as the value of `LLM_DOCKER_IMAGE`. To build and push a new Docker image for Jenkins, -define a new image name and tag in [`L0_MergeRequest.groovy`](../jenkins/L0_MergeRequest.groovy) and run +The full image names and tags are defined in [`current_image_tags.properties`](../jenkins/current_image_tags.properties). The `make` +system will parse the names/tags from this file. -```bash -make -C docker jenkins_push -``` +#### Running Start a new container using the same image as Jenkins using your local user account with @@ -134,6 +131,38 @@ make -C docker trtllm_run LOCAL_USER=1 DOCKER_PULL=1 The argument `DOCKER_PULL=1` instructs `make` to pull the latest version of the image before deploying it in the container. By default, the release images built in the above manner are tagged by their `git` branch name and may be frequently updated. +#### Building CI images + +To build and push a new Docker image for Jenkins, define new image names and tags in [`current_image_tags.properties`](../jenkins/current_image_tags.properties) and run + +```bash +# Commands assume an amd64 host +make -C docker jenkins_build +# +docker buildx create --name multi-builder +make -C docker jenkins-aarch64_build \ + DOCKER_BUILD_ARGS="--platform arm64 --builder=multi-builder" +# +# check jenkins/BuildDockerImage.groovy for current Python versions +make -C docker jenkins-rockylinux8_build PYTHON_VERSION=3.12.3 +make -C docker jenkins-rockylinux8_build PYTHON_VERSION=3.10.12 +``` + +The resulting images then need to be pushed: + +```bash +sh -c '. jenkins/current_image_tags.properties && echo $LLM_DOCKER_IMAGE $LLM_SBSA_DOCKER_IMAGE $LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE $LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE' | tr ' ' '\n' | xargs -I{} docker push {} +``` + +Alternatively, it is possible to trigger the image build by opening a new pull request and commenting + +```text +/bot run --stage-list "Build-Docker-Images" +``` + +The resulting images can then be re-tagged using `scripts/rename_docker_images.py` +and the new tags included in [`current_image_tags.properties`](../jenkins/current_image_tags.properties). + ### Docker rootless Some aspects require special treatment when using [Docker rootless mode](https://docs.docker.com/engine/security/rootless/). The `docker/Makefile` contains heuristics to detect Docker rootless mode. When assuming diff --git a/docker/common/install_nixl.sh b/docker/common/install_nixl.sh index 239434495132..18ee554f693a 100644 --- a/docker/common/install_nixl.sh +++ b/docker/common/install_nixl.sh @@ -2,56 +2,41 @@ set -ex GITHUB_URL="https://github.com" - -UCX_VERSION="v1.18.1" UCX_INSTALL_PATH="/usr/local/ucx/" CUDA_PATH="/usr/local/cuda" - -NIXL_VERSION="0.2.0" - -UCX_REPO="https://github.com/openucx/ucx.git" +NIXL_VERSION="0.3.1" NIXL_REPO="https://github.com/ai-dynamo/nixl.git" - -if [ ! -d ${UCX_INSTALL_PATH} ]; then - git clone --depth 1 -b ${UCX_VERSION} ${UCX_REPO} - cd ucx - ./autogen.sh - ./contrib/configure-release \ - --prefix=${UCX_INSTALL_PATH} \ - --enable-shared \ - --disable-static \ - --disable-doxygen-doc \ - --enable-optimizations \ - --enable-cma \ - --enable-devel-headers \ - --with-cuda=${CUDA_PATH} \ - --with-verbs \ - --with-dm \ - --enable-mt - make install -j$(nproc) - cd .. - rm -rf ucx # Remove UCX source to save space - echo "export LD_LIBRARY_PATH=${UCX_INSTALL_PATH}/lib:\$LD_LIBRARY_PATH" >> "${ENV}" -fi - ARCH_NAME="x86_64-linux-gnu" +GDS_PATH="$CUDA_PATH/targets/x86_64-linux" if [ "$(uname -m)" != "amd64" ] && [ "$(uname -m)" != "x86_64" ]; then ARCH_NAME="aarch64-linux-gnu" - EXTRA_NIXL_ARGS="-Ddisable_gds_backend=true" -fi - -if [ $ARCH_NAME != "x86_64-linux-gnu" ]; then - echo "The NIXL backend is temporarily unavailable on the aarch64 platform. Exiting script." - exit 0 + GDS_PATH="$CUDA_PATH/targets/sbsa-linux" fi pip3 install --no-cache-dir meson ninja pybind11 git clone --depth 1 -b ${NIXL_VERSION} ${NIXL_REPO} cd nixl -meson setup builddir -Ducx_path=${UCX_INSTALL_PATH} -Dstatic_plugins=UCX ${EXTRA_NIXL_ARGS} + +cuda_path=$(find / -name "libcuda.so.1" 2>/dev/null | head -n1) +if [[ -z "$cuda_path" ]]; then + echo "libcuda.so.1 not found " + exit 1 +fi + +ln -sf $cuda_path $CUDA_PATH/lib64/libcuda.so.1 + +meson setup builddir \ + -Ducx_path=$UCX_INSTALL_PATH \ + -Dcudapath_lib="$CUDA_PATH/lib64" \ + -Dcudapath_inc="$CUDA_PATH/include" \ + -Dgds_path="$GDS_PATH" \ + -Dinstall_headers=true \ + -Dstatic_plugins=UCX + cd builddir && ninja install cd ../.. rm -rf nixl* # Remove NIXL source tree to save space +rm $CUDA_PATH/lib64/libcuda.so.1 echo "export LD_LIBRARY_PATH=/opt/nvidia/nvda_nixl/lib/${ARCH_NAME}:/opt/nvidia/nvda_nixl/lib64:\$LD_LIBRARY_PATH" >> "${ENV}" diff --git a/docker/common/install_ucx.sh b/docker/common/install_ucx.sh new file mode 100644 index 000000000000..22f444d9746f --- /dev/null +++ b/docker/common/install_ucx.sh @@ -0,0 +1,30 @@ +#!/bin/bash +set -ex + +GITHUB_URL="https://github.com" +UCX_VERSION="v1.18.1" +UCX_INSTALL_PATH="/usr/local/ucx/" +CUDA_PATH="/usr/local/cuda" +UCX_REPO="https://github.com/openucx/ucx.git" + +if [ ! -d ${UCX_INSTALL_PATH} ]; then + git clone --depth 1 -b ${UCX_VERSION} ${UCX_REPO} + cd ucx + ./autogen.sh + ./contrib/configure-release \ + --prefix=${UCX_INSTALL_PATH} \ + --enable-shared \ + --disable-static \ + --disable-doxygen-doc \ + --enable-optimizations \ + --enable-cma \ + --enable-devel-headers \ + --with-cuda=${CUDA_PATH} \ + --with-verbs \ + --with-dm \ + --enable-mt + make install -j$(nproc) + cd .. + rm -rf ucx # Remove UCX source to save space + echo "export LD_LIBRARY_PATH=${UCX_INSTALL_PATH}/lib:\$LD_LIBRARY_PATH" >> "${ENV}" +fi diff --git a/docs/source/advanced/disaggregated-service.md b/docs/source/advanced/disaggregated-service.md index 757b1da81f43..e5c4a19ba4b7 100644 --- a/docs/source/advanced/disaggregated-service.md +++ b/docs/source/advanced/disaggregated-service.md @@ -16,8 +16,6 @@ An [architectural and performance overview](../../../docs/source/blogs/tech_blog TRT-LLM uses some environment variables to control the behavior of disaggregated service. -* `TRTLLM_USE_UCX_KVCACHE`: Specifies whether to use UCX for KV cache transfer. The default value is `0`. This must be enabled when using a disaggregated service. - * `TRTLLM_PARALLEL_CACHE_SEND`: If set to `1`, contextExecutor will attempt to send KV cache for multiple requests in parallel. The default value is `0`. * `TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP`: If set to `1`, generationExecutor will not overlap KV cache transfer with model inference. The default value is `0`. @@ -34,6 +32,10 @@ TRT-LLM uses some environment variables to control the behavior of disaggregated * `TRTLLM_KVCACHE_SEND_MAX_CONCURRENCY_NUM`: The maximum number of concurrent KV cache sends. The default value is `4`. This environment variable only takes effect when `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE` is greater than 0. +There are some other useful environment variables that may help when encountering failures or performance issues. + +* `NCCL_GRAPH_MIXING_SUPPORT`: With the default value `1`, the CUDA driver may create too many CUDA streams while working with one CUDA graph, leading to performance drop. Setting it to `0` will reduce the number of CUDA streams, but please make sure there are no other NCCL ops outside the one CUDA graph, otherwise it's unsafe. + ## Troubleshooting and FAQ ### General FAQs @@ -66,55 +68,29 @@ A. Yes, it's recommended that different executor use different GPUs . We support *Q. How to handle error `Disaggregated serving is not enabled, please check the configuration?`* -A. Please set the environment variables -``` -export TRTLLM_USE_UCX_KVCACHE=1 -``` +A. please set `backendType` of `CacheTransceiverConfig`. +```cpp +ExecutorConfig executorConfig{...}; -*Q. Why do some profiling tools show that TRT-LLM's KV cache transfer does not utilize NVLink even on devices equipped with NVLink?* +executorConfig.setCacheTransceiverConfig(texec::CacheTransceiverConfig(BackendType::DEFAULT)); +``` -A. Please check version of `UCX` with `ucx_info -v`. -If the version of UCX <=1.17, set the environment variables `UCX_RNDV_FRAG_MEM_TYPE=cuda` and `UCX_MEMTYPE_CACHE=n` to enable NVLink. For BlackWell architecture GPUs, UCX version >=1.19 is required to enable NVLink. -If the version of UCX >=1.18, there are several ways to enable NVLink: -1. Set the environment variables `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=0B`,`UCX_CUDA_COPY_ASYNC_MEM_TYPE=cuda`, `UCX_CUDA_COPY_DMABUF=no`, `UCX_MEMTYPE_CACHE=n` and `UCX_RNDV_PIPELINE_ERROR_HANDLING=y`. -2. Set the environment variables `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=$Size`, `UCX_MEMTYPE_CACHE=n` and `UCX_RNDV_PIPELINE_ERROR_HANDLING=y`. $Size represents the size of the buffer for KV cache transfer, which is recommended to be larger than the size of the KV cache for the longest request. +When the environment variable `TRTLLM_USE_MPI_KVCACHE=1` is set, TRT-LLM will transfer the KV cache using `CUDA-aware MPI`. All executor processes involved must share the same MPI world communicator. Consequently, with `TRTLLM_USE_MPI_KVCACHE=1`, TRT-LLM only supports launching multiple executors via `MPI`. Additionally, the `CommunicationMode` for the executors must be set to `kLEADER` or `kORCHESTRATOR` with `SpawnProcesses=false` for the `disaggregated-service`. These restrictions do not apply when `TRTLLM_USE_UCX_KVCACHE=1` is set. *Q. Does TRT-LLM support using GPU direct RDMA for inter-node KV Cache transfer?* -A. Yes, TRT-LLM supports using GPU direct RDMA for inter-node KV cache transfer, but it is not enabled by default. There are several ways to enable GPU direct RDMA: -1. Set the environment variables `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=0B`,`UCX_RNDV_FRAG_MEM_TYPE=cuda`, `UCX_MEMTYPE_CACHE=n` and `UCX_RNDV_PIPELINE_ERROR_HANDLING=y`. -2. Set the environment variables `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=$Size`, `UCX_MEMTYPE_CACHE=n` and `UCX_RNDV_PIPELINE_ERROR_HANDLING=y`, $Size represents the size of the buffer for KV cache transfer, which is recommended to be larger than the size of the KV cache for the longest request. +A. Yes, TRT-LLM supports using GPU direct RDMA for inter-node KV cache transfer. -*Q. Are there any guidelines for performance tuning of KV cache transfer?* +*Q. What causes the substantial bandwidth fluctuations in kvCache transfers, especially during the first few requests following service initialization?* -A. Depending on the user's use case, certain sets of environment variables can help avoid poor KV cache transfer performance. +A. The communication for kvCache transfer between executors are established dynamically. The connection establishment process incurs significant overhead, which explains the apparently lower kvCache transfer bandwidth observed during the initial requests after service startup. This lower bandwidth reflects the inclusion of connection establishment overhead. When conducting benchmarks, it is recommended to perform a warm-up phase to ensure accurate performance measurements. -Environment Variable Set A +*Q. When my servers are running on different NVLink domains, some servers hang or have a lower performance. How to fix that? -``` -export TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=0B -export UCX_RNDV_FRAG_MEM_TYPES=cuda -export UCX_MEMTYPE_CACHE=n -export UCX_RNDV_PIPELINE_ERROR_HANDLING=y -``` -This set allows KV cache transfers to utilize NVLink within nodes and GDRDMA between nodes. - -Environment Variable Set B +A. NVLink domain can be found with `nvidia-smi -q` in the `Fabric.ClusterUUID` field. A few UCX environment variables can be adjusted when your servers have different NVLink domains: -``` -export TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=0B -export UCX_CUDA_COPY_ASYNC_MEM_TYPE=cuda -export UCX_CUDA_COPY_DMABUF=no -export UCX_MEMTYPE_CACHE=n -export UCX_RNDV_PIPELINE_ERROR_HANDLING=y -``` -Set B may provide slightly better performance on a single node compared to Set A. However, when transferring KV cache across multiple nodes, it may cause program instability. +* `UCX_CUDA_IPC_ENABLE_MNNVL`: Set to `n`. This also can reduce UCX timeout error messages like `UCX ERROR cuMemImportFromShareableHandle failed: invalid resource handle`, although these errors don't necessarily cause your trtllm-serve to fail. -Environment Variable Set C +* `UCX_NET_DEVICES`: Check if this is set correctly, or unset this variable to allow UCX to use all possible devices. -``` -export TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=$Size -export UCX_MEMTYPE_CACHE=n -export UCX_RNDV_PIPELINE_ERROR_HANDLING=y -``` -Set C can achieve better performance than Sets A and B, both within and between nodes. However, if the KV cache size exceeds the specified $Size, performance may degrade. +* `UCX_RNDV_SCHEME`: Set to `get_zcopy` or `put_zcopy` on GB200 for better performance. The default value is `auto`. diff --git a/docs/source/advanced/speculative-decoding.md b/docs/source/advanced/speculative-decoding.md index a601d9dd24aa..85a87ae0624d 100644 --- a/docs/source/advanced/speculative-decoding.md +++ b/docs/source/advanced/speculative-decoding.md @@ -3,13 +3,14 @@ - [About Speculative Sampling](#about-speculative-sampling) - [Performance Improvements](#Performance-improvements) - [Draft-Target-Model](#Draft-Target-Model) -- [Prompt-Lookup-Decoding](#prompt-lookup-decoding) +- [NGram](#ngram) - [Medusa](#medusa) - [Medusa Tree](#medusa-tree) - [Using Medusa with TensorRT-LLM](#using-medusa-with-tensorrt-llm) - [Limitations](#limitations) - [ReDrafter](#redrafter) - [EAGLE](#eagle) + - [Disaggregated Serving](#disaggregated-serving) - [Lookahead decoding](#lookahead-decoding) ## About Speculative Sampling @@ -35,7 +36,7 @@ TensorRT-LLM supports several approaches for generating draft tokens, including: 1. [Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads paper](https://arxiv.org/abs/2401.10774). 2. [Recurrent Drafter for Fast Speculative Decoding in Large Language Models](https://arxiv.org/html/2403.09919v1). 3. [EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty](https://arxiv.org/pdf/2401.15077). -3. Utilizing prompt tokens as draft tokens. For more information, refer to [Prompt Lookup Decoding](https://github.com/apoorvumang/prompt-lookup-decoding/). +3. Utilizing prompt tokens as draft tokens. For more information, refer to [NGram](https://github.com/apoorvumang/prompt-lookup-decoding/). 4. Utilizing Jacobi-like decoding to predict and verify draft tokens using the same model which does not need additional fine-tuning. Refer to [Break the Sequential Dependency of LLM Inference Using Lookahead Decoding](https://arxiv.org/pdf/2402.02057). @@ -61,13 +62,13 @@ Subsequently, the prompt, now updated with the accepted tokens, is sent back to This iterative process continues until a predefined stop conditions are met. An example of this orchestration process can be found in the [TensorRT-LLM Triton backend](https://github.com/triton-inference-server/tensorrtllm_backend/blob/main/inflight_batcher_llm/client/e2e_grpc_speculative_decoding_client.py). -We provide two styles of running Draft-Target-Model now: using TensorRT-LLM-BLS in Triton Inference Server, or using TensorRT-LLM directly. Detailed steps of running can be found in [examples/draft_target_model/README.md](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/draft_target_model/README.md) and the code can be found in [examples/prompt_lookup/run_dtm_pld.py](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/prompt_lookup/run_dtm_pld.py). +We provide two styles of running Draft-Target-Model now: using TensorRT-LLM-BLS in Triton Inference Server, or using TensorRT-LLM directly. Detailed steps of running can be found in [examples/draft_target_model/README.md](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/draft_target_model/README.md) and the code can be found in [examples/ngram/run_dtm_ngram.py](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/ngram/run_dtm_ngram.py). -## Prompt-Lookup-Decoding +## NGram -The Prompt-Lookup speculative decoding directly copies from the input prompt and previous generated output as draft tokens while generating the later output. It works like Draft-Target-Model but involves only one Target LLM model without further fine-tuning. The Prompt-Lookup profit from the scenarios which have high n-gram overlap between input prompt and output, such as summarization, document QA, multi-turn chat, code editing, etc. +The NGram speculative decoding directly copies from the input prompt and previous generated output as draft tokens while generating the later output. It works like Draft-Target-Model but involves only one Target LLM model without further fine-tuning. The NGram profit from the scenarios which have high n-gram overlap between input prompt and output, such as summarization, document QA, multi-turn chat, code editing, etc. -See document in [examples/prompt_lookup/README.md](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/prompt_lookup/README.md) and the code can be found in [examples/prompt_lookup/run_dtm_pld.py](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/prompt_lookup/run_dtm_pld.py). +See document in [examples/ngram/README.md](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/ngram/README.md) and the code can be found in [examples/ngram/run_dtm_ngram.py](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/ngram/run_dtm_ngram.py). ## Medusa @@ -169,6 +170,10 @@ The EAGLE approach enhances the single-model Medusa method by predicting and ver Similarly to ReDrafter, TensorRT-LLM implements the EAGLE model such that logits prediction, draft tokens acceptance and draft token generation are performed inside of the TensorRT engine. EAGLE-1 and EAGLE-2 are both supported, while EAGLE-2 is currently in the experimental stage. Please, visit the [EAGLE README](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/eagle/README.md) for information about building and running the model. +### Disaggregated Serving + +[Disaggregated Serving](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/advanced/disaggregated-service.md) with EAGLE3 using the two model approach is supported in the Pytorch backend. Please refer to the following [Dynamo example](https://github.com/ai-dynamo/dynamo/blob/main/examples/tensorrt_llm/llama4_plus_eagle.md) on how to run EAGLE3 with Disaggregated Serving for Llama 4 Maverick. + ## Lookahead Decoding Lookahead decoding algorithm operates through two parallel computation branches within the same model: a lookahead branch that generates n-grams using a fixed-sized 2D window, and a verification branch that validates promising n-gram candidates. This approach eliminates the necessity for additional model training or fine-tuning and can be enabled for any autoregressive model. Refer to the [Lookahead decoding README](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/lookahead/README.md) for information about building and running the model. 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 98c72e700d6d..f13ef7315135 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 @@ -137,7 +137,6 @@ To do the benchmark, run the following command: YOUR_DATA_PATH= cat >./extra-llm-api-config.yml< cat >./extra-llm-api-config.yml< ./extra_llm_api_options.yaml < ./extra_llm_api_options_eplb.yaml </tensorrt_llm:main sh \ - -c "echo -e 'enable_attention_dp: false\nenable_min_latency: true\nenable_autotuner: false\ncuda_graph_config:\n max_batch_size: 8\nspeculative_config:\n decoding_type: Eagle\n max_draft_len: 3\n speculative_model_dir: /config/models/eagle\nkv_cache_config:\n enable_block_reuse: false' > c.yaml && \ + -c "echo -e 'enable_autotuner: false\nenable_attention_dp: false\nenable_min_latency: true\ncuda_graph_config:\n max_batch_size: 8\nspeculative_config:\n decoding_type: Eagle\n max_draft_len: 3\n speculative_model_dir: /config/models/eagle\n eagle3_one_model: true\nkv_cache_config:\n enable_block_reuse: false' > c.yaml && \ TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL=True \ trtllm-serve /config/models/maverick \ --host 0.0.0.0 --port 8000 \ @@ -141,7 +141,9 @@ docker kill ## Performance Tuning -The configuration provided is optimized for 8xB200 GPUs, but you can adjust several parameters for your specific workload: +The configuration provided is optimized for 8xB200 GPUs, but you can adjust several parameters for your specific workload. + +**Note:** This configuration is optimized for minimum latency (`enable_min_latency: true`). When increasing the concurrency of requests, the tokens per second (TPS) per user degrades rapidly. This setup is designed to maximize single-user performance rather than high-concurrency throughput. For workloads with many concurrent users, you may need to adjust the configuration accordingly. - `max_batch_size`: Controls how many requests can be batched together - `max_draft_len`: The number of tokens Eagle can speculate ahead diff --git a/docs/source/commands/trtllm-bench.rst b/docs/source/commands/trtllm-bench.rst new file mode 100644 index 000000000000..7f03c8dfc660 --- /dev/null +++ b/docs/source/commands/trtllm-bench.rst @@ -0,0 +1,164 @@ +trtllm-bench +=========================== + +trtllm-bench is a comprehensive benchmarking tool for TensorRT-LLM engines. It provides three main subcommands for different benchmarking scenarios: + +**Common Options for All Commands:** + +**Usage:** + +.. click:: tensorrt_llm.commands.bench:main + :prog: trtllm-bench + :nested: full + :commands: throughput, latency, build + + + +prepare_dataset.py +=========================== + +trtllm-bench is designed to work with the `prepare_dataset.py `_ script, which generates benchmark datasets in the required format. The prepare_dataset script supports: + +**Dataset Types:** + +- Real datasets from various sources +- Synthetic datasets with normal or uniform token distributions +- LoRA task-specific datasets + +**Key Features:** + +- Tokenizer integration for proper text preprocessing +- Configurable random seeds for reproducible results +- Support for LoRA adapters and task IDs +- Output in JSON format compatible with trtllm-bench + +.. important:: + The ``--stdout`` flag is **required** when using prepare_dataset.py with trtllm-bench to ensure proper data streaming format. + +**Usage:** + +prepare_dataset +------------------- + +.. code-block:: bash + + python prepare_dataset.py [OPTIONS] + +**Options** + +---- + +.. list-table:: + :widths: 20 80 + :header-rows: 1 + + * - Option + - Description + * - ``--tokenizer`` + - Tokenizer directory or HuggingFace model name (required) + * - ``--output`` + - Output JSON filename (default: preprocessed_dataset.json) + * - ``--stdout`` + - Print output to stdout with JSON dataset entry on each line (**required for trtllm-bench**) + * - ``--random-seed`` + - Random seed for token generation (default: 420) + * - ``--task-id`` + - LoRA task ID (default: -1) + * - ``--rand-task-id`` + - Random LoRA task range (two integers) + * - ``--lora-dir`` + - Directory containing LoRA adapters + * - ``--log-level`` + - Logging level: info or debug (default: info) + +dataset +------------------- + +Process real datasets from various sources. + +.. code-block:: bash + + python prepare_dataset.py dataset [OPTIONS] + +**Options** + +---- + +.. list-table:: + :widths: 20 80 + :header-rows: 1 + + * - Option + - Description + * - ``--input`` + - Input dataset file or directory (required) + * - ``--max-input-length`` + - Maximum input sequence length (default: 2048) + * - ``--max-output-length`` + - Maximum output sequence length (default: 512) + * - ``--num-samples`` + - Number of samples to process (default: all) + * - ``--format`` + - Input format: json, jsonl, csv, or txt (default: auto-detect) + + +token_norm_dist +------------------- + +Generate synthetic datasets with normal token distribution. + +.. code-block:: bash + + python prepare_dataset.py token_norm_dist [OPTIONS] + +**Options** + +---- + +.. list-table:: + :widths: 20 80 + :header-rows: 1 + + * - Option + - Description + * - ``--num-requests`` + - Number of requests to be generated (required) + * - ``--input-mean`` + - Normal distribution mean for input tokens (required) + * - ``--input-stdev`` + - Normal distribution standard deviation for input tokens (required) + * - ``--output-mean`` + - Normal distribution mean for output tokens (required) + * - ``--output-stdev`` + - Normal distribution standard deviation for output tokens (required) + + +token_unif_dist +------------------- + +Generate synthetic datasets with uniform token distribution + +.. code-block:: bash + + python prepare_dataset.py token_unif_dist [OPTIONS] + +**Options** + +---- + +.. list-table:: + :widths: 20 80 + :header-rows: 1 + + * - Option + - Description + * - ``--num-requests`` + - Number of requests to be generated (required) + * - ``--input-min`` + - Uniform distribution minimum for input tokens (required) + * - ``--input-max`` + - Uniform distribution maximum for input tokens (required) + * - ``--output-min`` + - Uniform distribution minimum for output tokens (required) + * - ``--output-max`` + - Uniform distribution maximum for output tokens (required) diff --git a/docs/source/commands/trtllm-serve.rst b/docs/source/commands/trtllm-serve.rst index ab7a67673009..ff9a7d07ece4 100644 --- a/docs/source/commands/trtllm-serve.rst +++ b/docs/source/commands/trtllm-serve.rst @@ -67,9 +67,14 @@ Another example uses ``curl``: :linenos: Multimodal Serving -~~~~~~~~~~~~~~~~~ +~~~~~~~~~~~~~~~~~~ -For multimodal models (e.g., Qwen2-VL), you'll need to create a configuration file and start the server with additional options: +For multimodal models, you need to create a configuration file and start the server with additional options due to the following limitations: + +* TRT-LLM multimodal is currently not compatible with ``kv_cache_reuse`` +* Multimodal models require ``chat_template``, so only the Chat API is supported + +To set up multimodal models: First, create a configuration file: @@ -78,7 +83,6 @@ First, create a configuration file: cat >./extra-llm-api-config.yml<`__ + for implementation details. + +**Video** + +* Using "video_url": + + .. code-block:: json + + {"role": "user", "content": [ + {"type": "text", "text": "What's in this video?"}, + {"type": "video_url", "video_url": {"url": "https://example.com/video.mp4"}} + ]} + +**Audio** + +* Using "audio_url": + + .. code-block:: json + + {"role": "user", "content": [ + {"type": "text", "text": "What's in this audio?"}, + {"type": "audio_url", "audio_url": {"url": "https://example.com/audio.mp3"}} + ]} + + Benchmark --------- diff --git a/docs/source/index.rst b/docs/source/index.rst index b63ec95a676c..50b9c1226786 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -77,6 +77,7 @@ Welcome to TensorRT-LLM's Documentation! :caption: Command-Line Reference :hidden: + commands/trtllm-bench commands/trtllm-build commands/trtllm-serve diff --git a/docs/source/installation/linux.md b/docs/source/installation/linux.md index 6f1383f3ef85..9bccba451c7f 100644 --- a/docs/source/installation/linux.md +++ b/docs/source/installation/linux.md @@ -32,6 +32,7 @@ ```bash pip3 install --upgrade pip setuptools && pip3 install tensorrt_llm ``` + **This project will download and install additional third-party open source software projects. Review the license terms of these open source projects before use.** 2. Sanity check the installation by running the following in Python (tested on Python 3.12): diff --git a/docs/source/performance/perf-overview.md b/docs/source/performance/perf-overview.md index 3f55a4e1095d..9e316617186b 100644 --- a/docs/source/performance/perf-overview.md +++ b/docs/source/performance/perf-overview.md @@ -28,101 +28,119 @@ nvidia/Llama-3.1-405B-Instruct-FP4 ``` #### Llama 3.3 70B FP4 + | | GPU | B200 | | | | -|:-----------------------------|:---|:----------|:----------|:----------|:----------| -| | TP Size | 1 | 2 | 4 | 8 | -| ISL, OSL| | | | | | -| | | | | | | -| 128, 128 | | 11,253.28 | 17,867.66 | 24,944.50 | 27,471.49 | -| 128, 2048 | | 9,925.00 | 15,459.71 | 23,608.58 | 30,742.86 | -| 128, 4096 | | 6,318.92 | 8,711.88 | 17,659.74 | 24,947.05 | -| 500, 2000 | | 7,559.88 | 10,602.27 | 20,910.23 | 28,182.34 | -| 1000, 1000 | | 6,866.96 | 10,838.01 | 16,567.86 | 19,991.64 | -| 1000, 2000 | | 6,736.88 | 9,132.08 | 15,737.02 | 20,518.04 | -| 1024, 2048 | | 6,580.56 | 8,767.45 | 15,722.55 | 20,437.96 | -| 2048, 128 | | 1,375.49 | 1,610.69 | 2,707.58 | 3,717.82 | -| 2048, 2048 | | 4,544.73 | 6,956.14 | 12,292.23 | 15,661.22 | -| 5000, 500 | | 1,488.19 | 2,379.73 | 3,588.45 | 4,810.21 | -| 20000, 2000 | | 580.96 | 1,043.58 | 1,957.84 | 3,167.30 | +|:------------------------|:--------|:----------|:----------|:----------|:----------| +| | TP Size | 1 | 2 | 4 | 8 | +| ISL, OSL | | | | | | +| | | | | | | +| 128, 128 | | 10,994.48 | 17,542.11 | 24,667.31 | 27,272.27 | +| 128, 2048 | | 9,580.46 | 15,432.35 | 23,568.12 | 31,174.31 | +| 128, 4096 | | 6,418.39 | 9,841.53 | 17,808.76 | 25,229.25 | +| 500, 2000 | | 7,343.32 | 11,850.57 | 20,709.67 | 28,038.78 | +| 1000, 1000 | | 6,752.53 | 10,815.88 | 16,413.04 | 20,060.66 | +| 1000, 2000 | | 6,670.07 | 9,830.73 | 15,597.49 | 20,672.37 | +| 1024, 2048 | | 6,636.75 | 9,807.13 | 15,519.23 | 20,617.28 | +| 2048, 128 | | 1,342.17 | 1,989.41 | 3,033.14 | 4,035.64 | +| 5000, 500 | | 1,429.67 | 2,419.67 | 3,686.84 | 5,182.96 | +| 20000, 2000 | | 629.77 | 1,177.01 | 2,120.66 | 3,429.03 | #### Llama 3.1 405B FP4 -| | GPU | B200 | -|:-----------------------------|:---|:----------| -| | TP Size | 8 | -| ISL, OSL| | | -| | | | -| 128, 128 | | 9,184.83 | -| 128, 2048 | | 10,387.23 | -| 128, 4096 | | 8,741.80 | -| 500, 2000 | | 9,242.34 | -| 1000, 1000 | | 7,565.50 | -| 1000, 2000 | | 7,696.76 | -| 1024, 2048 | | 7,568.93 | -| 2048, 128 | | 953.57 | -| 2048, 2048 | | 6,092.32 | -| 5000, 500 | | 1,332.22 | -| 20000, 2000 | | 961.58 | + +| | GPU | B200 | | +|:------------------------|:------- |:---------|:----------| +| | TP Size | 4 | 8 | +| ISL, OSL | | | | +| | | | | +| 128, 128 | | 6,163.81 | 9,002.90 | +| 128, 2048 | | 7,081.21 | 10,288.28 | +| 128, 4096 | | 6,028.37 | 8,713.77 | +| 500, 2000 | | 5,858.75 | 9,125.86 | +| 1000, 1000 | | 4,848.00 | 7,582.97 | +| 1000, 2000 | | 5,375.25 | 7,626.28 | +| 1024, 2048 | | 5,345.70 | 7,464.03 | +| 2048, 128 | | 693.55 | 1,086.56 | +| 5000, 500 | | 947.49 | 1,532.45 | +| 20000, 2000 | | 641.11 | 1,097.84 | ### FP8 Models: ``` nvidia/Llama-3.1-8B-Instruct-FP8 -nvidia/Llama-3.1-70B-Instruct-FP8 +nvidia/Llama-3.3-70B-Instruct-FP8 nvidia/Llama-3.1-405B-Instruct-FP8 +nvidia/Llama-4-Maverick-17B-128E-Instruct-FP8 ``` #### Llama 3.1 8B FP8 -| | GPU | H200 141GB HBM3 | H100 80GB HBM3 | + +| | GPU | H200 141GB HBM3 | H100 80GB HBM3 | |:-----------------------------|:---|:------------------|:-----------------| -| | TP Size | 1 | 1 | +| | TP Size | 1 | 1 | | ISL, OSL | | | | | | | | | -| 128, 128 | | 28,447.38 | 27,568.68 | -| 128, 2048 | | 23,294.74 | 22,003.62 | -| 128, 4096 | | 17,481.48 | 13,640.35 | -| 500, 2000 | | 21,462.57 | 17,794.39 | -| 1000, 1000 | | 17,590.60 | 15,270.02 | -| 1000, 2000 | | 17,139.51 | 13,850.22 | -| 1024, 2048 | | 16,970.63 | 13,374.15 | -| 2048, 128 | | 3,531.33 | 3,495.05 | -| 2048, 2048 | | 12,022.38 | 9,653.67 | -| 5000, 500 | | 3,851.65 | 3,371.16 | -| 20000, 2000 | | 1,706.06 | 1,340.92 | - -#### Llama 3.1 70B FP8 -| | GPU | H200 141GB HBM3 | | | | H100 80GB HBM3 | | | | +| 128, 128 | | 27,970.14 | 27,688.36 | +| 128, 2048 | | 23,326.38 | 21,841.15 | +| 128, 4096 | | 17,508.51 | 13,730.89 | +| 500, 2000 | | 21,390.41 | 17,833.34 | +| 1000, 1000 | | 17,366.89 | 15,270.62 | +| 1000, 2000 | | 16,831.31 | 13,798.08 | +| 1024, 2048 | | 16,737.03 | 13,385.50 | +| 2048, 128 | | 3,488.03 | 3,414.67 | +| 5000, 500 | | 3,813.69 | 3,394.54 | +| 20000, 2000 | | 1,696.66 | 1,345.42 | + +#### Llama 3.3 70B FP8 + +| | GPU | H200 141GB HBM3 | | | | H100 80GB HBM3 | | | | |:-----------------------------|:---|:------------------|:---------|:----------|:----------|:-----------------|:---------|:----------|:----------| -| | TP Size | 1 | 2 | 4 | 8 | 1 | 2 | 4 | 8 | -| ISL, OSL| | | | | | | | | | +| | TP Size | 1 | 2 | 4 | 8 | 1 | 2 | 4 | 8 | +| ISL, OSL | | | | | | | | | | | | | | | | | | | | | -| 128, 128 | | 3,657.58 | 6,477.50 | 10,466.04 | 15,554.57 | 3,191.27 | 6,183.41 | 10,260.68 | 14,686.01 | -| 128, 2048 | | 4,351.07 | 8,450.31 | 13,438.71 | 20,750.58 | 745.19 | 5,822.02 | 11,442.01 | 17,463.99 | -| 128, 4096 | | 2,696.61 | 5,598.92 | 11,524.93 | 16,634.90 | | 3,714.87 | 8,209.91 | 12,598.55 | -| 500, 2000 | | 3,475.58 | 6,712.35 | 12,332.32 | 17,311.28 | | 4,704.31 | 10,278.02 | 14,630.41 | -| 1000, 1000 | | 2,727.42 | 5,097.36 | 8,698.15 | 12,794.92 | 734.67 | 4,191.26 | 7,427.35 | 11,082.48 | -| 1000, 2000 | | 2,913.54 | 5,841.15 | 9,016.49 | 13,174.68 | 526.31 | 3,920.44 | 7,590.35 | 11,108.11 | -| 1024, 2048 | | 2,893.02 | 5,565.28 | 9,017.72 | 13,117.34 | 525.43 | 3,896.14 | 7,557.32 | 11,028.32 | -| 2048, 128 | | 433.30 | 772.97 | 1,278.26 | 1,947.33 | 315.90 | 747.51 | 1,240.12 | 1,840.12 | -| 2048, 2048 | | 1,990.25 | 3,822.83 | 7,068.68 | 10,529.06 | 357.98 | 2,732.86 | 5,640.31 | 8,772.88 | -| 5000, 500 | | 543.88 | 1,005.81 | 1,714.77 | 2,683.22 | 203.27 | 866.77 | 1,571.92 | 2,399.78 | -| 20000, 2000 | | 276.99 | 618.01 | 1,175.35 | 2,021.08 | | 408.43 | 910.77 | 1,568.84 | +| 128, 128 | | 3,605.47 | 6,427.69 | 10,407.42 | 15,434.37 | 3,128.33 | 6,216.91 | | | +| 128, 2048 | | 4,315.80 | 8,464.03 | 13,508.59 | 20,759.72 | 756.42 | 5,782.57 | 11,464.94 | 17,424.32 | +| 128, 4096 | | 2,701.17 | 5,573.55 | 11,458.56 | 16,668.75 | | 3,868.37 | 8,206.39 | 12,624.61 | +| 500, 2000 | | 3,478.76 | 6,740.06 | 12,200.18 | | | 4,684.06 | 9,903.53 | 14,553.93 | +| 1000, 1000 | | 2,744.32 | 5,119.72 | 8,685.44 | 12,744.51 | 742.14 | 4,247.19 | 7,435.65 | 11,018.81 | +| 1000, 2000 | | 2,896.44 | 5,847.26 | 9,031.21 | 13,141.17 | 533.74 | 3,866.53 | 7,611.12 | 11,139.22 | +| 1024, 2048 | | 2,874.18 | 5,568.61 | 8,946.71 | 13,082.62 | 530.16 | 3,796.68 | 7,575.24 | 11,004.31 | +| 2048, 128 | | 435.90 | 772.67 | 1,264.76 | | | 736.89 | 1,213.33 | 1,839.22 | +| 2048, 2048 | | | | | 10,412.85 | | | | | +| 5000, 500 | | 545.96 | 997.15 | 1,698.22 | 2,655.28 | 204.94 | 862.91 | 1,552.68 | 2,369.84 | +| 20000, 2000 | | 276.66 | 620.33 | 1,161.29 | 1,985.85 | | 416.13 | 903.66 | 1,554.10 | #### Llama 3.1 405B FP8 -| | GPU | H200 141GB HBM3 | H100 80GB HBM3 | + +| | GPU | H200 141GB HBM3 | H100 80GB HBM3 | |:-----------------------------|:---|:------------------|:-----------------| -| | TP Size | 8 | 8 | +| | TP Size | 8 | 8 | | ISL, OSL | | | | | | | | | -| 128, 128 | | 3,800.11 | 3,732.40 | -| 128, 2048 | | 5,661.13 | 4,572.23 | -| 128, 4096 | | 5,167.18 | 2,911.42 | -| 500, 2000 | | 4,854.29 | 3,661.85 | -| 1000, 1000 | | 3,332.15 | 2,963.36 | -| 1000, 2000 | | 3,682.15 | 3,253.17 | -| 1024, 2048 | | 3,685.56 | 3,089.16 | -| 2048, 128 | | 453.42 | 448.89 | -| 2048, 2048 | | 3,055.73 | 2,139.94 | -| 5000, 500 | | 656.11 | 579.14 | -| 20000, 2000 | | 514.02 | 370.26 | +| 128, 2048 | | 5,567.87 | | +| 128, 4096 | | 5,136.85 | | +| 500, 2000 | | 4,787.61 | 3,673.91 | +| 1000, 1000 | | 3,286.30 | 3,012.22 | +| 1000, 2000 | | 3,636.76 | 3,262.20 | +| 1024, 2048 | | 3,618.66 | 3,109.70 | +| 2048, 128 | | 443.10 | 449.02 | +| 5000, 500 | | 645.46 | | +| 20000, 2000 | | | 372.12 | + +#### Llama 4 Maverick FP8 + +| | GPU | H200 141GB HBM3 | H100 80GB HBM3 | +|:-----------------------------|:---|:------------------|:-----------------| +| | TP Size | 8 | 8 | +| ISL, OSL | | | | +| | | | | +| 128, 2048 | | 27,543.87 | | +| 128, 4096 | | 18,541.01 | 11,163.12 | +| 500, 2000 | | 21,117.34 | | +| 1000, 2000 | | | 10,556.00 | +| 1024, 2048 | | 16,859.45 | 11,584.33 | +| 2048, 128 | | 4,364.06 | 3,832.38 | +| 2048, 2048 | | 12,800.89 | | +| 5000, 500 | | 5,128.60 | | +| 20000, 2000 | | 1,764.27 | 1,400.79 | ## Reproducing Benchmarked Results @@ -198,6 +216,8 @@ a model name (HuggingFace reference or path to a local model), a [generated data trtllm-bench --model $model_name throughput --dataset $dataset_file --backend pytorch --extra_llm_api_options $llm_options ``` +The data collected for the v0.20 benchmarks was run with the following file: + `llm_options.yml` ```yaml cuda_graph_config: @@ -220,7 +240,7 @@ cuda_graph_config: - 8192 ``` -In majority of cases, we also use a higher KV cache percentage by setting `--kv_cache_free_gpu_mem_fraction 0.95` in the benchmark command. This allows us to obtain better performance than the default setting of `0.90`. We fall back to `0.90` if we hit an out of memory issue. +In a majority of cases, we also use a higher KV cache percentage by setting `--kv_cache_free_gpu_mem_fraction 0.95` in the benchmark command. This allows us to obtain better performance than the default setting of `0.90`. We fall back to `0.90` if we hit an out of memory issue. The results will be printed to the terminal upon benchmark completion. For example, diff --git a/docs/source/quick-start-guide.md b/docs/source/quick-start-guide.md index b3027e0737ae..12b9a5ec0379 100644 --- a/docs/source/quick-start-guide.md +++ b/docs/source/quick-start-guide.md @@ -8,13 +8,15 @@ This is the starting point to try out TensorRT-LLM. Specifically, this Quick Sta There are multiple ways to install and run TensorRT-LLM. For most users, the options below should be ordered from simple to complex. The approaches are equivalent in terms of the supported features. +Note: **This project will download and install additional third-party open source software projects. Review the license terms of these open source projects before use.** + 1. [](installation/containers) 1. Pre-built release wheels on [PyPI](https://pypi.org/project/tensorrt-llm) (see [](installation/linux)) 1. [Building from source](installation/build-from-source-linux) -The following examples can most easily be executed using the prebuilt [Docker release container available on NGC](https://registry.ngc.nvidia.com/orgs/nvstaging/teams/tensorrt-llm/containers/release) (see also [release.md](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docker/release.md) on GitHub). +The following examples can most easily be executed using the prebuilt [Docker release container available on NGC](https://registry.ngc.nvidia.com/orgs/nvstaging/teams/tensorrt-llm/containers/release) (see also [release.md](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docker/release.md) on GitHub). Ensure to run these commands as a user with appropriate permissions, preferably `root`, to streamline the setup process. ## LLM API @@ -92,7 +94,7 @@ For detailed examples and command syntax, refer to the [trtllm-serve](commands/t 2. Open a new terminal and use the following command to directly attach to the running container: -```bash +```bash:docs/source/quick-start-guide.md docker exec -it bash ``` diff --git a/docs/source/reference/ci-overview.md b/docs/source/reference/ci-overview.md index 9002ae6ab333..30cc613a2e38 100644 --- a/docs/source/reference/ci-overview.md +++ b/docs/source/reference/ci-overview.md @@ -55,9 +55,27 @@ The array elements are: GPU type, YAML file (without extension), shard index, an 2. Search `jenkins/L0_Test.groovy` for a stage whose YAML file matches (for example `l0_a100`) and whose name contains `[Post-Merge]` if the YAML entry uses `stage: post_merge`. 3. The resulting stage name(s) are what you pass to Jenkins via the `stage_list` parameter when triggering a job. -### Example +### Using `test_to_stage_mapping.py` + +Manually searching YAML and Groovy files can be tedious. The helper script +`scripts/test_to_stage_mapping.py` automates the lookup: + +```bash +python scripts/test_to_stage_mapping.py --tests "triton_server/test_triton.py::test_gpt_ib_ptuning[gpt-ib-ptuning]" +python scripts/test_to_stage_mapping.py --tests gpt_ib_ptuning +python scripts/test_to_stage_mapping.py --stages A100X-Triton-Post-Merge-1 +python scripts/test_to_stage_mapping.py --test-list my_tests.txt +python scripts/test_to_stage_mapping.py --test-list my_tests.yml +``` + +The first two commands print the Jenkins stages that run the specified tests or +patterns. Patterns are matched by substring, so partial test names are +supported out of the box. The third lists every test executed in the given stage. When +providing tests on the command line, quote each test string so the shell does +not interpret the `[` and `]` characters as globs. Alternatively, store the +tests in a newline‑separated text file or a YAML list and supply it with +`--test-list`. -`triton_server/test_triton.py::test_gpt_ib_ptuning[gpt-ib-ptuning]` appears in `l0_a100.yml` under `stage: post_merge` and `backend: triton`. The corresponding Jenkins stages are `A100X-Triton-[Post-Merge]-1` and `A100X-Triton-[Post-Merge]-2` (two shards). To run the same tests on your pull request, comment: @@ -67,6 +85,7 @@ To run the same tests on your pull request, comment: This executes the same tests that run post-merge for this hardware/backend. + ## Waiving tests Sometimes a test is known to fail due to a bug or unsupported feature. Instead diff --git a/docs/source/reference/support-matrix.md b/docs/source/reference/support-matrix.md index 37fada2c0ded..0c59baf992bc 100644 --- a/docs/source/reference/support-matrix.md +++ b/docs/source/reference/support-matrix.md @@ -25,6 +25,8 @@ TensorRT-LLM optimizes the performance of a range of well-known models on NVIDIA | `Qwen2ForRewardModel` | Qwen2-based | `Qwen/Qwen2.5-Math-RM-72B` | L | | `Qwen2VLForConditionalGeneration` | Qwen2-VL | `Qwen/Qwen2-VL-7B-Instruct` | L + V | | `Qwen2_5_VLForConditionalGeneration` | Qwen2.5-VL | `Qwen/Qwen2.5-VL-7B-Instruct` | L + V | +| `Qwen3ForCausalLM` | Qwen3 | `Qwen/Qwen3-8B` | L | +| `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B` | L | Note: - L: Language only @@ -72,7 +74,7 @@ Note: - [mT5](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) - [OPT](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/opt) - [Phi-1.5/Phi-2/Phi-3](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/phi) -- [Qwen/Qwen1.5/Qwen2](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/qwen) +- [Qwen/Qwen1.5/Qwen2/Qwen3](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/qwen) - [Qwen-VL](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/qwenvl) - [RecurrentGemma](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/recurrentgemma) - [Replit Code](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/mpt) [^replitcode] diff --git a/docs/source/release-notes.md b/docs/source/release-notes.md index bb663aba7d23..dee84ecfde50 100644 --- a/docs/source/release-notes.md +++ b/docs/source/release-notes.md @@ -4,6 +4,152 @@ All published functionality in the Release Notes has been fully tested and verified with known limitations documented. To share feedback about this release, access our [NVIDIA Developer Forum](https://forums.developer.nvidia.com/). +## TensorRT-LLM Release 0.21.0 + +### Key Features and Enhancements +- **Model Support** + - Added Gemma3 VLM support +- **Features** + - Added large-scale EP support + - Integrated NIXL into the communication layer of the disaggregated service + - Added fabric Memory support for KV Cache Transfer + - Added MCP in ScaffoldingLLM + - Added support for w4a8_mxfp4_fp8 quantization + - Added support for fp8 rowwise quantization + - Added generation logits support in TRTLLM Sampler + - Added log probs support in TRTLLM Sampler + - Optimized TRTLLM Sampler perf single beam single step + - Enabled Disaggregated serving for Qwen-3 + - Added EAGLE3 support for Qwen-3 + - Fused finalize and allreduce for Qwen-MoE model + - Refactored Fused MoE module + - Added support for chunked attention on Blackwell and Hopper + - Introduced sliding-window attention kernels for the generation phase on Blackwell + - Updated DeepSeek FP8 TRT-LLM Gen cubins to improve performance in large batch size scenarios + - Added FP8 block-scale GEMM support on SM89 + - Enabled overlap scheduler between draft forwards + - Added Piecewise cuda graph support for MLA + - Added model-agnostic one-engine eagle3 + - Enabled Finalize + Allreduce + add + rmsnorm fusion + - Integrated TRT-LLM Gen FP8 block scale MoE with Pytorch workflow kernel autotuner + - Added support for Eagle3 + disaggregated serving in two model speculative decoding flow + - Validated Llama 3.1 models on H200 NVL +- Benchmark: + - Added all_reduce.py benchmark script for testing + - Added beam width to trtllm-bench latency command + - Fixed trtllm-bench iter_stats and cuda_graph_batch_sizes errors + - Enabled trtllm-bench to run LoRA and add basic e2e perf testing capability for LoRA + - Supported post_proc for bench + - Added no_kv_cache_reuse option and streaming support for trtllm serve bench + +### Infrastructure Changes +- The base Docker image for TensorRT-LLM is updated to `nvcr.io/nvidia/pytorch:25.05-py3`. +- The base Docker image for TensorRT-LLM Backend is updated to `nvcr.io/nvidia/tritonserver:25.05-py3`. +- The dependent public PyTorch version is updated to 2.7.1. +- The dependent TensorRT version is updated to 10.11. +- The dependent NVIDIA ModelOpt version is updated to 0.31. +- The dependent NCCL version is updated to 2.27.5. + +### API Changes +- Set _AutoDeployLlmArgs as primary config object +- Removed decoder request from decoder interface +- Enhanced the torch_compile_config in llm args +- Removed the redundant use_kv_cache field from PytorchConfig +- Moved allreduce_strategy from committed api to reference + +### Fixed Issues +- Fixed disaggregated service hang when MNNVL two-shot AllReduce is enabled (#4678) +- Fixed EP load balancer with MTP layer and route offset by EP rank (#4767) +- Fixed cuda graph padding for spec decoding (#4853) +- Fixed llama 4 long context issue (#4809) +- Fixed max_num_sequences calculation with overlap scheduling (#4532) +- Fixed chunked prefill + overlap scheduling (#5761) +- Fixed trtllm-bench hang issue due to LLM API IPC (#4798) +- Fixed index out of bounds error in spec decoding (#5954) +- Fixed MTP illegal memory access in cuda graph warmup (#5947) +- Fixed no free slots error with spec decode + disagg (#5975) +- Fixed one-off attention window size for Gemma3 1B (#5564) + +### Known Issues +- accuracy/test_cli_flow::TestGpt2::test_beam_search_large is broken. +- Enabling disaggregated serving, MTP, and the overlap scheduler at the same time can lead to accuracy problems. + +## TensorRT-LLM Release 0.20.0 + +### Key Features and Enhancements +- **Model Support** + - Added Qwen3 support.Refer to “Qwen3” section in `examples/models/core/qwen/README.md`. + - Added HyperCLOVAX-SEED-Vision support in PyTorch flow. Refer to `examples/models/contrib/hyperclovax/README.md` + - Added Dynasor-CoT in scaffolding examples. Refer to `examples/scaffolding/contrib/Dynasor/README.md` + - Added Mistral Small 3.1 24B VLM support in TRT workflow + - Added Gemma3-1b-it support in PyTorch workflow + - Added Nemotron-H model support + - Added Eagle-3 support for LLAMA4 +- **PyTorch workflow** + - Added lora support + - Added return logits support + - Adopt new logprob definition in PyTorch flow + - Enabled per-request stats with PyTorch backend + - Enabled LogitsProcessor in PyTorch backend +- Benchmark: + - Add beam width to low latency. + - Fix trtllm-bench iter_stats and cuda_graph_batch_sizes errors. + - Remove deprecated Python runtime benchmark + - Add benchmark support for scaffolding +- Multimodal models + - Added support in trtllm-serve + - Added support in trtllm-bench, the support is limited to image only for now +- Supported DeepSeek-R1 W4A8 on Hopper +- Add the RTX Pro 6000 support on single GPU +- Integrated Llama4 input processor +- Added CGA reduction FHMA kernels on Blackwell +- Enabled chunked context for FlashInfer +- Supported KV cache reuse for MLA +- Added Piecewise CUDA Graph support +- Supported multiple LoRA adapters and TP +- Added KV cache-aware router for disaggregated serving +- Unfused attention for native support +- Added group_rms_norm kernel to normalize multiple inputs in a single operator +- Added smart router for the MoE module +- Added head size 72 support for QKV preprocessing kernel +- Added MNNVL MoE A2A support +- Optimized Large Embedding Tables in Multimodal Models +- Supported Top-K logprobs and prompt_logprobs in LLMAPI +- Enabled overlap scheduler in TRT workflow via executor API + +### Infrastructure Changes +- **TRT-LLM team formally releases docker image on [NGC](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tensorrt-llm/containers/release/tags)**. +- The pre-built TensorRT-LLM wheel on PyPI is linked against PyTorch 2.7.0 now, which uses the CXX11 ABI +- The dependent TensorRT version is updated to 10.10.0 +- The dependent CUDA version is updated to 12.9.0 +- The dependent public PyTorch version is updated to 2.7.0 +- The dependent NVIDIA ModelOpt version is updated to 0.29.0 +- The dependent NCCL version is maintained at 2.25.1 +- Open-sourced XQA kernels +- Dependent datasets version was upgraded to 3.1.0 +- Migrate Triton Backend to TensorRT LLM repo to TensorRT LLM submodule +- Downgrade gcc toolset version from 13 to 11 + +### API Changes +- [Breaking Change]:Enable scheduling overlap by default +- Remove deprecated GptSession/V1 from TRT workflow +- Set _AutoDeployLlmArgs as primary config object +- Allow overriding CLI arguments with YAML file in trtllm-serve +- Introduced multimodal embedding field in LlmRequest + + +### Fixed Issues +- Fix hang bug when context server doesn't have enough capacity for KV Cache (#3095) +- Fix C++ decoder synchronization in PyTorch (#3106) +- Fix bug of create cuda stream as default parameter which will be initialized during importing (#3764) +- Fix bug related to creating CUDA stream as default parameter, which will be initialized during importing (#3764) +- Fix attention DP bug on Qwen3 MoE model (#4141) +- Fix illegal memory access when running LLaMA 4 with CUDA Graph enabled (#4101) +- Reset planned states to avoid memory leak in TrtllmAttentionWrapper (#4227) + +### Known Issues +- multi-GPU model support on RTX Pro 6000 + ## TensorRT-LLM Release 0.19.0 diff --git a/docs/source/scripts/disaggregated/disaggr_torch.slurm b/docs/source/scripts/disaggregated/disaggr_torch.slurm deleted file mode 100644 index ae047c23552f..000000000000 --- a/docs/source/scripts/disaggregated/disaggr_torch.slurm +++ /dev/null @@ -1,112 +0,0 @@ -#!/bin/bash -#SBATCH --nodes=2 -#SBATCH --ntasks=8 -#SBATCH --ntasks-per-node=4 -#SBATCH --partition=batch -#SBATCH --account=${account} -#SBATCH --time=02:00:00 -#SBATCH --job-name="${account}:disaggr-test" - -isl=8192 -osl=256 -multi_round=10 -gen_yaml_file=gen_yaml.py -container_image=${docker_image} -mount_dir=/${account}/${user}/ -workdir=/${account}/${user}/8k-${osl}/disaggr-e2e/ -model_dir=/${account}/${user}/DeepSeek-R1-nvfp4_allmoe/ -logdir=$workdir/bm_deepseek-r1-8k-${osl}-disaggr-e2e-nostream -streaming=false -mkdir -p ${logdir} - -dep_dir=${workdir} -run_benchmark_cmd="bash ${dep_dir}/run_benchmark.sh" - -container_name=disaggr-test - -num_ctx_servers=$1 -ctx_tp_size=$2 -ctx_batch_size=$3 -ctx_max_num_tokens=$4 -ctx_enable_attention_dp=$5 -num_gen_servers=$6 -gen_tp_size=$7 -gen_batch_size=$8 -gen_max_num_tokens=$9 -gen_enable_attention_dp=${10} -gen_gpu_memory_fraction=${11} -concurrency_list=${12} -sub_file=${13} - -# concurrency=$((concurrency * gen_tp_size)) -echo "concurrency_list: ${concurrency_list}" - -ctx_gpus=$((num_ctx_servers * ctx_tp_size)) -gen_gpus=$((num_gen_servers * gen_tp_size)) - -echo "enable_attention_dp: ${ctx_enable_attention_dp}, ${gen_enable_attention_dp}, gpu_memory_fraction: ${gen_gpu_memory_fraction}" - -enable_pdl=false -if [ "${gen_enable_attention_dp}" = "false" ]; then - enable_pdl=true -fi - -full_logdir=${logdir}/${sub_file} -mkdir -p ${full_logdir} - -# start the container -srun -l --container-image=${container_image} \ - --container-name=${container_name} \ - --container-mounts=${mount_dir}:${mount_dir} \ - --mpi=pmix \ - echo "Container up." - -# generate the yaml file -srun -l --container-name=${container_name} \ - --container-mounts=${mount_dir}:${mount_dir} \ - --mpi=pmix --overlap \ - python3 ${dep_dir}/${gen_yaml_file} --config ${full_logdir}/config.yaml \ - --model ${model_dir} \ - --num_ctx_servers ${num_ctx_servers} \ - --ctx_tp_size ${ctx_tp_size} \ - --ctx_batch_size ${ctx_batch_size} \ - --ctx_max_num_tokens ${ctx_max_num_tokens} \ - --num_gen_servers ${num_gen_servers} \ - --gen_tp_size ${gen_tp_size} \ - --gen_batch_size ${gen_batch_size} \ - --gen_max_num_tokens ${gen_max_num_tokens} \ - --gen_gpu_memory_fraction ${gen_gpu_memory_fraction} \ - $(if [ "${gen_enable_attention_dp}" = "true" ]; then echo "--gen_enable_attention_dp"; fi) \ - $(if [ "${ctx_enable_attention_dp}" = "true" ]; then echo "--ctx_enable_attention_dp"; fi) - -echo "YAML file generated." - -hostname_value=$(grep '^hostname:' ${full_logdir}/config.yaml | awk -F': ' '{print $2}') -echo "server host name: $hostname_value" - -nsys_on="" -# nsys_on=${full_logdir} - -# start the workers -srun -l --container-name=${container_name} \ - --container-mounts=${mount_dir}:${mount_dir} \ - --mpi=pmix --overlap \ - bash ${dep_dir}/start_worker.sh ${full_logdir}/config.yaml "${enable_pdl}" ${ctx_gpus} ${nsys_on} &> ${full_logdir}/output_workers.log & -# start the server -srun -l --container-name=${container_name} \ - --container-mounts=${mount_dir}:${mount_dir} \ - --mpi=pmix --overlap -N 1 -n 1 \ - bash trtllm-serve disaggregated -c ${full_logdir}/config.yaml -t 1800 -r 1800 &> ${full_logdir}/output_server.log & -# start benchmark -srun -l --container-name=${container_name} \ - --container-mounts=${mount_dir}:${mount_dir} \ - --mpi=pmix --overlap -N 1 -n 1 \ - --nodelist=${hostname_value} \ - ${run_benchmark_cmd} ${isl} ${osl} ${multi_round} ${model_dir} "${concurrency_list}" ${streaming} ${full_logdir}/ > ${full_logdir}/benchmark.log 2>&1 -wait - -# try to kill the server and workers -srun -l --container-name=${container_name} \ - --container-mounts=${mount_dir}:${mount_dir} \ - --mpi=pmix --overlap \ - pkill -f "trtllm-serve" || true diff --git a/docs/source/scripts/disaggregated/gen_yaml.py b/docs/source/scripts/disaggregated/gen_yaml.py deleted file mode 100644 index 1d198a9766db..000000000000 --- a/docs/source/scripts/disaggregated/gen_yaml.py +++ /dev/null @@ -1,301 +0,0 @@ -import argparse -import os -import re -from typing import Dict, List - -import yaml - - -def process_node_and_task() -> tuple[int, List[str], List[str]]: - """ - Process SLURM node and task environment variables. - - Returns: - tuple: (max_tasks_per_node, nodes, task_nodes) - """ - slurm_job_nodelist = os.getenv('SLURM_JOB_NODELIST', '') - print(f"SLURM_JOB_NODELIST: {slurm_job_nodelist}") - if not slurm_job_nodelist: - raise ValueError(f"Environment variable SLURM_JOB_NODELIST not found.") - - slurm_tasks_per_node = os.getenv('SLURM_TASKS_PER_NODE', '') - print(f"SLURM_TASKS_PER_NODE: {slurm_tasks_per_node}") - if not slurm_tasks_per_node: - raise ValueError( - f"Environment variable SLURM_TASKS_PER_NODE not found.") - - # Generate list of nodes - if '[' in slurm_job_nodelist: - # Handle nodelist with range format (e.g., "ptyche[0065-0066]") - node_prefix = re.match(r'^[a-zA-Z]+', slurm_job_nodelist).group(0) - node_range = re.search(r'\[(.*?)\]', slurm_job_nodelist).group(1) - nodes = [] - for part in node_range.split(','): - if '-' in part: - start, end = part.split('-') - # Get the width of the number format from the first number - width = len(start) - # Convert to integers after getting the width - start, end = int(start), int(end) - # Format numbers with leading zeros - nodes.extend([ - f"{node_prefix}{str(i).zfill(width)}" - for i in range(start, end + 1) - ]) - else: - # Preserve the original format for single numbers - nodes.append(f"{node_prefix}{part}") - else: - # Handle single node format (e.g., "ptyche0065") - nodes = [slurm_job_nodelist] - print(f"Nodes: {nodes}") - - # Generate tasks per node - tasks_per_node = [] - for part in slurm_tasks_per_node.split(','): - if '(x' in part: - count, repeat = map(int, re.findall(r'\d+', part)) - tasks_per_node.extend([count] * repeat) - else: - tasks_per_node.append(int(part)) - print(f"Tasks per node: {tasks_per_node}") - - if (len(tasks_per_node) != len(nodes)): - raise ValueError( - f"Number of nodes and tasks per node do not match. Number of nodes: {len(nodes)}, Number of tasks per node: {len(tasks_per_node)}" - ) - - max_tasks_per_node = max(tasks_per_node) - task_nodes = [] - for node, tasks in zip(nodes, tasks_per_node): - task_nodes.extend([node] * tasks) - - return max_tasks_per_node, nodes, task_nodes - - -def generate_urls(ctx_or_gen: str, - num_instances: int, - tensor_parallel_size: int, - pipeline_parallel_size: int, - max_tasks_per_node: int, - nodes: List[str], - task_nodes: List[str], - node_to_port: Dict[str, int], - task_nodes_offset: int = 0) -> tuple[List[str], int]: - """ - Generate URLs for context or generation servers. - - Returns: - tuple: (urls, updated_task_nodes_offset) - """ - urls = [] - - for instance in range(num_instances): - tasks_needed = tensor_parallel_size * pipeline_parallel_size - - if (task_nodes_offset + tasks_needed) > len(task_nodes): - print(f"{ctx_or_gen} urls so far: {urls}") - raise ValueError( - f"For {ctx_or_gen} instance {instance}, there are not enough tasks available. task_nodes_offset: {task_nodes_offset}, tasks_needed: {tasks_needed}, len(task_nodes): {len(task_nodes)}" - ) - - min_node = (tasks_needed + max_tasks_per_node - 1) / max_tasks_per_node - instance_nodes = set(task_nodes[task_nodes_offset:task_nodes_offset + - tasks_needed]) - if len(instance_nodes) > min_node: - raise ValueError( - f"Tasks for a instance {instance} of {ctx_or_gen} instances use more node than expected. Nodes used: {instance_nodes}, number of nodes expected: {min_node}, max_tasks_per_node: {max_tasks_per_node}" - ) - - node = task_nodes[task_nodes_offset] - port = node_to_port[node] - node_to_port[node] += 1 - task_nodes_offset += tasks_needed - - urls.append(f"{node}:{port}") - - print(f"{ctx_or_gen} urls: {urls}") - return urls, task_nodes_offset - - -def gen_config_file(config_path: str, - model_path: str, - num_ctx_servers: int, - ctx_tp_size: int, - ctx_batch_size: int, - ctx_max_num_tokens: int, - ctx_enable_attention_dp: bool, - num_gen_servers: int, - gen_tp_size: int, - gen_batch_size: int, - gen_max_num_tokens: int, - gen_enable_attention_dp: bool, - gen_gpu_memory_fraction: float, - worker_start_port: int = 8001, - server_port: int = 8000) -> None: - """ - Generate configuration YAML file for disaggregated inference. - - Args: - config_path: Path to save the config file - model_path: Path to the model - num_ctx_servers: Number of context servers - ctx_tp_size: Tensor parallel size for context servers - ctx_batch_size: Batch size for context servers - ctx_max_num_tokens: Max number of tokens for context servers - ctx_enable_attention_dp: Enable attention DP for context servers - num_gen_servers: Number of generation servers - gen_tp_size: Tensor parallel size for generation servers - gen_batch_size: Batch size for generation servers - gen_max_num_tokens: Max number of tokens for generation servers - gen_enable_attention_dp: Enable attention DP for generation servers - gen_gpu_memory_fraction: GPU memory fraction for generation servers - worker_start_port: Start port for workers - server_port: Server port - """ - gen_cuda_graph_batch_sizes = [ - 1, 2, 4, 8, 16, 32, 64, 128, 256, gen_batch_size - ] - - config = { - 'model': model_path, - 'hostname': 'localhost', - 'port': server_port, - 'backend': 'pytorch', - 'context_servers': { - 'num_instances': num_ctx_servers, - 'max_batch_size': ctx_batch_size, - 'max_num_tokens': ctx_max_num_tokens, - 'max_seq_len': 8300, - 'free_gpu_memory_fraction': 0.7, - 'tensor_parallel_size': ctx_tp_size, - 'moe_expert_parallel_size': ctx_tp_size, - 'enable_attention_dp': ctx_enable_attention_dp, - 'pipeline_parallel_size': 1, - 'print_iter_log': True, - 'disable_overlap_scheduler': True, - 'kv_cache_dtype': 'fp8', - 'cache_transceiver_config': { - 'max_num_tokens': 8320, - }, - }, - 'generation_servers': { - 'num_instances': num_gen_servers, - 'tensor_parallel_size': gen_tp_size, - 'moe_expert_parallel_size': gen_tp_size, - 'enable_attention_dp': gen_enable_attention_dp, - 'pipeline_parallel_size': 1, - 'max_batch_size': gen_batch_size, - 'max_num_tokens': gen_max_num_tokens, - 'max_seq_len': 8576, - 'free_gpu_memory_fraction': gen_gpu_memory_fraction, - 'cuda_graph_config': { - 'enable_padding': True, - 'batch_sizes': gen_cuda_graph_batch_sizes, - }, - 'print_iter_log': True, - 'kv_cache_dtype': 'fp8', - 'moe_config': { - 'backend': 'TRTLLM', - }, - 'cache_transceiver_config': { - 'max_num_tokens': 8320, - }, - } - } - - # Process nodes and generate URLs - max_tasks_per_node, nodes, task_nodes = process_node_and_task() - node_ports = {node: worker_start_port for node in nodes} - - # Generate URLs for context and generation servers - ctx_urls, task_nodes_offset = generate_urls("ctx", num_ctx_servers, - ctx_tp_size, 1, - max_tasks_per_node, nodes, - task_nodes, node_ports) - if num_ctx_servers > 0: - config['context_servers']['urls'] = ctx_urls - - gen_urls, _ = generate_urls("gen", num_gen_servers, gen_tp_size, 1, - max_tasks_per_node, nodes, task_nodes, - node_ports, task_nodes_offset) - config['generation_servers']['urls'] = gen_urls - - # set the hostname to the first node - config['hostname'] = nodes[0] - - # Write config to file - with open(config_path, 'w') as f: - yaml.dump(config, f, default_flow_style=False, sort_keys=False) - - -# gen main and args -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--config", type=str, default="/tmp/config.yaml") - parser.add_argument("--model", - type=str, - required=True, - help="Path to the model") - parser.add_argument("--num_ctx_servers", - type=int, - required=True, - help="Number of context servers") - parser.add_argument("--ctx_tp_size", - type=int, - required=True, - help="Tensor parallel size for context servers") - parser.add_argument("--ctx_batch_size", - type=int, - required=True, - help="Batch size for context servers") - parser.add_argument("--ctx_max_num_tokens", - type=int, - required=True, - help="Max number of tokens for context servers") - parser.add_argument("--ctx_enable_attention_dp", - dest='ctx_enable_attention_dp', - action='store_true', - help="Enable attention DP for context servers") - parser.add_argument("--num_gen_servers", - type=int, - required=True, - help="Number of generation servers") - parser.add_argument("--gen_tp_size", - type=int, - required=True, - help="Tensor parallel size for generation servers") - parser.add_argument("--gen_batch_size", - type=int, - required=True, - help="Batch size for generation servers") - parser.add_argument("--gen_max_num_tokens", - type=int, - required=True, - help="Max number of tokens for generation servers") - parser.add_argument("--gen_enable_attention_dp", - dest='gen_enable_attention_dp', - action='store_true', - help="Enable attention DP for generation servers") - parser.add_argument("--gen_gpu_memory_fraction", - type=float, - required=True, - help="GPU memory fraction for generation servers") - parser.add_argument("--worker_start_port", - type=int, - default=8336, - help="Start port for workers") - parser.add_argument("--server_port", - type=int, - default=8333, - help="Server port") - - args = parser.parse_args() - - gen_config_file(args.config, args.model, args.num_ctx_servers, - args.ctx_tp_size, args.ctx_batch_size, - args.ctx_max_num_tokens, args.ctx_enable_attention_dp, - args.num_gen_servers, args.gen_tp_size, args.gen_batch_size, - args.gen_max_num_tokens, args.gen_enable_attention_dp, - args.gen_gpu_memory_fraction, args.worker_start_port, - args.server_port) diff --git a/docs/source/scripts/disaggregated/run_benchmark.sh b/docs/source/scripts/disaggregated/run_benchmark.sh deleted file mode 100644 index 00c213499961..000000000000 --- a/docs/source/scripts/disaggregated/run_benchmark.sh +++ /dev/null @@ -1,98 +0,0 @@ -#!/bin/bash - -set -e -set -u -trap 'echo "Error occurred at line $LINENO"; exit 1' ERR - -if [ "$#" -lt 7 ]; then - echo "Error: Missing required arguments" - echo "Usage: $0 isl osl multi_round model_name concurrency_list streaming log_path" - exit 1 -fi - -isl=$1 -osl=$2 -multi_round=$3 -model_name=$4 -concurrency_list=$5 -streaming=$6 -log_path=$7 - -set -x -config_file=${log_path}/config.yaml - -# check if the config file exists every 10 seconds timeout 1800 seconds -timeout=1800 -start_time=$(date +%s) -while [ ! -f ${config_file} ]; do - current_time=$(date +%s) - elapsed=$((current_time - start_time)) - if [ $elapsed -ge $timeout ]; then - echo "Error: Config file ${config_file} not found within ${timeout} seconds" - exit 1 - fi - if [ $((elapsed % 30)) -eq 0 ]; then - echo "Waiting for config file... (${elapsed}s elapsed)" - fi - sleep 10 -done - -# grep the host and port from the config file -hostname=$(grep -i "hostname:" ${config_file} | awk '{print $2}') -port=$(grep -i "port:" ${config_file} | awk '{print $2}') -if [ -z "$hostname" ] || [ -z "$port" ]; then - echo "Error: Failed to extract hostname or port from config file" - exit 1 -fi -echo "Hostname: ${hostname}, Port: ${port}" - -# check server is health by curl every 10 seconds timeout 1800 seconds -timeout=1800 -start_time=$(date +%s) -while ! curl -s -o /dev/null -w "%{http_code}" http://${hostname}:${port}/health; do - hostname=$(grep -i "hostname:" ${config_file} | awk '{print $2}') - port=$(grep -i "port:" ${config_file} | awk '{print $2}') - echo "Hostname: ${hostname}, Port: ${port}" - current_time=$(date +%s) - elapsed=$((current_time - start_time)) - if [ $elapsed -ge $timeout ]; then - echo "Error: Server is not healthy after ${timeout} seconds" - exit 1 - fi - if [ $((elapsed % 30)) -eq 0 ]; then - echo "Waiting for server to be healthy... (${elapsed}s elapsed)" - fi - sleep 10 -done - -# run the benchmark -for concurrency in ${concurrency_list}; do - mkdir -p ${log_path}/concurrency_${concurrency} - max_count=$((${concurrency} * ${multi_round})) - echo "Running benchmark with concurrency: ${concurrency}, max_count: ${max_count}" - python -m tensorrt_llm.serve.scripts.benchmark_serving \ - --model ${model_name} \ - --tokenizer ${model_name} \ - --dataset-name random \ - --random-ids \ - --random-input-len ${isl} \ - --random-output-len ${osl} \ - --random-prefix-len 0 \ - --num-prompts ${max_count} \ - --max-concurrency ${concurrency} \ - --host ${hostname} \ - --port ${port} \ - --ignore-eos - echo "done for ${concurrency} in folder ${log_path}/concurrency_${concurrency}" -done - -echo "Benchmark done, gracefully shutting down server and workers..." -pkill -f "start_worker.sh" || true -pkill -f "trtllm-serve" || true -sleep 20 # - -if pgrep -f "trtllm-serve"; then - echo "Warning: Some processes may still be running" -else - echo "All processes successfully terminated" -fi diff --git a/docs/source/scripts/disaggregated/start_worker.sh b/docs/source/scripts/disaggregated/start_worker.sh deleted file mode 100644 index 6ba61d4906e0..000000000000 --- a/docs/source/scripts/disaggregated/start_worker.sh +++ /dev/null @@ -1,32 +0,0 @@ -#! /bin/bash - -config_file=$1 -enable_pdl=$2 -ctx_gpus=$3 -work_dir=$4 - -export TLLM_LOG_LEVEL=INFO -export TRTLLM_USE_MPI_KVCACHE=1 -export TRTLLM_MNNVL_AR_ENABLED=1 - -if [ "${enable_pdl}" = "true" ]; then - export TRTLLM_ENABLE_PDL=1 -fi - -#check if work_dir is provided -if [ -z "${work_dir}" ]; then - trtllm-serve disaggregated_mpi_worker -c ${config_file} -else - nsys_prefix="" - nsys_file=${work_dir}/nsys_worker_proc_${SLURM_PROCID} - export TLLM_PROFILE_RECORD_GC=1 - export TLLM_NVTX_DEBUG=1 - if [ ${SLURM_PROCID} -ge ${ctx_gpus} ]; then - export TLLM_PROFILE_START_STOP=300-400 - else - export TLLM_PROFILE_START_STOP=25-100 - fi - nsys_prefix="nsys profile -e \"NSYS_MPI_STORE_TEAMS_PER_RANK=1\" -o ${nsys_file} -f true -t cuda,nvtx,python-gil -c cudaProfilerApi --cuda-graph-trace node --capture-range-end=stop --gpu-metrics-devices=all" - - ${nsys_prefix} trtllm-serve disaggregated_mpi_worker -c ${config_file} -fi diff --git a/docs/source/scripts/disaggregated/submit.sh b/docs/source/scripts/disaggregated/submit.sh deleted file mode 100644 index 9757dc7d32f1..000000000000 --- a/docs/source/scripts/disaggregated/submit.sh +++ /dev/null @@ -1,36 +0,0 @@ -#! /bin/bash - -slurm_file=disaggr_torch.slurm - -# ctx1dep4_gen1tep4, max_batch16 -for c in 1 2 4 8 16 32 48 64; do - sbatch --nodes=2 --ntasks=8 --ntasks-per-node=4 ${slurm_file} 1 4 1 8300 true 1 4 32 32 false "0.95" "$c" ctx1dep4_gen1tep4_${c} -done - -# ctx2dep4_gen1tep4, max_batch 64 -for c in 64 96 128; do - sbatch --nodes=3 --ntasks=12 --ntasks-per-node=4 ${slurm_file} 2 4 1 8300 true 1 4 64 64 false "0.9" "$c" ctx2dep4_gen1tep4_${c} -done - -for c in 128 192 256; do - sbatch --nodes=4 --ntasks=16 --ntasks-per-node=4 ${slurm_file} 3 4 1 8300 true 1 4 32 32 true "0.9" "$c" ctx3dep4_gen1dep4_${c} -done - -for c in 256 384 512; do - sbatch --nodes=5 --ntasks=20 --ntasks-per-node=4 ${slurm_file} 4 4 1 8300 true 1 4 64 64 true "0.9" "$c" ctx4dep4_gen1dep4_${c} -done - -# ctx5dep4_gen1dep4, max_batch -for c in 256 384 512; do - sbatch --nodes=6 --ntasks=24 --ntasks-per-node=4 ${slurm_file} 5 4 1 8300 true 1 4 64 64 true "0.9" "$c" ctx5dep4_gen1dep4_${c} -done - -# ctx7dep4_gen1dep4 -for c in 512 768 1024; do - sbatch --nodes=8 --ntasks=32 --ntasks-per-node=4 ${slurm_file} 7 4 1 8300 true 1 4 128 128 true "0.9" "$c" ctx7dep4_gen1dep4_${c} -done - -# ctx8dep4_gen1dep4 -for c in 512 768 1024; do - sbatch --nodes=9 --ntasks=36 --ntasks-per-node=4 ${slurm_file} 8 4 1 8300 true 1 4 128 128 true "0.9" "$c" ctx8dep4_gen1dep4_${c} -done diff --git a/docs/source/torch.md b/docs/source/torch.md index 601ab06d8c89..b04c98db1d9c 100644 --- a/docs/source/torch.md +++ b/docs/source/torch.md @@ -13,7 +13,7 @@ The PyTorch backend of TensorRT-LLM is available in version 0.17 and later. You Here is a simple example to show how to use `tensorrt_llm.LLM` API with Llama model. -```{literalinclude} ../../examples/pytorch/quickstart.py +```{literalinclude} ../../examples/llm-api/quickstart_example.py :language: python :linenos: ``` diff --git a/docs/source/torch/arch_overview.md b/docs/source/torch/arch_overview.md index 11b12781cea5..ec7f6e51abf7 100644 --- a/docs/source/torch/arch_overview.md +++ b/docs/source/torch/arch_overview.md @@ -37,7 +37,7 @@ The single-step flow of PyExecutor involves: The core component of `PyExecutor` is the `ModelEngine`, responsible for executing the model's forward pass efficiently on the GPU. The key method of `ModelEngine` is `forward`, which handles the forward pass computation. -For the PyTorch backend, the derived class is `PyTorchModelEngine`, declared in [pytorch_model_engine.py](../../../tensorrt_llm/_torch/pyexecutor/pytorch_model_engine.py). +For the PyTorch backend, the derived class is `PyTorchModelEngine`, declared in [model_engine.py](../../../tensorrt_llm/_torch/pyexecutor/model_engine.py). ## Decoder diff --git a/docs/source/torch/features/feature_combination_matrix.md b/docs/source/torch/features/feature_combination_matrix.md index 8f8d5defe806..f62c1d33aa4d 100644 --- a/docs/source/torch/features/feature_combination_matrix.md +++ b/docs/source/torch/features/feature_combination_matrix.md @@ -15,4 +15,4 @@ | KV Cache Reuse | Yes | Yes | Yes | Untested | Untested | Untested | Yes | No | Yes | Yes | --- | | | | | Slide Window Attention | Yes | Yes | Yes | Untested | Untested | Untested | Untested | Untested | Yes | Yes | WIP | --- | | | | Logits Post Processor | No | Yes | Yes | No | Untested | No | No | No | Yes | Yes | Yes | Yes | --- | | -| Guided Decoding | No | Yes | Yes | Untested | Yes | No | No | No | Yes | Yes | Yes | Yes | Yes | --- | +| Guided Decoding | Yes | Yes | Yes | No | Yes | No | No | No | Yes | Yes | Yes | Yes | Yes | --- | diff --git a/examples/auto_deploy/.vscode/launch.json b/examples/auto_deploy/.vscode/launch.json index fb0e7e64270e..44bc25e6cb3c 100644 --- a/examples/auto_deploy/.vscode/launch.json +++ b/examples/auto_deploy/.vscode/launch.json @@ -16,8 +16,10 @@ "--args.model-factory=AutoModelForCausalLM", "--benchmark.enabled=false", "--prompt.batch-size=2", - "--args.model-kwargs", - "num_hidden_layers=3,num_attention_heads=32", + "--args.model-kwargs.num-hidden-layers=3", + "--args.model-kwargs.num-attention-heads=32", + "--prompt.sp-kwargs.max-tokens=128", + // "--dry-run", // uncomment to print the final config and return ], "console": "integratedTerminal", "justMyCode": false, diff --git a/examples/auto_deploy/README.md b/examples/auto_deploy/README.md index 553ce6e4db54..399d31ce36bd 100644 --- a/examples/auto_deploy/README.md +++ b/examples/auto_deploy/README.md @@ -6,7 +6,7 @@
-AutoDeploy is designed to simplify and accelerate the deployment of PyTorch models, including off-the-shelf models like those from Hugging Face, to TensorRT-LLM. It automates graph transformations to integrate inference optimizations such as tensor parallelism, KV-caching and quantization. AutoDeploy supports optimized in-framework deployment, minimizing the amount of manual modification needed. +AutoDeploy is an experimental feature in beta stage designed to simplify and accelerate the deployment of PyTorch models, including off-the-shelf models like those from Hugging Face, to TensorRT-LLM. It automates graph transformations to integrate inference optimizations such as tensor parallelism, KV-caching and quantization. AutoDeploy supports optimized in-framework deployment, minimizing the amount of manual modification needed. ______________________________________________________________________ @@ -146,7 +146,7 @@ Below is a non-exhaustive list of common config options: | `--args.skip-loading-weights` | Only load the architecture, not the weights | | `--args.model-kwargs` | Extra kwargs that are being passed to the model initializer in the model factory | | `--args.tokenizer-kwargs` | Extra kwargs that are being passed to the tokenizer initializer in the model factory | -| `--args.world-size` | The number of GPUs for Tensor Parallel | +| `--args.world-size` | The number of GPUs used for auto-sharding the model | | `--args.runtime` | Specifies which type of Engine to use during runtime (`"demollm"` or `"trtllm"`) | | `--args.compile-backend` | Specifies how to compile the graph at the end | | `--args.attn-backend` | Specifies kernel implementation for attention | @@ -157,7 +157,7 @@ Below is a non-exhaustive list of common config options: | `--prompt.batch-size` | Number of queries to generate | | `--benchmark.enabled` | Whether to run the built-in benchmark (true/false) | -For default values and additional configuration options, refer to the `ExperimentConfig` class in [build_and_run_ad.py](./build_and_run_ad.py) file. +For default values and additional configuration options, refer to the [`ExperimentConfig`](./build_and_run_ad.py) class in [build_and_run_ad.py](./build_and_run_ad.py) file. Here is a more complete example of using the script: @@ -172,7 +172,7 @@ python build_and_run_ad.py \ --benchmark.enabled True ``` -#### Logging Level +### Logging Level Use the following env variable to specify the logging level of our built-in logger ordered by decreasing verbosity; @@ -223,9 +223,6 @@ AutoDeploy can be seamlessly integrated into your existing workflows using TRT-L Here is an example of how you can build an LLM object with AutoDeploy integration: -
-Click to expand the example - ``` from tensorrt_llm._torch.auto_deploy import LLM @@ -233,7 +230,7 @@ from tensorrt_llm._torch.auto_deploy import LLM # Construct the LLM high-level interface object with autodeploy as backend llm = LLM( model=, - world_size=, + world_size=, compile_backend="torch-compile", model_kwargs={"num_hidden_layers": 2}, # test with smaller model configuration attn_backend="flashinfer", # choose between "triton" and "flashinfer" @@ -249,28 +246,207 @@ llm = LLM( ``` +Please consult the [AutoDeploy `LLM` API](../../tensorrt_llm/_torch/auto_deploy/llm.py) and the +[`AutoDeployConfig` class](../../tensorrt_llm/_torch/auto_deploy/llm_args.py) +for more detail on how AutoDeploy is configured via the `**kwargs` of the `LLM` API. + +### Expert Configuration of LLM API + +For expert TensorRT-LLM users, we also expose the full set of [`LlmArgs`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py) +*at your own risk* (the argument list diverges from TRT-LLM's argument list): + +
+Click to expand for more details on using LlmArgs directly + +- All config fields that are used by the AutoDeploy core pipeline (i.e. the `InferenceOptimizer`) are + _exclusively_ exposed in the [`AutoDeployConfig` class](../../tensorrt_llm/_torch/auto_deploy/llm_args.py). + Please make sure to refer to those first. +- For expert users we expose the full set of [`LlmArgs`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py) + that can be used to configure the [AutoDeploy `LLM` API](../../tensorrt_llm/_torch/auto_deploy/llm.py) including runtime options. +- Note that some fields in the full [`LlmArgs`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py) + object are overlapping, duplicated, and/or _ignored_ in AutoDeploy, particularly arguments + pertaining to configuring the model itself since AutoDeploy's model ingestion+optimize pipeline + significantly differs from the default manual workflow in TensorRT-LLM. +- However, with the proper care the full [`LlmArgs`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py) + objects can be used to configure advanced runtime options in TensorRT-LLM. +- Note that any valid field can be simply provided as keyword argument ("`**kwargs`") to the + [AutoDeploy `LLM` API](../../tensorrt_llm/_torch/auto_deploy/llm.py). +
-For more examples on TRT-LLM LLM API, visit [`this page`](https://nvidia.github.io/TensorRT-LLM/examples/llm_api_examples.html). +### Expert Configuration of `build_and_run_ad.py` -______________________________________________________________________ +For expert users, `build_and_run_ad.py` provides advanced configuration capabilities through a flexible argument parser powered by PyDantic Settings and OmegaConf. You can use dot notation for CLI arguments, provide multiple YAML configuration files, and leverage sophisticated configuration precedence rules to create complex deployment configurations. -## Roadmap +
+Click to expand for detailed configuration examples -1. **Model Coverage:** +#### CLI Arguments with Dot Notation - - Expand support for additional LLM variants and features: - - LoRA - - Speculative Decoding - - Model specialization for disaggregated serving +The script supports flexible CLI argument parsing using dot notation to modify nested configurations dynamically. You can target any field in both the [`ExperimentConfig`](./build_and_run_ad.py) and nested [`AutoDeployConfig`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py)/[`LlmArgs`](../../tensorrt_llm/_torch/auto_deploy/llm_args.) objects: -1. **Performance Optimization:** +```bash +# Configure model parameters +# NOTE: config values like num_hidden_layers are automatically resolved into the appropriate nested +# dict value ``{"args": {"model_kwargs": {"num_hidden_layers": 10}}}`` although not explicitly +# specified as CLI arg +python build_and_run_ad.py \ + --model "meta-llama/Meta-Llama-3.1-8B-Instruct" \ + --args.model-kwargs.num-hidden-layers=10 \ + --args.model-kwargs.hidden-size=2048 \ + --args.tokenizer-kwargs.padding-side=left - - Enhance inference speed and efficiency with: - - MoE fusion and all-reduce fusion techniques - - Reuse of TRT-LLM PyTorch operators for greater efficiency +# Configure runtime and backend settings +python build_and_run_ad.py \ + --model "TinyLlama/TinyLlama-1.1B-Chat-v1.0" \ + --args.world-size=2 \ + --args.compile-backend=torch-opt \ + --args.attn-backend=flashinfer -______________________________________________________________________ +# Configure prompting and benchmarking +python build_and_run_ad.py \ + --model "microsoft/phi-4" \ + --prompt.batch-size=4 \ + --prompt.sp-kwargs.max-tokens=200 \ + --prompt.sp-kwargs.temperature=0.7 \ + --benchmark.enabled=true \ + --benchmark.bs=8 \ + --benchmark.isl=1024 +``` + +#### YAML Configuration Files + +Both [`ExperimentConfig`](./build_and_run_ad.py) and [`AutoDeployConfig`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py)/[`LlmArgs`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py) inherit from [`DynamicYamlMixInForSettings`](../../tensorrt_llm/_torch/auto_deploy/utils/_config.py), enabling you to provide multiple YAML configuration files that are automatically deep-merged at runtime. + +Create a YAML configuration file (e.g., `my_config.yaml`): + +```yaml +# my_config.yaml +args: + model_kwargs: + num_hidden_layers: 12 + hidden_size: 1024 + world_size: 4 + compile_backend: torch-compile + attn_backend: triton + max_seq_len: 2048 + max_batch_size: 16 + transforms: + sharding: + strategy: auto + quantization: + enabled: false + +prompt: + batch_size: 8 + sp_kwargs: + max_tokens: 150 + temperature: 0.8 + top_k: 50 + +benchmark: + enabled: true + num: 20 + bs: 4 + isl: 1024 + osl: 256 +``` + +Create an additional override file (e.g., `production.yaml`): + +```yaml +# production.yaml +args: + world_size: 8 + compile_backend: torch-opt + max_batch_size: 32 + +benchmark: + enabled: false +``` + +Then use these configurations: + +```bash +# Using single YAML config +python build_and_run_ad.py \ + --model "meta-llama/Meta-Llama-3.1-8B-Instruct" \ + --yaml-configs my_config.yaml + +# Using multiple YAML configs (deep merged in order, later files have higher priority) +python build_and_run_ad.py \ + --model "meta-llama/Meta-Llama-3.1-8B-Instruct" \ + --yaml-configs my_config.yaml production.yaml + +# Targeting nested AutoDeployConfig with separate YAML +python build_and_run_ad.py \ + --model "meta-llama/Meta-Llama-3.1-8B-Instruct" \ + --yaml-configs my_config.yaml \ + --args.yaml-configs autodeploy_overrides.yaml +``` + +#### Configuration Precedence and Deep Merging + +The configuration system follows a strict precedence order where higher priority sources override lower priority ones: + +1. **CLI Arguments** (highest priority) - Direct command line arguments +1. **YAML Configs** - Files specified via `--yaml-configs` and `--args.yaml-configs` +1. **Default Settings** (lowest priority) - Built-in defaults from the config classes + +**Deep Merging**: Unlike simple overwriting, deep merging intelligently combines nested dictionaries recursively. For example: + +```yaml +# Base config +args: + model_kwargs: + num_hidden_layers: 10 + hidden_size: 1024 + max_seq_len: 2048 +``` + +```yaml +# Override config +args: + model_kwargs: + hidden_size: 2048 # This will override + # num_hidden_layers: 10 remains unchanged + world_size: 4 # This gets added +``` + +**Nested Config Behavior**: When using nested configurations, outer YAML configs become init settings for inner objects, giving them higher precedence: + +```bash +# The outer yaml-configs affects the entire ExperimentConfig +# The inner args.yaml-configs affects only the AutoDeployConfig +python build_and_run_ad.py \ + --model "meta-llama/Meta-Llama-3.1-8B-Instruct" \ + --yaml-configs experiment_config.yaml \ + --args.yaml-configs autodeploy_config.yaml \ + --args.world-size=8 # CLI override beats both YAML configs +``` + +#### Built-in Default Configuration + +Both [`AutoDeployConfig`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py) and [`LlmArgs`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py) classes automatically load a built-in [`default.yaml`](../../tensorrt_llm/_torch/auto_deploy/config/default.yaml) configuration file that provides sensible defaults for the AutoDeploy inference optimizer pipeline. This file is specified in the [`_get_config_dict()`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py) function and defines default transform configurations for graph optimization stages. + +The built-in defaults are automatically merged with your configurations at the lowest priority level, ensuring that your custom settings always override the defaults. You can inspect the current default configuration to understand the baseline transform pipeline: + +```bash +# View the default configuration +cat tensorrt_llm/_torch/auto_deploy/config/default.yaml + +# Override specific transform settings +python build_and_run_ad.py \ + --model "TinyLlama/TinyLlama-1.1B-Chat-v1.0" \ + --args.transforms.export-to-gm.strict=true +``` + +
+ +## Roadmap + +Check out our [Github Project Board](https://github.com/orgs/NVIDIA/projects/83) to learn more about +the current progress in AutoDeploy and where you can help. ## Disclaimer diff --git a/examples/auto_deploy/build_and_run_ad.py b/examples/auto_deploy/build_and_run_ad.py index 414074ef9a15..35879834db0c 100644 --- a/examples/auto_deploy/build_and_run_ad.py +++ b/examples/auto_deploy/build_and_run_ad.py @@ -1,13 +1,23 @@ """Main entrypoint to build, test, and prompt AutoDeploy inference models.""" -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, Iterator, List, Optional, Union import torch -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator -from pydantic_settings import BaseSettings, CliApp, CliImplicitFlag - -from tensorrt_llm._torch.auto_deploy import LLM, DemoLLM, LlmArgs -from tensorrt_llm._torch.auto_deploy.llm_args import _try_decode_dict_with_str_values +from omegaconf import OmegaConf +from pydantic import BaseModel, Field, field_validator, model_validator +from pydantic_settings import ( + BaseSettings, + CliApp, + CliImplicitFlag, + CliUnknownArgs, + SettingsConfigDict, +) + +from tensorrt_llm._torch.auto_deploy import LLM, AutoDeployConfig, DemoLLM +from tensorrt_llm._torch.auto_deploy.utils._config import ( + DynamicYamlMixInForSettings, + deep_merge_dicts, +) from tensorrt_llm._torch.auto_deploy.utils.benchmark import benchmark, store_benchmark_results from tensorrt_llm._torch.auto_deploy.utils.logger import ad_logger from tensorrt_llm.llmapi.llm import RequestOutput @@ -18,7 +28,11 @@ class PromptConfig(BaseModel): - """Prompt configuration.""" + """Prompt configuration. + + This configuration class can be used for this example script to configure the example prompts + and the sampling parameters. + """ batch_size: int = Field(default=2, description="Number of queries") queries: Union[str, List[str]] = Field( @@ -54,13 +68,16 @@ def model_post_init(self, __context: Any): @classmethod def validate_sp_kwargs(cls, sp_kwargs): """Insert desired defaults for sampling params and try parsing string values as JSON.""" - sp_kwargs = {**cls.model_fields["sp_kwargs"].default_factory(), **sp_kwargs} - sp_kwargs = _try_decode_dict_with_str_values(sp_kwargs) - return sp_kwargs + default = cls.model_fields["sp_kwargs"].get_default(call_default_factory=True) + return deep_merge_dicts(default, sp_kwargs) class BenchmarkConfig(BaseModel): - """Benchmark configuration.""" + """Benchmark configuration. + + This configuration class can be used for this example script to configure the simple + benchmarking we run at the end of the script. + """ enabled: bool = Field(default=False, description="If true, run simple benchmark") num: int = Field(default=10, ge=1, description="By default run 10 times and get average") @@ -73,18 +90,26 @@ class BenchmarkConfig(BaseModel): ) -class ExperimentConfig(BaseSettings): - """Experiment Configuration based on Pydantic BaseModel.""" +class ExperimentConfig(DynamicYamlMixInForSettings, BaseSettings): + """Experiment Configuration for the example script. - model_config = ConfigDict( + This configuration aggregates all relevant configurations for this example script. It is also + used to auto-generate the CLI interface. + """ + + model_config = SettingsConfigDict( extra="forbid", cli_kebab_case=True, + cli_ignore_unknown_args=True, + nested_model_default_partial_update=True, ) + extra_cli_args: CliUnknownArgs ### CORE ARGS ################################################################################## - # The main LLM arguments - contains model, tokenizer, backend configs, etc. - args: LlmArgs = Field( - description="The main LLM arguments containing model, tokenizer, backend configs, etc." + # The main AutoDeploy arguments - contains model, tokenizer, backend configs, etc. + args: AutoDeployConfig = Field( + description="The main AutoDeploy arguments containing model, tokenizer, backend configs, etc. " + "Please check `tensorrt_llm._torch.auto_deploy.llm_args.AutoDeployConfig` for more details." ) # Optional model field for convenience - if provided, will be used to initialize args.model @@ -119,16 +144,50 @@ def setup_args_from_model(cls, data: Dict) -> Dict: data["args"]["model"] = data["model"] return data + @model_validator(mode="before") + @classmethod + def process_extra_cli_args(cls, data: Dict) -> Dict: + """Process extra CLI args. + + This model validator enables the user to provide additional CLI args that may not be + auto-generated by the CLI app. A common use case for this would to modify graph transforms + dynamically via CLI arguments. + + For example, the user can provide a CLI argument for raw dictionaries like this, e.g., for + ``model_kwargs``: ``--args.model-kwargs.num-hidden-layers=10``. + """ + # build a clean dotlist: ["a.b=1","c.d.e=foo",…] + raw: List[str] = data.pop("extra_cli_args", []) + dotlist = [] + it: Iterator[str] = iter(raw) + for tok in it: + if not tok.startswith("--"): + continue + body = tok[2:] + if "=" in body: + body, val = body.split("=", 1) + else: + # flag + separate value + val = next(it, None) + # ensure kebab-case is converted to snake_case + dotlist.append(f"{body.replace('-', '_')}={val}") + + return deep_merge_dicts(data, OmegaConf.from_dotlist(dotlist)) + @field_validator("model", mode="after") @classmethod def sync_model_with_args(cls, model_value, info): - args: LlmArgs = info.data["args"] - return args.model if args is not None else model_value + if "args" not in info.data: + return model_value + args: AutoDeployConfig = info.data["args"] + return args.model @field_validator("prompt", mode="after") @classmethod def sync_prompt_batch_size_with_args_max_batch_size(cls, prompt: PromptConfig, info): - args: LlmArgs = info.data["args"] + if "args" not in info.data: + return prompt + args: AutoDeployConfig = info.data["args"] if args.max_batch_size < prompt.batch_size: args.max_batch_size = prompt.batch_size return prompt @@ -136,7 +195,9 @@ def sync_prompt_batch_size_with_args_max_batch_size(cls, prompt: PromptConfig, i @field_validator("benchmark", mode="after") @classmethod def adjust_args_for_benchmark(cls, benchmark: BenchmarkConfig, info): - args: LlmArgs = info.data["args"] + if "args" not in info.data: + return benchmark + args: AutoDeployConfig = info.data["args"] if benchmark.enabled: # propagate benchmark settings to args args.max_batch_size = max(benchmark.bs, args.max_batch_size) @@ -151,7 +212,6 @@ def build_llm_from_config(config: ExperimentConfig) -> LLM: "demollm": DemoLLM, "trtllm": LLM, } - ad_logger.info(f"{config.args._parallel_config=}") llm = llm_lookup[config.args.runtime](**config.args.to_dict()) return llm diff --git a/examples/auto_deploy/build_and_run_flux.py b/examples/auto_deploy/build_and_run_flux.py index 4170974b4532..a2a647764f31 100644 --- a/examples/auto_deploy/build_and_run_flux.py +++ b/examples/auto_deploy/build_and_run_flux.py @@ -6,7 +6,7 @@ from diffusers import DiffusionPipeline from tensorrt_llm._torch.auto_deploy.compile import compile_and_capture -from tensorrt_llm._torch.auto_deploy.transformations.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm from tensorrt_llm._torch.auto_deploy.transformations.library.fusion import fuse_gemms from tensorrt_llm._torch.auto_deploy.transformations.library.quantization import quantize from tensorrt_llm._torch.auto_deploy.utils.logger import ad_logger @@ -138,10 +138,10 @@ def main(): if args.restore_from: quant_state_dict = model.state_dict() - gm = quantize(gm, {}).to("cuda") + quantize(gm, {}).to("cuda") gm.load_state_dict(quant_state_dict, strict=False) - gm = fuse_gemms(gm) + fuse_gemms(gm) gm = compile_and_capture(gm, backend="torch-opt", args=(), kwargs=flux_kwargs) diff --git a/examples/constraints.txt b/examples/constraints.txt index bec7f2270a3f..5a14c8a137ca 100644 --- a/examples/constraints.txt +++ b/examples/constraints.txt @@ -1,3 +1,3 @@ -tensorrt_llm==1.0.0rc3 +tensorrt_llm==1.0.0rc5 evaluate~=0.4.1 rouge_score~=0.1.2 diff --git a/examples/disaggregated/README.md b/examples/disaggregated/README.md index 120706dd01af..5f34cc810a5c 100644 --- a/examples/disaggregated/README.md +++ b/examples/disaggregated/README.md @@ -1,33 +1,50 @@ -# TRT-LLM Disaggregated Serving +# Disaggregated Serving -To run TRT-LLM in disaggregated mode, you must first launch context (prefill) and generation (decode) servers using `trtllm-serve`. +To run TensorRT-LLM in disaggregated mode, you must first launch context (prefill) and generation (decode) servers using `trtllm-serve`. -## Launching context and generation servers using multiple independent `trtllm-serve` commands +## Launching disaggregated servers locally on single node + +We use the `cache_transceiver_config` configuration to set up disaggregated serving, which includes the following parameters: + +```yaml +cache_transceiver_config: + backend: + max_tokens_in_buffer: +``` + +`backend` specifies the communication backend for transferring the kvCache, valid options include `DEFAULT`,`UCX`, `NIXL`, and `MPI`, the default backend is UCX. + +`max_tokens_in_buffer` defines the buffer size for kvCache transfers, it is recommended to set this value greater than or equal to the maximum ISL (Input Sequence Length) of all requests for optimal performance. You can use multiple `trtllm-serve` commands to launch the context and generation servers that will be used for disaggregated serving. For example, you could launch two context servers and one generation servers as follows: -``` -echo -e "disable_overlap_scheduler: True\ncache_transceiver_config:\n max_num_tokens: 2048" > context_extra-llm-api-config.yml -echo -e "cache_transceiver_config:\n max_num_tokens: 2048" > gen_extra-llm-api-config.yml +```bash +# Generate context_extra-llm-api-config.yml +# Overlap scheduler for context servers are disabled because it's not supported for disaggregated context servers yet +echo -e "disable_overlap_scheduler: True\ncache_transceiver_config:\n backend: UCX\n max_tokens_in_buffer: 2048" > context_extra-llm-api-config.yml -export TRTLLM_USE_UCX_KVCACHE=1 -#Context servers +# Start context servers CUDA_VISIBLE_DEVICES=0 trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8001 --backend pytorch --extra_llm_api_options ./context_extra-llm-api-config.yml &> log_ctx_0 & CUDA_VISIBLE_DEVICES=1 trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8002 --backend pytorch --extra_llm_api_options ./context_extra-llm-api-config.yml &> log_ctx_1 & -#Generation servers + +# Generate gen_extra-llm-api-config.yml +echo -e "cache_transceiver_config:\n backend: UCX\n max_tokens_in_buffer: 2048" > gen_extra-llm-api-config.yml + +# Start generation servers CUDA_VISIBLE_DEVICES=2 trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8003 --backend pytorch --extra_llm_api_options ./gen_extra-llm-api-config.yml &> log_gen_0 & ``` + Once the context and generation servers are launched, you can launch the disaggregated server, which will accept requests from clients and do the orchestration between context and generation servers. The disaggregated server can be launched with: -``` +```bash trtllm-serve disaggregated -c disagg_config.yaml ``` where `disagg_config.yaml` contains information about the context and generation servers. For the current example, it would look like: -``` +```yaml hostname: localhost port: 8000 backend: pytorch @@ -42,13 +59,19 @@ generation_servers: - "localhost:8003" ``` -Clients can then send requests to the disaggregated server at `localhost:8000`, which is an OpenAI compatible endpoint. +Clients can then send requests to the disaggregated server at `localhost:8000`, which is an OpenAI API compatible endpoint. + +## Launching disaggregated servers on SLURM clusters + +Refer to [Disaggregated Inference Benchmark Scripts](./slurm/). ## Sending requests to the disaggregated server Once the context, generation and disaggregated servers are launched, you can send requests to the disaggregated server using curl: -``` -curl http://localhost:8000/v1/completions -H "Content-Type: application/json" -d '{ +```bash +curl http://localhost:8000/v1/completions \ + -H "Content-Type: application/json" \ + -d '{ "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "prompt": "NVIDIA is a great company because", "max_tokens": 16, @@ -64,25 +87,28 @@ python3 ./clients/disagg_client.py -c disagg_config.yaml -p ./clients/prompts.js Currently, trtllm supports dynamic addition and removal of servers by leveraging ETCD. To enable this feature, you should start the context and generation servers with an additional flag ```--metadata_server_config_file``` and ```--server_role```. Before launching the context and generation servers, you should first start the ETCD server. By default, the ETCD server listens for client requests at ```localhost:2379```. -``` +```bash etcd ``` After this, you can enable the dynamic scaling feature for the use case above as follows: -``` +```bash export TRTLLM_USE_UCX_KVCACHE=1 -#Context servers + +# Context servers CUDA_VISIBLE_DEVICES=0 trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8001 --backend pytorch --server_role CONTEXT --extra_llm_api_options ./context_extra-llm-api-config.yml --metadata_server_config_file ./metadata_config.yml &> log_ctx_0 & CUDA_VISIBLE_DEVICES=1 trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8002 --backend pytorch --server_role CONTEXT --extra_llm_api_options ./context_extra-llm-api-config.yml --metadata_server_config_file ./metadata_config.yml &> log_ctx_1 & -#Generation servers + +# Generation servers CUDA_VISIBLE_DEVICES=2 trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8003 --backend pytorch --server_role GENERATION --extra_llm_api_options ./gen_extra-llm-api-config.yml --metadata_server_config_file ./metadata_config.yml &> log_gen_0 & ``` + As for the disaggregated server, you should also specify the --metadata_server_config_file like the following -``` +```bash trtllm-serve disaggregated -c disagg_config.yaml -m ./metadata_config.yml ``` The metadata_config file looks like -``` +```yaml hostname: "localhost" port: 2379 health_check_timeout: 5.0 @@ -94,10 +120,14 @@ The ```hostname``` and ```port``` must match those used when starting the ETCD s ### Dynamically adding servers Users can add servers by directly launching them with trtllm-serve. For example, you can start an additional generation server as follows: +```bash +CUDA_VISIBLE_DEVICES=3 trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 \ + --host localhost --port 8004 \ + --backend pytorch --server_role GENERATION \ + --extra_llm_api_options ./gen_extra-llm-api-config.yml \ + --metadata_server_config_file ./metadata_config.yml &> log_gen_0 & ``` -CUDA_VISIBLE_DEVICES=3 trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8004 --backend pytorch --server_role GENERATION --extra_llm_api_options ./gen_extra-llm-api-config.yml --metadata_server_config_file ./metadata_config.yml &> log_gen_0 & -``` -Trtllm will automatically register any newly launched server with the ETCD server, allowing the router to send new requests to the added server. +TensorRT-LLM will automatically register any newly launched server with the ETCD server, allowing the router to send new requests to the added server. ### Dynamically removing servers @@ -106,7 +136,7 @@ When removing servers, special attention is required in the current version. You ## Launching context and generation servers using MPI (Deprecated) One can also launch all context and generation servers using MPI. This can be done by issuing the following command: -``` +```bash export TRTLLM_USE_MPI_KVCACHE=1 mpirun -n trtllm-serve disaggregated_mpi_worker -c disagg_config.yaml ``` @@ -128,6 +158,8 @@ context_servers: pipeline_parallel_size: 1 kv_cache_config: free_gpu_memory_fraction: 0.9 + cache_transceiver_config: + backend: UCX urls: - "localhost:8001" - "localhost:8002" @@ -135,11 +167,17 @@ generation_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 + cache_transceiver_config: + backend: UCX urls: - "localhost:8003" ``` Once the context and generation servers are launched, you can again launch the disaggregated server with -``` +```bash trtllm-serve disaggregated -c disagg_config.yaml ``` + +## Know Issues + +The MPI communication backend for kvCache transfer has been deprecated and may not be supported in the future. When using the MPI backend, the environment variable `TRTLLM_USE_MPI_KVCACHE=1` should be set to avoid conflicts between mpi4py and kvCache transfer. diff --git a/examples/disaggregated/disagg_config.yaml b/examples/disaggregated/disagg_config.yaml index 6d5314f235c2..ae72c1b074e0 100644 --- a/examples/disaggregated/disagg_config.yaml +++ b/examples/disaggregated/disagg_config.yaml @@ -10,11 +10,15 @@ context_servers: pipeline_parallel_size: 1 kv_cache_config: free_gpu_memory_fraction: 0.2 + cache_transceiver_config: + backend: "default" urls: - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 + cache_transceiver_config: + backend: "default" urls: - "localhost:8002" diff --git a/docs/source/scripts/disaggregated/README.md b/examples/disaggregated/slurm/README.md similarity index 84% rename from docs/source/scripts/disaggregated/README.md rename to examples/disaggregated/slurm/README.md index ed21b998ddd2..a81607b8bd41 100644 --- a/docs/source/scripts/disaggregated/README.md +++ b/examples/disaggregated/slurm/README.md @@ -81,13 +81,14 @@ This script orchestrates the execution of the benchmark client. It waits for the ## Workflow -1. The user runs `./submit.sh`. -2. `submit.sh` submits one or more jobs to SLURM by calling `sbatch disaggr_torch.slurm` with different parameters. -3. For each job, SLURM allocates resources and runs `disaggr_torch.slurm`. -4. `disaggr_torch.slurm` runs `gen_yaml.py` to create a `config.yaml`. -5. `disaggr_torch.slurm` uses `srun` to launch `start_worker.sh` on all nodes, starting the MPI workers. -6. `disaggr_torch.slurm` starts the main `trtllm-serve` process. -7. `disaggr_torch.slurm` runs `run_benchmark.sh` which waits for the server to be ready. -8. `run_benchmark.sh` executes the benchmark for each concurrency level specified. -9. After the benchmark, `run_benchmark.sh` and `disaggr_torch.slurm` attempt to kill the server and worker processes. -10. Logs for each run are stored in a subdirectory specified by the `sub_file` parameter. +1. Make sure that SLURM parameters are correctly set in `disaggr_torch.slurm`. +2. The user runs `./submit.sh`. +3. `submit.sh` submits one or more jobs to SLURM by calling `sbatch disaggr_torch.slurm` with different parameters. +4. For each job, SLURM allocates resources and runs `disaggr_torch.slurm`. +5. `disaggr_torch.slurm` runs `gen_yaml.py` to create a `config.yaml`. +6. `disaggr_torch.slurm` uses `srun` to launch `start_worker.sh` on all nodes, starting the MPI workers. +7. `disaggr_torch.slurm` starts the main `trtllm-serve` process. +8. `disaggr_torch.slurm` runs `run_benchmark.sh` which waits for the server to be ready. +9. `run_benchmark.sh` executes the benchmark for each concurrency level specified. +10. After the benchmark, `run_benchmark.sh` and `disaggr_torch.slurm` attempt to kill the server and worker processes. +11. Logs for each run are stored in a subdirectory specified by the `sub_file` parameter. diff --git a/examples/wide_ep/slurm_scripts/disaggr_torch.slurm b/examples/disaggregated/slurm/disaggr_torch.slurm similarity index 83% rename from examples/wide_ep/slurm_scripts/disaggr_torch.slurm rename to examples/disaggregated/slurm/disaggr_torch.slurm index 4d3e6d801210..941978a56565 100644 --- a/examples/wide_ep/slurm_scripts/disaggr_torch.slurm +++ b/examples/disaggregated/slurm/disaggr_torch.slurm @@ -4,19 +4,21 @@ #SBATCH --ntasks-per-node=4 #SBATCH --partition=${partition} # add your partition here #SBATCH --account=${account} # add your account here -#SBATCH --time=01:00:00 +#SBATCH --time=02:00:00 #SBATCH --job-name=${job_name} # add your job name here isl=1024 osl=1024 -multi_round=1 +multi_round=10 gen_yaml_file=gen_yaml.py +streaming=true container_image=${container_image} # add your container image here mount_dir=${mount_dir} # add your mount directory here -workdir=${mount_dir}/bench-large-ep/slurm_scripts/ +workdir=${workdir} # add your path to the slurm scripts here model_dir=${model_dir} # add your model directory here -logdir=${workdir}/bm_20250703_deepseek-r1-${isl}-${osl}/ -streaming=false + +mounts=${mount_dir}:${mount_dir} +logdir=${workdir}/benchmark-${isl}-${osl}/ mkdir -p ${logdir} container_name=disaggr-test @@ -36,7 +38,7 @@ eplb_num_slots=${12} mtp_size=${13} concurrency=${14} -sub_dir=${logdir}/dep${gen_tp_size}_concurrency${concurrency}_eplb${eplb_num_slots}_mtp${mtp_size} +full_logdir=${logdir}/dep${gen_tp_size}_concurrency${concurrency}_eplb${eplb_num_slots}_mtp${mtp_size} ctx_gpus=$((num_ctx_servers * ctx_tp_size)) gen_gpus=$((num_gen_servers * gen_tp_size)) @@ -47,22 +49,23 @@ enable_pdl=false if [ "${gen_enable_attention_dp}" = "false" ]; then enable_pdl=true echo "enable_pdl: ${enable_pdl}" - sub_dir=${logdir}/tep${gen_tp_size}_concurrency${concurrency}_eplb${eplb_num_slots}_mtp${mtp_size} + full_logdir=${logdir}/tep${gen_tp_size}_concurrency${concurrency}_eplb${eplb_num_slots}_mtp${mtp_size} fi - -full_logdir=${sub_dir} mkdir -p ${full_logdir} +nsys_on="" +# nsys_on=${full_logdir} # Uncomment this line to enable Nsys profiling + # start the container srun -l --container-image=${container_image} \ --container-name=${container_name} \ - --container-mounts=${mount_dir}:${mount_dir} \ + --container-mounts=${mounts} \ --mpi=pmix \ echo "Container up." # generate the yaml file srun -l --container-name=${container_name} \ - --container-mounts=${mount_dir}:${mount_dir} \ + --container-mounts=${mounts} \ --mpi=pmix --overlap \ python3 ${workdir}/${gen_yaml_file} --config ${full_logdir}/config.yaml \ --model ${model_dir} \ @@ -87,33 +90,32 @@ echo "server host name: $hostname_value" # try to kill the server and workers srun -l --container-name=${container_name} \ - --container-mounts=${mount_dir}:${mount_dir} \ + --container-mounts=${mounts} \ --mpi=pmix --overlap \ pkill -f "trtllm-serve" || true -nsys_on="" -# nsys_on=${full_logdir} - # start the workers srun -l --container-name=${container_name} \ - --container-mounts=${mount_dir}:${mount_dir} \ + --container-mounts=${mounts} \ --mpi=pmix --overlap \ bash ${workdir}/start_worker.sh ${full_logdir}/config.yaml "${concurrency}" "${enable_pdl}" ${ctx_gpus} ${nsys_on} &> ${full_logdir}/output_workers.log & + # start the server srun -l --container-name=${container_name} \ - --container-mounts=${mount_dir}:${mount_dir} \ + --container-mounts=${mounts} \ --mpi=pmix --overlap -N 1 -n 1 \ -w ${hostname_value} \ bash ${workdir}/start_server.sh ${full_logdir}/config.yaml &> ${full_logdir}/output_server.log & + # start benchmarking srun -l --container-name=${container_name} \ - --container-mounts=${mount_dir}:${mount_dir} \ + --container-mounts=${mounts} \ --mpi=pmix --overlap -N 1 -n 1 \ bash ${workdir}/run_benchmark.sh ${isl} ${osl} ${multi_round} ${model_dir} "${concurrency}" ${streaming} ${full_logdir}/ > ${full_logdir}/benchmark.log 2>&1 # try to kill the server and workers srun -l --container-name=${container_name} \ - --container-mounts=${mount_dir}:${mount_dir} \ + --container-mounts=${mounts} \ --mpi=pmix --overlap \ kill -9 $(ps aux | grep '[t]rtllm-serve' | awk '{print $2}') >/dev/null 2>&1 || true wait diff --git a/examples/wide_ep/slurm_scripts/gen_yaml.py b/examples/disaggregated/slurm/gen_yaml.py similarity index 98% rename from examples/wide_ep/slurm_scripts/gen_yaml.py rename to examples/disaggregated/slurm/gen_yaml.py index 121f614d8700..0ef7a3ecf503 100644 --- a/examples/wide_ep/slurm_scripts/gen_yaml.py +++ b/examples/disaggregated/slurm/gen_yaml.py @@ -182,7 +182,8 @@ def gen_config_file(config_path: str, 'disable_overlap_scheduler': True, 'kv_cache_dtype': 'fp8', 'cache_transceiver_config': { - 'max_num_tokens': 4608, + 'backend': 'default', + 'max_tokens_in_buffer': 8320, }, }, 'generation_servers': { @@ -203,7 +204,8 @@ def gen_config_file(config_path: str, 'kv_cache_dtype': 'fp8', 'moe_backend': gen_moe_backend, 'cache_transceiver_config': { - 'max_num_tokens': 4608, + 'backend': 'default', + 'max_tokens_in_buffer': 8320, }, } } diff --git a/examples/wide_ep/slurm_scripts/run_benchmark.sh b/examples/disaggregated/slurm/run_benchmark.sh similarity index 100% rename from examples/wide_ep/slurm_scripts/run_benchmark.sh rename to examples/disaggregated/slurm/run_benchmark.sh diff --git a/examples/disaggregated/slurm/slurm_populate_urls.py b/examples/disaggregated/slurm/slurm_populate_urls.py deleted file mode 100644 index abe8122dbe56..000000000000 --- a/examples/disaggregated/slurm/slurm_populate_urls.py +++ /dev/null @@ -1,164 +0,0 @@ -import argparse -import os -import re - -import yaml - -# Parse command line arguments -parser = argparse.ArgumentParser( - description='Update YAML configuration with SLURM node information.') -parser.add_argument( - '--nodelist_env_var', - type=str, - default='SLURM_JOB_NODELIST', - help= - 'Name of the env var that provides the list of nodes as dev[7-8,11,13] for example' -) -parser.add_argument( - '--tasks_per_node_env_var', - type=str, - default='SLURM_TASKS_PER_NODE', - help= - 'Name of the env var that provides the tasks per node as 8(x3),2 for example' -) -parser.add_argument('--disagg_server_port', - type=int, - default=8000, - help='The port to use for disagg server') -parser.add_argument('--worker_start_port', - type=int, - default=8001, - help='The starting port to use for workers') -parser.add_argument('--input_yaml', - type=str, - default='config.yaml', - help='Path to the input YAML file') -parser.add_argument('--output_yaml', - type=str, - default='output_config.yaml', - help='Path to the output YAML file') -args = parser.parse_args() - -# Parse SLURM_JOB_NODELIST and SLURM_TASKS_PER_NODE from environment variables -print("---") -slurm_job_nodelist = os.getenv(args.nodelist_env_var, '') -if not slurm_job_nodelist: - raise ValueError(f"Environment variable {args.nodelist_env_var} not found.") -print(f"{args.nodelist_env_var}: {slurm_job_nodelist}") -slurm_tasks_per_node = os.getenv(args.tasks_per_node_env_var, '') -if not slurm_tasks_per_node: - raise ValueError( - f"Environment variable {args.tasks_per_node_env_var} not found.") -print(f"{args.tasks_per_node_env_var}: {slurm_tasks_per_node}") -print("---") - -# Generate list of nodes -node_prefix = re.match(r'^[a-zA-Z]+', slurm_job_nodelist).group(0) -node_range = re.search(r'\[(.*?)\]', slurm_job_nodelist).group(1) -nodes = [] -for part in node_range.split(','): - if '-' in part: - start, end = map(int, part.split('-')) - nodes.extend([f"{node_prefix}{i}" for i in range(start, end + 1)]) - else: - nodes.append(f"{node_prefix}{part}") -print(f"Nodes: {nodes}") - -# Generate tasks per node -tasks_per_node = [] -for part in slurm_tasks_per_node.split(','): - if '(x' in part: - count, repeat = map(int, re.findall(r'\d+', part)) - tasks_per_node.extend([count] * repeat) - else: - tasks_per_node.append(int(part)) -print(f"Tasks_per_node: {tasks_per_node}") - -if (len(tasks_per_node) != len(nodes)): - raise ValueError( - f"Number of nodes and tasks per node do not match. Number of nodes: {len(nodes)}, Number of tasks per node: {len(tasks_per_node)}" - ) - -max_tasks_per_node = max(tasks_per_node) -task_nodes = [] -for node, tasks in zip(nodes, tasks_per_node): - task_nodes.extend([node] * tasks) - -print(f"Task nodes: {task_nodes}") -print("---") - - -# Function to generate URLs -def generate_urls(ctx_or_gen, - num_instances, - tensor_parallel_size, - pipeline_parallel_size, - max_task_per_node, - nodes, - task_nodes, - node_to_port, - task_nodes_offset=0): - urls = [] - - for instance in range(num_instances): - tasks_needed = tensor_parallel_size * pipeline_parallel_size - - if (task_nodes_offset + tasks_needed) > len(task_nodes): - print(f"{ctx_or_gen} urls so far: {urls}") - raise ValueError( - f"For {ctx_or_gen} instance {instance}, there are not enough tasks available. task_nodes_offset: {task_nodes_offset}, tasks_needed: {tasks_needed}, len(task_nodes): {len(task_nodes)}" - ) - - # Minimum number of nodes needed for that instance - min_node = (tasks_needed + max_tasks_per_node - 1) / max_tasks_per_node - instance_nodes = set(task_nodes[task_nodes_offset:task_nodes_offset + - tasks_needed]) - if len(instance_nodes) > min_node: - raise ValueError( - f"Tasks for a instance {instance} of {ctx_or_gen} instances use more node than expected. Nodes used: {instance_nodes}, number of nodes expected: {min_node}, max_tasks_per_node: {max_tasks_per_node}" - ) - - node = task_nodes[task_nodes_offset] - port = node_to_port[node] - node_to_port[node] += 1 - task_nodes_offset += tasks_needed - - urls.append(f"{node}:{port}") - - print(f"{ctx_or_gen} urls: {urls}") - return urls, task_nodes_offset - - -# Load the YAML file -with open(args.input_yaml, 'r') as file: - config = yaml.safe_load(file) - -# Keep track of the port number for each node -node_ports = {} -for node in nodes: - node_ports[node] = args.worker_start_port - -# Generate URLs for context_servers and generation_servers -context_urls, task_node_offset = generate_urls( - "ctx", config['context_servers']['num_instances'], - config['context_servers']['tensor_parallel_size'], - config['context_servers']['pipeline_parallel_size'], max_tasks_per_node, - nodes, task_nodes, node_ports) - -generation_urls, _ = generate_urls( - "gen", config['generation_servers']['num_instances'], - config['generation_servers']['tensor_parallel_size'], - config['generation_servers']['pipeline_parallel_size'], max_tasks_per_node, - nodes, task_nodes, node_ports, task_node_offset) - -# Update the YAML configuration -config['hostname'] = nodes[0] -config['port'] = args.disagg_server_port -config['context_servers']['urls'] = context_urls -config['generation_servers']['urls'] = generation_urls - -# Save the updated YAML file -with open(args.output_yaml, 'w') as file: - yaml.safe_dump(config, file, sort_keys=False) - -print("YAML file updated successfully.") diff --git a/examples/wide_ep/slurm_scripts/start_server.sh b/examples/disaggregated/slurm/start_server.sh similarity index 100% rename from examples/wide_ep/slurm_scripts/start_server.sh rename to examples/disaggregated/slurm/start_server.sh diff --git a/examples/wide_ep/slurm_scripts/start_worker.sh b/examples/disaggregated/slurm/start_worker.sh similarity index 100% rename from examples/wide_ep/slurm_scripts/start_worker.sh rename to examples/disaggregated/slurm/start_worker.sh diff --git a/examples/disaggregated/slurm/submit.sh b/examples/disaggregated/slurm/submit.sh new file mode 100644 index 000000000000..8412b3eb754e --- /dev/null +++ b/examples/disaggregated/slurm/submit.sh @@ -0,0 +1,39 @@ +#!/bin/bash + +# !!! +# Make sure that SLURM parameters are correctly set in `disaggr_torch.slurm` before executing this script. +# !!! + +# concurrency 8 +concurrency=8 +ctx_num=1 +total_node_num=8 +ntasks_per_node=4 # 4 GPUs per GB200 node +ntasks=$((total_node_num * ntasks_per_node)) + +# `--segment` makes sure that all nodes are in the same NVLink domain +# disaggr_torch.slurm arguments: +# num_ctx_servers=$1 +# ctx_tp_size=$2 +# ctx_batch_size=$3 +# ctx_max_num_tokens=$4 +# ctx_enable_attention_dp=$5 +# num_gen_servers=$6 +# gen_tp_size=$7 +# gen_batch_size=$8 +# gen_max_num_tokens=$9 +# gen_enable_attention_dp=${10} +# gen_gpu_memory_fraction=${11} +# eplb_num_slots=${12} +# mtp_size=${13} +# concurrency=${14} + +# This command starts a job with 8 nodes, 32 GPUs in total. +# The server will include 4 context workers with DEP4, and 1 generation worker with DEP8. +sbatch --nodes=${total_node_num} \ + --ntasks=${ntasks} \ + --ntasks-per-node=${ntasks_per_node} \ + --gres=gpu:${ntasks_per_node} \ + --segment=${total_node_num} \ + disaggr_torch.slurm \ + ${ctx_num} 4 4 4480 true 1 8 1024 1024 true "0.8" 0 0 "$concurrency" diff --git a/examples/llm-api/README.md b/examples/llm-api/README.md index 98c02d227137..2012406fd4d1 100644 --- a/examples/llm-api/README.md +++ b/examples/llm-api/README.md @@ -40,9 +40,10 @@ python3 quickstart_multimodal.py --model_dir Efficient-Large-Model/NVILA-8B --mo python3 quickstart_advanced.py \ --model_dir meta-llama/Llama-3.1-8B-Instruct \ --spec_decode_algo NGRAM \ - --max_matching_ngram_size=2 \ - --spec_decode_nextn=4 \ - --disable_overlap_scheduler + --spec_decode_max_draft_len 4 \ + --max_matching_ngram_size 2 \ + --disable_overlap_scheduler \ + --disable_kv_cache_reuse ``` ```bash @@ -50,8 +51,8 @@ python3 quickstart_advanced.py \ python3 quickstart_advanced.py \ --model_dir meta-llama/Llama-3.1-8B-Instruct \ --spec_decode_algo draft_target \ - --spec_decode_nextn 5 \ + --spec_decode_max_draft_len 5 \ --draft_model_dir meta-llama/Llama-3.2-1B-Instruct \ - --disable_overlap_scheduler + --disable_overlap_scheduler \ --disable_kv_cache_reuse ``` diff --git a/examples/llm-api/llm_guided_decoding.py b/examples/llm-api/llm_guided_decoding.py index a5e0f89244d3..e5df98e5da3a 100644 --- a/examples/llm-api/llm_guided_decoding.py +++ b/examples/llm-api/llm_guided_decoding.py @@ -7,12 +7,9 @@ def main(): - # Specify the guided decoding backend; xgrammar is supported currently. - llm = LLM( - model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - guided_decoding_backend='xgrammar', - disable_overlap_scheduler=True # Not supported by xgrammar mode - ) + # Specify the guided decoding backend; xgrammar and llguidance are supported currently. + llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", + guided_decoding_backend='xgrammar') # An example from json-mode-eval schema = '{"title": "WirelessAccessPoint", "type": "object", "properties": {"ssid": {"title": "SSID", "type": "string"}, "securityProtocol": {"title": "SecurityProtocol", "type": "string"}, "bandwidth": {"title": "Bandwidth", "type": "string"}}, "required": ["ssid", "securityProtocol", "bandwidth"]}' diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index 4abb501b6d99..5e447e6a0e42 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -50,7 +50,7 @@ def add_llm_args(parser): parser.add_argument('--moe_backend', type=str, default='CUTLASS', - choices=['CUTLASS', 'TRTLLM', 'VANILLA']) + choices=['CUTLASS', 'TRTLLM', 'VANILLA', 'WIDEEP']) parser.add_argument('--enable_attention_dp', default=False, action='store_true') @@ -108,11 +108,8 @@ def add_llm_args(parser): # Speculative decoding parser.add_argument('--spec_decode_algo', type=str, default=None) - parser.add_argument('--spec_decode_nextn', type=int, default=1) - parser.add_argument('--draft_model_dir', - '--eagle_model_dir', - type=str, - default=None) + parser.add_argument('--spec_decode_max_draft_len', type=int, default=1) + parser.add_argument('--draft_model_dir', type=str, default=None) parser.add_argument('--max_matching_ngram_size', type=int, default=5) parser.add_argument('--use_one_model', default=False, action='store_true') @@ -145,10 +142,11 @@ def parse_arguments(): return args -def setup_llm(args): +def setup_llm(args, **kwargs): kv_cache_config = KvCacheConfig( enable_block_reuse=not args.disable_kv_cache_reuse, free_gpu_memory_fraction=args.kv_cache_fraction, + dtype=args.kv_cache_dtype, ) spec_decode_algo = args.spec_decode_algo.upper( @@ -161,23 +159,23 @@ def setup_llm(args): ) spec_config = MTPDecodingConfig( - num_nextn_predict_layers=args.spec_decode_nextn, + num_nextn_predict_layers=args.spec_decode_max_draft_len, use_relaxed_acceptance_for_thinking=args. use_relaxed_acceptance_for_thinking, relaxed_topk=args.relaxed_topk, relaxed_delta=args.relaxed_delta) elif spec_decode_algo == "EAGLE3": spec_config = EagleDecodingConfig( - max_draft_len=args.spec_decode_nextn, + max_draft_len=args.spec_decode_max_draft_len, speculative_model_dir=args.draft_model_dir, eagle3_one_model=args.use_one_model) elif spec_decode_algo == "DRAFT_TARGET": spec_config = DraftTargetDecodingConfig( - max_draft_len=args.spec_decode_nextn, + max_draft_len=args.spec_decode_max_draft_len, speculative_model_dir=args.draft_model_dir) elif spec_decode_algo == "NGRAM": spec_config = NGramDecodingConfig( - max_draft_len=args.spec_decode_nextn, + max_draft_len=args.spec_decode_max_draft_len, max_matching_ngram_size=args.max_matching_ngram_size, is_keep_all=True, is_use_oldest=True, @@ -194,7 +192,6 @@ def setup_llm(args): model=args.model_dir, backend='pytorch', disable_overlap_scheduler=args.disable_overlap_scheduler, - kv_cache_dtype=args.kv_cache_dtype, kv_cache_config=kv_cache_config, attn_backend=args.attention_backend, cuda_graph_config=cuda_graph_config, @@ -222,7 +219,9 @@ def setup_llm(args): speculative_config=spec_config, trust_remote_code=args.trust_remote_code, gather_generation_logits=args.return_generation_logits, - max_beam_width=args.max_beam_width) + max_beam_width=args.max_beam_width, + **kwargs, + ) sampling_params = SamplingParams( max_tokens=args.max_tokens, diff --git a/examples/llm-api/quickstart_multimodal.py b/examples/llm-api/quickstart_multimodal.py index 3859580fcde2..c4d40655d3dc 100644 --- a/examples/llm-api/quickstart_multimodal.py +++ b/examples/llm-api/quickstart_multimodal.py @@ -7,24 +7,56 @@ from tensorrt_llm.inputs import (ALL_SUPPORTED_MULTIMODAL_MODELS, default_multimodal_input_loader) -example_images = [ - "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore.png", - "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png", - "https://huggingface.co/datasets/Sayali9141/traffic_signal_images/resolve/main/61.jpg", -] -example_image_prompts = [ - "Describe the natural environment in the image.", - "Describe the object and the weather condition in the image.", - "Describe the traffic condition on the road in the image.", -] -example_videos = [ - "https://huggingface.co/datasets/Efficient-Large-Model/VILA-inference-demos/resolve/main/OAI-sora-tokyo-walk.mp4", - "https://huggingface.co/datasets/Efficient-Large-Model/VILA-inference-demos/resolve/main/world.mp4", -] -example_video_prompts = [ - "Tell me what you see in the video briefly.", - "Describe the scene in the video briefly.", -] +example_medias_and_prompts = { + "image": { + "media": [ + "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore.png", + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png", + "https://huggingface.co/datasets/Sayali9141/traffic_signal_images/resolve/main/61.jpg", + ], + "prompt": [ + "Describe the natural environment in the image.", + "Describe the object and the weather condition in the image.", + "Describe the traffic condition on the road in the image.", + ] + }, + "video": { + "media": [ + "https://huggingface.co/datasets/Efficient-Large-Model/VILA-inference-demos/resolve/main/OAI-sora-tokyo-walk.mp4", + "https://huggingface.co/datasets/Efficient-Large-Model/VILA-inference-demos/resolve/main/world.mp4", + ], + "prompt": [ + "Tell me what you see in the video briefly.", + "Describe the scene in the video briefly.", + ] + }, + "audio": { + "media": [ + "https://huggingface.co/microsoft/Phi-4-multimodal-instruct/resolve/main/examples/what_is_the_traffic_sign_in_the_image.wav", + "https://huggingface.co/microsoft/Phi-4-multimodal-instruct/resolve/main/examples/what_is_shown_in_this_image.wav", + ], + "prompt": [ + "Transcribe the audio clip into text, please don't add other text.", + "Transcribe the audio clip into text, please don't add other text.", + ] + }, + "image_audio": { + "media": [ + [ + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png", + "https://huggingface.co/microsoft/Phi-4-multimodal-instruct/resolve/main/examples/what_is_shown_in_this_image.wav" + ], + [ + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png", + "https://huggingface.co/microsoft/Phi-4-multimodal-instruct/resolve/main/examples/what_is_shown_in_this_image.wav" + ], + ], + "prompt": [ + "Describe the scene in the image briefly.", + "", + ] + } +} def add_multimodal_args(parser): @@ -34,7 +66,7 @@ def add_multimodal_args(parser): help="Model type.") parser.add_argument("--modality", type=str, - choices=["image", "video"], + choices=["image", "video", "audio", "image_audio"], default="image", help="Media type.") parser.add_argument("--media", @@ -53,11 +85,24 @@ def add_multimodal_args(parser): return parser +def add_lora_args(parser): + parser.add_argument("--load_lora", + default=False, + action='store_true', + help="Whether to load the LoRA model.") + parser.add_argument("--auto_model_name", + type=str, + default=None, + help="The auto model name in TRTLLM repo.") + return parser + + def parse_arguments(): parser = argparse.ArgumentParser( description="Multimodal models with the PyTorch workflow.") parser = add_llm_args(parser) parser = add_multimodal_args(parser) + parser = add_lora_args(parser) args = parser.parse_args() args.disable_kv_cache_reuse = True # kv cache reuse does not work for multimodal, force overwrite @@ -71,11 +116,19 @@ def main(): args = parse_arguments() # set prompts and media to example prompts and images if they are not provided if args.prompt is None: - args.prompt = example_image_prompts if args.modality == "image" else example_video_prompts + args.prompt = example_medias_and_prompts[args.modality]["prompt"] if args.media is None: - args.media = example_images if args.modality == "image" else example_videos + args.media = example_medias_and_prompts[args.modality]["media"] + + lora_config = None + if args.load_lora: + assert args.auto_model_name is not None, "Please provide the auto model name to load LoRA config." + import importlib + models_module = importlib.import_module('tensorrt_llm._torch.models') + model_class = getattr(models_module, args.auto_model_name) + lora_config = model_class.lora_config(args.model_dir) - llm, sampling_params = setup_llm(args) + llm, sampling_params = setup_llm(args, lora_config=lora_config) image_format = args.image_format if args.model_type is not None: @@ -85,7 +138,7 @@ def main(): open(os.path.join(llm._hf_model_dir, 'config.json')))['model_type'] assert model_type in ALL_SUPPORTED_MULTIMODAL_MODELS, f"Unsupported model_type: {model_type}" - device = "cuda" + device = "cpu" inputs = default_multimodal_input_loader(tokenizer=llm.tokenizer, model_dir=llm._hf_model_dir, model_type=model_type, @@ -96,7 +149,16 @@ def main(): num_frames=args.num_frames, device=device) - outputs = llm.generate(inputs, sampling_params) + lora_request = None + if args.load_lora: + lora_request = model_class.lora_request(len(inputs), args.modality, + llm._hf_model_dir) + + outputs = llm.generate( + inputs, + sampling_params, + lora_request=lora_request, + ) for i, output in enumerate(outputs): prompt = args.prompt[i] diff --git a/examples/models/core/deepseek_v3/README.md b/examples/models/core/deepseek_v3/README.md index fa4561066dc0..3434c24f652f 100644 --- a/examples/models/core/deepseek_v3/README.md +++ b/examples/models/core/deepseek_v3/README.md @@ -77,7 +77,7 @@ git clone https://huggingface.co/deepseek-ai/DeepSeek-V3 ## Quick Start ### Run a single inference -To quickly run DeepSeek-V3, [examples/llm-api/quickstart_advanced.py](../pytorch/quickstart_advanced.py): +To quickly run DeepSeek-V3, [examples/llm-api/quickstart_advanced.py](../llm-api/quickstart_advanced.py): ```bash cd examples/llm-api @@ -94,10 +94,10 @@ Prompt: 'The future of AI is', Generated text: ' a topic of great interest and s ``` ### Multi-Token Prediction (MTP) -To run with MTP, use [examples/llm-api/quickstart_advanced.py](../pytorch/quickstart_advanced.py) with additional options, see +To run with MTP, use [examples/llm-api/quickstart_advanced.py](../../../llm-api/quickstart_advanced.py) with additional options, see ```bash cd examples/llm-api -python quickstart_advanced.py --model_dir --spec_decode_algo MTP --spec_decode_nextn N +python quickstart_advanced.py --model_dir --spec_decode_algo MTP --spec_decode_max_draft_len N ``` `N` is the number of MTP modules. When `N` is equal to `0`, which means that MTP is not used (default). When `N` is greater than `0`, which means that `N` MTP modules are enabled. In the current implementation, the weight of each MTP module is shared. @@ -124,7 +124,7 @@ When verifying and receiving draft tokens, there are two ways: ```bash cd examples/llm-api - python quickstart_advanced.py --model_dir --spec_decode_algo MTP --spec_decode_nextn N --use_relaxed_acceptance_for_thinking --relaxed_topk 15 --relaxed_delta 0.5 + python quickstart_advanced.py --model_dir --spec_decode_algo MTP --spec_decode_max_draft_len N --use_relaxed_acceptance_for_thinking --relaxed_topk 15 --relaxed_delta 0.5 ``` ### Long context support @@ -193,7 +193,6 @@ Evaluate the model accuracy using `trtllm-eval`. 1. (Optional) Prepare an advanced configuration file: ```bash cat >./extra-llm-api-config.yml < #### Run a single inference -To quickly run Qwen3, [examples/llm-api/quickstart_advanced.py](../../../pytorch/quickstart_advanced.py): +To quickly run Qwen3, [examples/llm-api/quickstart_advanced.py](../../../llm-api/quickstart_advanced.py): ```bash python3 examples/llm-api/quickstart_advanced.py --model_dir Qwen3-30B-A3B/ --kv_cache_fraction 0.6 diff --git a/examples/models/core/qwen2audio/run.py b/examples/models/core/qwen2audio/run.py index e0d495a67f81..93e161c7e083 100644 --- a/examples/models/core/qwen2audio/run.py +++ b/examples/models/core/qwen2audio/run.py @@ -122,7 +122,8 @@ def get_model(self): num_kv_heads = config["pretrained_config"].get("num_key_value_heads", num_heads) if "kv_cache_type" in config["build_config"]: - kv_cache_type = KVCacheType(config["build_config"]["kv_cache_type"]) + kv_cache_type = KVCacheType.from_string( + config["build_config"]["kv_cache_type"]) else: kv_cache_type = KVCacheType.CONTINUOUS diff --git a/examples/models/core/qwenvl/run.py b/examples/models/core/qwenvl/run.py index a04c2b142e37..06ce341a9a03 100644 --- a/examples/models/core/qwenvl/run.py +++ b/examples/models/core/qwenvl/run.py @@ -118,7 +118,8 @@ def get_model(self): num_kv_heads = config["pretrained_config"].get("num_key_value_heads", num_heads) if "kv_cache_type" in config["build_config"]: - kv_cache_type = KVCacheType(config["build_config"]["kv_cache_type"]) + kv_cache_type = KVCacheType.from_string( + config["build_config"]["kv_cache_type"]) else: kv_cache_type = KVCacheType.CONTINUOUS diff --git a/examples/prompt_lookup/README.md b/examples/ngram/README.md similarity index 54% rename from examples/prompt_lookup/README.md rename to examples/ngram/README.md index ae33e0f6c0a2..60201ce063fd 100644 --- a/examples/prompt_lookup/README.md +++ b/examples/ngram/README.md @@ -1,17 +1,17 @@ -# Prompt-Lookup Speculative Decoding +# NGram Speculative Decoding -This document shows how to build and run a model using Prompt-Lookup speculative decoding (supported as `ASSISTED_GENERATION` in transformers and vLLM, source: [GitHub](https://github.com/apoorvumang/prompt-lookup-decoding/tree/main)) in TensorRT-LLM on single GPU, or single node multiple GPU. +This document shows how to build and run a model using NGram speculative decoding (supported as `ASSISTED_GENERATION` in transformers and vLLM, source: [GitHub](https://github.com/apoorvumang/prompt-lookup-decoding/tree/main)) in TensorRT-LLM on single GPU, or single node multiple GPU. ## Overview -We provide two styles of workflow to run Prompt-Lookup (named V1 and V2 respectively) now. V1 is in TRT workflow and similar to the Draft-Target-Model workflow, running in orchestrator mode and calling `runner.generate()` multiple times to get outputs, which is more flexible for customizing but slightly more overhead. V2 is in pytorch workflow and similar to the Look-Ahead workflow, running in leader mode and calling `runner.generate()` only one time to get outputs, which provides higher performance but fixed process. +We provide two styles of workflow to run NGram (named V1 and V2 respectively) now. V1 is in TRT workflow and similar to the Draft-Target-Model workflow, running in orchestrator mode and calling `runner.generate()` multiple times to get outputs, which is more flexible for customizing but slightly more overhead. V2 is in pytorch workflow and similar to the Look-Ahead workflow, running in leader mode and calling `runner.generate()` only one time to get outputs, which provides higher performance but fixed process. -The Prompt-Lookup has 3 additional hyperparameters that you need to specify to control the process of generation: -- `prompt_lookup_num_tokens`: the maximum number of tokens provided as draft tokens in one iteration, which is usually from 4 to 10 in common usage (default value: 4). Empirically, the larger the value is, the higher acceptance rate but higher overhead is expected at the same time, so the right balance based on the models and application scenarios needs to be found. +The NGram has 3 additional hyperparameters that you need to specify to control the process of generation: +- `max_draft_len`: the maximum number of tokens provided as draft tokens in one iteration, which is usually from 4 to 10 in common usage (default value: 4). Empirically, the larger the value is, the higher acceptance rate but higher overhead is expected at the same time, so the right balance based on the models and application scenarios needs to be found. - `max_matching_ngram_size`: the maximum number of tokens extracted from the tail of the input prompt or generated output as a pattern, which is used to search corresponding draft tokens (default value: 2). Empirically, the larger the value is, the more precise context can be matched from the existed sequence, indicating higher acceptance rate, but the higher probability of miss-match and higher overhead appear, which fall back to normal generation (one token per iteration). - `device_list`: the index list of device(s) to run the model in V1 workflow. The length of it must be the same as the TP size of the draft model engine. For instances, `device_list=[0]` means using tp_size=1 and GPU 0 for the model, `device_list=[4,5,6,7]` means using tp=4 and GPU from 4 to 7 for the model. This parameter is neddless in V2 workflow. -+ For example, the process of getting draft tokens using `prompt_lookup_num_tokens=2` and `max_matching_ngram_size=4` with a sentence `prefix=[..., t1, t2, t3, t4]` is like below: ++ For example, the process of getting draft tokens using `max_draft_len=2` and `max_matching_ngram_size=4` with a sentence `prefix=[..., t1, t2, t3, t4]` is like below: ```Python pattern = prefix[:-2] # pattern=[t3, t4] (length=2) @@ -40,9 +40,9 @@ return None # No any candidate exists + We use an open-source `llama-v2-13B` models in this example. + `--use_paged_context_fmha=enable` must be specified since we need KVcache reuse in this approach. + `--speculative_decoding_mode=draft_tokens_external` must be specified. -+ `--max_draft_len` must be specified larger or equal to `prompt_lookup_num_tokens`. -+ `---prompt_lookup_config` is corresponding configuration of Prompt-Lookup, we can see its usage in [util.py](../util.py). - + As an example, `[10,2,[0]]` means `prompt_lookup_num_tokens=10`, `max_matching_ngram_size=2`, and device of target model is `GPU0`. ++ `--max_draft_len` must be specified as the length maximum of the draft tokens. ++ `--ngram_config` is corresponding configuration of NGram, we can see its usage in [util.py](../util.py). + + As an example, `[10,2,[0]]` means `max_draft_len=10`, `max_matching_ngram_size=2`, and device of target model is `GPU0`. + `--kv_cache_enable_block_reuse` must be specified for this approach. + Only CPP session is supported, so `--use_py_session` must not be specified. + `--num_beams` can not be specified as larger than 1 since beam search is not supported in this approach yet. @@ -50,29 +50,29 @@ return None # No any candidate exists ```bash # Build engine python3 examples/models/core/llama/convert_checkpoint.py \ - --model_dir= \ - --output_dir=./ckpt-target \ - --dtype=float16 + --model_dir \ + --output_dir ./ckpt-target \ + --dtype float16 trtllm-build \ - --checkpoint_dir=./ckpt-target \ - --output_dir=./target-engine \ - --gemm_plugin=float16 \ - --use_paged_context_fmha=enable \ - --speculative_decoding_mode=draft_tokens_external \ - --max_draft_len=10 \ - --max_batch_size=4 \ - --max_input_len=3200 \ - --max_seq_len=4800 + --checkpoint_dir ./ckpt-target \ + --output_dir ./target-engine \ + --gemm_plugin float16 \ + --use_paged_context_fmha enable \ + --speculative_decoding_mode draft_tokens_external \ + --max_draft_len 10 \ + --max_batch_size 4 \ + --max_input_len 3200 \ + --max_seq_len 4800 # Run decoding python3 examples/run.py \ --tokenizer_dir \ --engine_dir ./target-engine \ - --prompt_lookup_config="[10,2,[0]]" \ - --max_output_len=256 \ + --ngram_config "[10,2,[0]]" \ + --max_output_len 256 \ --kv_cache_enable_block_reuse \ - --input_text="How does Draft-Sampling work?" + --input_text "How does Draft-Sampling work?" # Run summarization tasks python examples/summarize.py \ @@ -81,8 +81,8 @@ python examples/summarize.py \ --check_accuracy \ --hf_model_dir \ --engine_dir ./target-engine \ - --batch_size=1 \ - --prompt_lookup_config="[10,2,[0]]" \ + --batch_size 1 \ + --ngram_config "[10,2,[0]]" \ --kv_cache_enable_block_reuse ``` @@ -90,6 +90,8 @@ python examples/summarize.py \ ```bash python3 examples/llm-api/quickstart_advanced.py \ - --max_matching_ngram_size=2 \ - --spec_decode_nextn=4 + --spec_decode_max_draft_len 4 \ + --max_matching_ngram_size 2 \ + --disable_overlap_scheduler \ + --disable_kv_cache_reuse ``` diff --git a/examples/prompt_lookup/requirements.txt b/examples/ngram/requirements.txt similarity index 100% rename from examples/prompt_lookup/requirements.txt rename to examples/ngram/requirements.txt diff --git a/examples/prompt_lookup/run_dtm_pld.py b/examples/ngram/run_dtm_ngram.py similarity index 89% rename from examples/prompt_lookup/run_dtm_pld.py rename to examples/ngram/run_dtm_ngram.py index 559c1e7bbef9..d0cd8687ef86 100644 --- a/examples/prompt_lookup/run_dtm_pld.py +++ b/examples/ngram/run_dtm_ngram.py @@ -23,12 +23,12 @@ from tensorrt_llm.runtime import ModelRunnerCpp -class PLDPool: # Ngrams pool for Prompt-Lookup-Decoding +class NgramPool: # Ngrams pool for Ngram def __init__( self, input_batch_size: int, - prompt_lookup_num_tokens: int, + max_draft_len: int, max_matching_ngram_size: int, end_id: int, max_seq_len: list[int], @@ -36,7 +36,7 @@ def __init__( is_use_oldest: bool = True, ): self.input_batch_size = input_batch_size - self.prompt_lookup_num_tokens = prompt_lookup_num_tokens + self.max_draft_len = max_draft_len self.max_matching_ngram_size = max_matching_ngram_size self.end_id = end_id self.max_seq_len = max_seq_len @@ -45,7 +45,7 @@ def __init__( self.pool = [{} for _ in range(input_batch_size)] self.start_index = [0 for _ in range(input_batch_size)] - assert self.prompt_lookup_num_tokens > 0, f"prompt_lookup_num_tokens must be greater than 0, but got {self.prompt_lookup_num_tokens}" + assert self.max_draft_len > 0, f"max_draft_len must be greater than 0, but got {self.max_draft_len}" assert self.max_matching_ngram_size > 0, f"max_matching_ngram_size must be greater than 0, but got {self.max_matching_ngram_size}" def print_pool(self): @@ -82,16 +82,15 @@ def get_draft_tokens(self, prefix: list[torch.Tensor], -1): # Find each possible key-value combination, and use tuple for hash for l in range(len(sequence) - size): - r = min(l + size + self.prompt_lookup_num_tokens, - len(sequence)) + r = min(l + size + self.max_draft_len, len(sequence)) key = tuple(sequence[l:l + size]) value = tuple(sequence[l + size:r]) if key not in self.pool[gbi] or not self.is_keep_all or \ - len(self.pool[gbi][key][0]) < self.prompt_lookup_num_tokens: + len(self.pool[gbi][key][0]) < self.max_draft_len: # Update the value if # 1. the key does not exist # 2. we only keep the newest one value for each key (MRU) - # 3. the length of the value saved before is less than `prompt_lookup_num_tokens` + # 3. the length of the value saved before is less than `max_draft_len` self.pool[gbi][key] = OrderedSet((value, )) elif value not in self.pool[gbi][key]: # Extend the value if the key is already existed but count of values is not enough @@ -113,26 +112,26 @@ def get_draft_tokens(self, prefix: list[torch.Tensor], break draft_tokens.append(chosen_ids) self.start_index[gbi] = max( - 0, prefix_len[bi] - (self.prompt_lookup_num_tokens + - self.max_matching_ngram_size - 1)) + 0, prefix_len[bi] - + (self.max_draft_len + self.max_matching_ngram_size - 1)) return draft_tokens, None -def run_dtm_pld(batch_input_ids, - args, - runtime_rank, - end_id, - pad_id, - stop_words_list, - bad_words_list, - vocab_size, - *, - target_runner=None): - # `dtm` for Draft-Target-Model, `pld` for Prompt-Lookup-Decoding +def run_dtm_ngram(batch_input_ids, + args, + runtime_rank, + end_id, + pad_id, + stop_words_list, + bad_words_list, + vocab_size, + *, + target_runner=None): + # `dtm` for Draft-Target-Model, `ngram` for NGram is_dtm = (args.draft_target_model_config is not None) - is_pld = (args.prompt_lookup_config is not None) - assert is_dtm ^ is_pld, "`--draft_target_model_config` and `--prompt_lookup_config` can not be specified at the same time." + is_ngram = (args.ngram_config is not None) + assert is_dtm ^ is_ngram, "`--draft_target_model_config` and `--ngram_config` can not be specified at the same time." if is_dtm: assert args.draft_engine_dir is not None, "`--draft_engine_dir` must be specified in Draft-Target-Model." draft_len, draft_device_list, target_device_list, use_logits = ast.literal_eval( @@ -142,12 +141,11 @@ def run_dtm_pld(batch_input_ids, logger.info(f"Device(s) for draft model: {draft_device_list}") logger.info(f"Device(s) for target model: {target_device_list}") logger.info(f"Use logits to accept tokens: {use_logits}") - if is_pld: - logger.info( - f"Using Prompt-Lookup-Decoding speculative decoding V1 workflow") - prompt_lookup_num_tokens, max_matching_ngram_size, target_device_list = ast.literal_eval( - args.prompt_lookup_config) - logger.info(f"prompt_lookup_num_tokens: {prompt_lookup_num_tokens}") + if is_ngram: + logger.info(f"Using NGram speculative decoding V1 workflow") + max_draft_len, max_matching_ngram_size, target_device_list = ast.literal_eval( + args.ngram_config) + logger.info(f"max_draft_len: {max_draft_len}") logger.info(f"max_matching_ngram_size: {max_matching_ngram_size}") logger.info(f"Device(s) for the model: {target_device_list}") use_logits = False # `logits` is useless in this approach yet @@ -166,9 +164,9 @@ def run_dtm_pld(batch_input_ids, n_draft_token = [0 for _ in range(input_batch_size)] n_accept_token = [0 for _ in range(input_batch_size)] - if is_pld: - pld_pool = PLDPool(input_batch_size, prompt_lookup_num_tokens, - max_matching_ngram_size, end_id, max_seq_len) + if is_ngram: + ngram_pool = NgramPool(input_batch_size, max_draft_len, + max_matching_ngram_size, end_id, max_seq_len) # Repack the output like the output of function `generate` outputs = {} @@ -297,8 +295,8 @@ def run_dtm_pld(batch_input_ids, if use_logits: d_logits[bi] = draft["generation_logits"][bi, 0, -d_len[bi]:, :] - if is_pld: - d_ids, d_logits = pld_pool.get_draft_tokens(prefix, batch_slot) + if is_ngram: + d_ids, d_logits = ngram_pool.get_draft_tokens(prefix, batch_slot) d_len = [len(i) for i in d_ids] # Run target model @@ -310,8 +308,8 @@ def run_dtm_pld(batch_input_ids, draft_logits_list=d_logits) if is_dtm: max_new_tokens = draft_len + 1 - if is_pld: - max_new_tokens = prompt_lookup_num_tokens + 1 + if is_ngram: + max_new_tokens = max_draft_len + 1 target_generation_kwargs.update(max_new_tokens=max_new_tokens) target = target_runner.generate(**target_generation_kwargs) torch.cuda.synchronize() diff --git a/examples/run.py b/examples/run.py index fed6c3851d5d..3e46e9d9f6c0 100755 --- a/examples/run.py +++ b/examples/run.py @@ -35,7 +35,7 @@ if PYTHON_BINDINGS: from tensorrt_llm.runtime import ModelRunnerCpp -from prompt_lookup.run_dtm_pld import run_dtm_pld +from ngram.run_dtm_ngram import run_dtm_ngram def parse_arguments(args=None): @@ -430,17 +430,17 @@ def main(args): logger.info(f"Using {'Python' if args.use_py_session else 'C++'} session") - if args.draft_target_model_config is not None or args.prompt_lookup_config is not None: - # Speculative-Decoding of Draft-Target-Model (DTM) and Prompt-Lookup-Decoding (PLD) - # If the parameters of `runner_kwargs` and `runner.generate()` in the "else" branch change, the same change should be done for `examples/prompt_lookup/run_dtm_pld.py` + if args.draft_target_model_config is not None or args.ngram_config is not None: + # Speculative-Decoding of Draft-Target-Model (DTM) and NGram + # If the parameters of `runner_kwargs` and `runner.generate()` in the "else" branch change, the same change should be done for `examples/ngram/run_dtm_ngram.py` assert args.kv_cache_enable_block_reuse, "`--kv_cache_enable_block_reuse` must be specified in speculative decoding." assert not args.use_py_session, "`--use_py_session` is not supported in Speculative decoding." assert not is_enc_dec, "Encoder-Decoder model is not supported in Speculative decoding." assert args.num_beams == 1, "`--num_beams>1` is not supported in Speculative decoding." - outputs = run_dtm_pld(batch_input_ids, args, runtime_rank, end_id, - pad_id, stop_words_list, bad_words_list, - len(tokenizer)) + outputs = run_dtm_ngram(batch_input_ids, args, runtime_rank, end_id, + pad_id, stop_words_list, bad_words_list, + len(tokenizer)) if not args.streaming: # Unpack runner from the return value in No-Streaming mode outputs, runner = list(outputs)[0] diff --git a/examples/scaffolding/run_best_of_n_with_reward.py b/examples/scaffolding/run_best_of_n_with_reward.py index e451cf6b2c03..6ff9ed1228a3 100644 --- a/examples/scaffolding/run_best_of_n_with_reward.py +++ b/examples/scaffolding/run_best_of_n_with_reward.py @@ -60,7 +60,7 @@ def main(): prompts = [query] results = llm.generate(prompts) - print(results[0].output.output_str) + print(results[0].outputs[0].text) llm.shutdown(shutdown_workers=True) print(f'main shut down done') diff --git a/examples/scaffolding/run_majority_vote_aime24.py b/examples/scaffolding/run_majority_vote_aime24.py index 64b4510b19dd..a3587a136639 100644 --- a/examples/scaffolding/run_majority_vote_aime24.py +++ b/examples/scaffolding/run_majority_vote_aime24.py @@ -101,9 +101,8 @@ def main(): result = results[i] test_case = test_dataset[i] ref_answer = int(test_case["answer"]) - result.result() - output = result.output - extracted_answer = extract_answer_from_boxed(output.output_str) + output = result.outputs[0] + extracted_answer = extract_answer_from_boxed(output.text) try: # print(f"[QUESTION]:\n{prompt}\n\n[OUTPUT]\n\n{output.output_str}\n\n") answer = int(extracted_answer) diff --git a/examples/summarize.py b/examples/summarize.py index d984ce65666c..273c1700015b 100644 --- a/examples/summarize.py +++ b/examples/summarize.py @@ -41,7 +41,7 @@ if PYTHON_BINDINGS: from tensorrt_llm.runtime import ModelRunnerCpp -from prompt_lookup.run_dtm_pld import run_dtm_pld +from ngram.run_dtm_ngram import run_dtm_ngram def ensemble_mrope_params(batch_input_ids, max_position_embeddings, @@ -318,17 +318,17 @@ def eval_trt_llm(datapoint, return [], [], [], {} input_lengths = [x.size(0) for x in batch_input_ids] - if args.prompt_lookup_config is not None: - # Speculative decoding of Prompt-Lookup-Decoding (PLD) - outputs = run_dtm_pld(batch_input_ids, - args, - runtime_rank, - end_id, - pad_id, - stop_words_list, - bad_words_list, - tokenizer.vocab_size, - target_runner=runner) + if args.ngram_config is not None: + # Speculative decoding of NGram + outputs = run_dtm_ngram(batch_input_ids, + args, + runtime_rank, + end_id, + pad_id, + stop_words_list, + bad_words_list, + tokenizer.vocab_size, + target_runner=runner) if not args.streaming: # Unpack runner from the return value in No-Streaming mode outputs, runner = list(outputs)[0] else: # Normal run @@ -596,18 +596,17 @@ def eval_hf(datapoint, args.lookahead_config ) == 3, "Lookahead needs [max_window_size, max_ngram_size, max_verification_set_size]" runner_kwargs.update(lookahead_config=args.lookahead_config) - if args.prompt_lookup_config is not None: + if args.ngram_config is not None: assert args.kv_cache_enable_block_reuse, "`--kv_cache_enable_block_reuse` must be specified in speculative decoding." assert not args.use_py_session, "`--use_py_session` is not supported in Speculative decoding." - assert not is_enc_dec, "Encoder-Decoder model is not supported in Speculative decoding." assert args.num_beams == 1, "`--num_beams>1` is not supported in Speculative decoding." - prompt_lookup_num_tokens, _, target_device_list = ast.literal_eval( - args.prompt_lookup_config) - args.max_output_len = output_len # Specialization for PLD + max_draft_len, _, target_device_list = ast.literal_eval( + args.ngram_config) + args.max_output_len = output_len # Specialization for NGram runner_kwargs.update(is_orchestrator_mode=True, device_ids=target_device_list, - max_input_len=test_token_num + - prompt_lookup_num_tokens + output_len) + max_input_len=test_token_num + max_draft_len + + output_len) runner = runner_cls.from_dir(**runner_kwargs) assert not (args.eval_ppl and not runner.gather_context_logits), \ diff --git a/examples/utils.py b/examples/utils.py index c7556298bc24..509b734ebeaa 100644 --- a/examples/utils.py +++ b/examples/utils.py @@ -439,12 +439,12 @@ def add_common_args(parser): " E.g.: [4, [0], [1], False] for [draft_len, draft_model_device_list, target_model_device_list, use_logits]." ) parser.add_argument( - '--prompt_lookup_config', + '--ngram_config', type=str, default=None, help= - "Configuration of Prompt-Lookup decoding, see `examples/prompt_lookup/README.md` for more information." - " E.g.: [10,2,[0]] for [prompt_lookup_num_tokens, max_matching_ngram_size, device_list].", + "Configuration of NGram decoding, see `examples/ngram/README.md` for more information." + " E.g.: [10,2,[0]] for [max_draft_len, max_matching_ngram_size, device_list].", ) parser.add_argument( '--medusa_choices', diff --git a/examples/wide_ep/slurm_scripts/README.md b/examples/wide_ep/slurm_scripts/README.md index 752373bdc6fe..3bd5e926b210 100644 --- a/examples/wide_ep/slurm_scripts/README.md +++ b/examples/wide_ep/slurm_scripts/README.md @@ -17,13 +17,10 @@ Please note that: ### Core Scripts -1. **`submit.sh`** - Main entry point for submitting benchmark jobs -2. **`disaggr_torch.slurm`** - SLURM job script orchestrating the entire benchmark -3. **`gen_yaml.py`** - Generates configuration files for serving setup -4. **`start_server.sh`** - Starts the inference server -5. **`start_worker.sh`** - Starts the worker processes -6. **`run_benchmark.sh`** - Executes the benchmark workload -7. **`process_gen_iterlog.py`** - Processes benchmark results and generates reports +Note that, core implementation of the slurm scripts are included in `examples/disaggregated/slurm`. + +1. `submit.sh` - Main entry point for submitting benchmark jobs +2. `process_gen_iterlog.py` - Processes benchmark results and generates reports ## Usage @@ -35,94 +32,18 @@ Before running the scripts, ensure you have: - Model files accessible on the cluster - Required environment variables set -### Configuration - -Edit the following variables in `submit.sh` and `disaggr_torch.slurm`: +### Running Benchmarks ```bash -# In disaggr_torch.slurm -container_image=${container_image} # Your container image -mount_dir=${mount_dir} # Mount directory path -model_dir=${model_dir} # Model directory path +# Refer to `examples/disaggregated/slurm/` +# Please find the `disaggr_torch.slurm` script in the `examples/disaggregated/slurm/` directory. +# Make sure that SLURM parameters are correctly set in `disaggr_torch.slurm` before executing this script. +./submit.sh ``` -### Running Benchmarks -1. **Submit benchmark jobs**: - ```bash - ./submit.sh - ``` - -2. **Monitor job progress**: - ```bash - squeue -u $USER - ``` - -3. **View results**: - Results are saved in `bm_20250703_deepseek-r1-{isl}-{osl}/` directory - -## Script Details - -### `submit.sh` -Main entry script that submits multiple SLURM jobs with different configurations: -- **DEP8**: 8-way parallelism for decode servers -- **DEP16**: 16-way parallelism with different EPLB slot configurations -- **DEP32**: 32-way parallelism for high-throughput scenarios - -Parameters tested: -- Concurrency levels: 1x, 64x, 1024x multipliers -- EPLB slots: 0, 256, 288 -- Different parallelism sizes - -### `disaggr_torch.slurm` -SLURM job script that: -1. Sets up container environment -2. Generates configuration files -3. Starts server and workers -4. Executes benchmarks -5. Cleans up processes - -**Key parameters**: -- `num_ctx_servers`: Number of context servers -- `ctx_tp_size`: Tensor parallel size for context servers -- `num_gen_servers`: Number of generation servers -- `gen_tp_size`: Tensor parallel size for generation servers -- `concurrency`: Number of concurrent requests - -### `gen_yaml.py` -Generates YAML configuration files with: -- Server topology and resource allocation -- Network configuration (hostnames, ports) -- Memory and batch size settings -- Optimization parameters (CUDA graphs, KV cache) - -**Key features**: -- Automatic node and task allocation -- Support for attention data parallelism -- MoE load balancing configuration -- Speculative decoding (MTP) support - -### `start_server.sh` & `start_worker.sh` -- **Server**: Starts the main inference server with API endpoint -- **Workers**: Starts MPI workers for distributed processing -- Support for profiling with NSight Systems -- Environment variable configuration for optimizations - -### `run_benchmark.sh` -Executes benchmarking using TensorRT-LLM's benchmark_serving tool: -- Downloads ShareGPT dataset for realistic workloads -- Waits for server health checks -- Runs load testing with specified concurrency -- Collects performance metrics -- Gracefully shuts down services - -**Metrics collected**: -- Throughput (tokens/second) -- Latency (request completion time) -- Context vs generation only statistics - -### `process_gen_iterlog.py` -Post-processes benchmark results: +### Post-processes benchmark results using `process_gen_iterlog.py` + - Parses iteration logs from workers - Calculates throughput metrics - Generates CSV reports diff --git a/examples/wide_ep/slurm_scripts/submit.sh b/examples/wide_ep/slurm_scripts/submit.sh index 47ca87fd1cbe..1ede3ee3d29e 100644 --- a/examples/wide_ep/slurm_scripts/submit.sh +++ b/examples/wide_ep/slurm_scripts/submit.sh @@ -1,31 +1,38 @@ #!/bin/bash + +# !!! +# Please find the `disaggr_torch.slurm` script in the `examples/disaggregated/slurm/` directory. +# Make sure that SLURM parameters are correctly set in `disaggr_torch.slurm` before executing this script. +# !!! + mtp_size=0 +ntasks_per_node=4 # 4 GPUs per GB200 node # dep8 for b in 1 64 1024; do concurrency=$((b * 8)) ctx_num=$(((concurrency + 5499)/5500)) - total_gpu_num=$((ctx_num + 2)) - total_tasks=$((total_gpu_num * 4)) - sbatch --nodes=${total_gpu_num} --ntasks=${total_tasks} --ntasks-per-node=4 --segment=${total_gpu_num} disaggr_torch.slurm ${ctx_num} 4 4 4480 true 1 8 1024 1024 true "0.8" 0 "$mtp_size" "$concurrency" + total_node_num=$((ctx_num + 2)) + ntasks=$((total_node_num * ntasks_per_node)) + sbatch --nodes=${total_node_num} --ntasks=${ntasks} --ntasks-per-node=${ntasks_per_node} --segment=${total_node_num} disaggr_torch.slurm ${ctx_num} 4 4 4480 true 1 8 1024 1024 true "0.8" 0 "$mtp_size" "$concurrency" done # dep16 eplb0, 256, 288 for b in 1 64 1024; do concurrency=$((b * 16)) ctx_num=$(((concurrency + 5499)/5500)) - total_gpu_num=$((ctx_num + 4)) - total_tasks=$((total_gpu_num * 4)) - sbatch --nodes=${total_gpu_num} --ntasks=${total_tasks} --ntasks-per-node=4 --segment=${total_gpu_num} disaggr_torch.slurm ${ctx_num} 4 4 4480 true 1 16 1024 1024 true "0.7" 0 "$mtp_size" "$concurrency" - sbatch --nodes=${total_gpu_num} --ntasks=${total_tasks} --ntasks-per-node=4 --segment=${total_gpu_num} disaggr_torch.slurm ${ctx_num} 4 4 4480 true 1 16 1024 1024 true "0.7" 256 "$mtp_size" "$concurrency" - sbatch --nodes=${total_gpu_num} --ntasks=${total_tasks} --ntasks-per-node=4 --segment=${total_gpu_num} disaggr_torch.slurm ${ctx_num} 4 4 4480 true 1 16 1024 1024 true "0.7" 288 "$mtp_size" "$concurrency" + total_node_num=$((ctx_num + 4)) + ntasks=$((total_node_num * ntasks_per_node)) + sbatch --nodes=${total_node_num} --ntasks=${ntasks} --ntasks-per-node=${ntasks_per_node} --segment=${total_node_num} disaggr_torch.slurm ${ctx_num} 4 4 4480 true 1 16 1024 1024 true "0.7" 0 "$mtp_size" "$concurrency" + sbatch --nodes=${total_node_num} --ntasks=${ntasks} --ntasks-per-node=${ntasks_per_node} --segment=${total_node_num} disaggr_torch.slurm ${ctx_num} 4 4 4480 true 1 16 1024 1024 true "0.7" 256 "$mtp_size" "$concurrency" + sbatch --nodes=${total_node_num} --ntasks=${ntasks} --ntasks-per-node=${ntasks_per_node} --segment=${total_node_num} disaggr_torch.slurm ${ctx_num} 4 4 4480 true 1 16 1024 1024 true "0.7" 288 "$mtp_size" "$concurrency" done # dep32 eplb288 for b in 512; do concurrency=$((b * 32)) ctx_num=$(((concurrency + 5499)/5500)) - total_gpu_num=$((ctx_num + 8)) - total_tasks=$((total_gpu_num * 4)) - sbatch --nodes=${total_gpu_num} --ntasks=${total_tasks} --ntasks-per-node=4 --segment=${total_gpu_num} disaggr_torch.slurm ${ctx_num} 4 4 4480 true 1 32 1024 1024 true "0.7" 288 "$mtp_size" "$concurrency" + total_node_num=$((ctx_num + 8)) + ntasks=$((total_node_num * ntasks_per_node)) + sbatch --nodes=${total_node_num} --ntasks=${ntasks} --ntasks-per-node=${ntasks_per_node} --segment=${total_node_num} disaggr_torch.slurm ${ctx_num} 4 4 4480 true 1 32 1024 1024 true "0.7" 288 "$mtp_size" "$concurrency" done diff --git a/jenkins/Build.groovy b/jenkins/Build.groovy index 51cd760425a6..77e12ee51003 100644 --- a/jenkins/Build.groovy +++ b/jenkins/Build.groovy @@ -47,6 +47,12 @@ CONFIG_LINUX_AARCH64 = "linux_aarch64" @Field def CONFIG_LINUX_AARCH64_LLVM = "linux_aarch64_LLVM" +@Field +def CONFIG_LINUX_X86_64_NANOBIND = "linux_x86_64_Nanobind" + +@Field +def CONFIG_LINUX_AARCH64_NANOBIND = "linux_aarch64_Nanobind" + @Field def BUILD_CONFIGS = [ // Vanilla TARNAME is used for packaging in runLLMPackage @@ -56,6 +62,11 @@ def BUILD_CONFIGS = [ (TARNAME) : "TensorRT-LLM.tar.gz", (WHEEL_ARCHS): "80-real;86-real;89-real;90-real;100-real;120-real", ], + (CONFIG_LINUX_X86_64_NANOBIND) : [ + (WHEEL_EXTRA_ARGS) : "--binding_type nanobind --extra-cmake-vars ENABLE_MULTI_DEVICE=1 --extra-cmake-vars WARNING_IS_ERROR=ON --extra-cmake-vars NIXL_ROOT=/opt/nvidia/nvda_nixl --micro_benchmarks", + (TARNAME) : "nanobind-TensorRT-LLM.tar.gz", + (WHEEL_ARCHS): "80-real;86-real;89-real;90-real;100-real;120-real", + ], (CONFIG_LINUX_X86_64_SINGLE_DEVICE) : [ (WHEEL_EXTRA_ARGS) : "--extra-cmake-vars ENABLE_MULTI_DEVICE=0 --extra-cmake-vars WARNING_IS_ERROR=ON --extra-cmake-vars ENABLE_UCX=0 --micro_benchmarks", (TARNAME) : "single-device-TensorRT-LLM.tar.gz", @@ -71,6 +82,11 @@ def BUILD_CONFIGS = [ (TARNAME) : "TensorRT-LLM-GH200.tar.gz", (WHEEL_ARCHS): "90-real;100-real;120-real", ], + (CONFIG_LINUX_AARCH64_NANOBIND): [ + (WHEEL_EXTRA_ARGS) : "--binding_type nanobind --extra-cmake-vars WARNING_IS_ERROR=ON", + (TARNAME) : "nanobind-TensorRT-LLM-GH200.tar.gz", + (WHEEL_ARCHS): "90-real;100-real;120-real", + ], (CONFIG_LINUX_AARCH64_LLVM) : [ (WHEEL_EXTRA_ARGS) : "--extra-cmake-vars WARNING_IS_ERROR=ON -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_CUDA_HOST_COMPILER=clang -DCMAKE_LINKER_TYPE=LLD", (TARNAME) : "llvm-TensorRT-LLM-GH200.tar.gz", @@ -306,18 +322,19 @@ def uploadArtifacts(artifacts, prefix = UPLOAD_PATH, retryTimes = 2, serverId = for (it in artifacts) { def uploadpath = it.key def filepath = it.value - echo "uploading ${filepath} as ${uploadpath}" - trtllm_utils.llmRetry(retryTimes, "uploadArtifacts", { - rtUpload ( - serverId: serverId, - spec: """{ + def spec = """{ "files": [ { "pattern": "${filepath}", "target": "${prefix}/${uploadpath}" } ] - }""", + }""" + echo "Uploading ${filepath} as ${uploadpath}. Spec: ${spec}" + trtllm_utils.llmRetry(retryTimes, "uploadArtifacts", { + rtUpload ( + serverId: serverId, + spec: spec, ) }) } @@ -428,7 +445,8 @@ def runLLMBuild(pipeline, buildFlags, tarName, is_linux_x86_64) // Build tritonserver artifacts def llmPath = sh (script: "realpath ${LLM_ROOT}",returnStdout: true).trim() - sh "cd ${LLM_ROOT}/triton_backend/inflight_batcher_llm && mkdir build && cd build && cmake .. -DTRTLLM_DIR=${llmPath} -DUSE_CXX11_ABI=ON && make -j${BUILD_JOBS} install" + // TODO: Remove after the cmake version is upgraded to 3.31.8 + sh "cd ${LLM_ROOT}/triton_backend/inflight_batcher_llm && mkdir build && cd build && cmake .. -DTRTLLM_DIR=${llmPath} -DTRITON_COMMON_REPO_TAG=r25.05 -DTRITON_CORE_REPO_TAG=r25.05 -DTRITON_THIRD_PARTY_REPO_TAG=r25.05 -DTRITON_BACKEND_REPO_TAG=r25.05 -DUSE_CXX11_ABI=ON && make -j${BUILD_JOBS} install" // Step 3: packaging wheels into tarfile sh "cp ${LLM_ROOT}/build/tensorrt_llm-*.whl TensorRT-LLM/" @@ -521,6 +539,8 @@ def launchStages(pipeline, cpu_arch, enableFailFast, globalVars) pipeline, cpu_arch == AARCH64_TRIPLE ? CONFIG_LINUX_AARCH64 : CONFIG_LINUX_X86_64_VANILLA), "Build TRT-LLM LLVM": [LLM_DOCKER_IMAGE] + prepareLLMBuild( pipeline, cpu_arch == AARCH64_TRIPLE ? CONFIG_LINUX_AARCH64_LLVM : CONFIG_LINUX_X86_64_LLVM), + "Build TRT-LLM Nanobind": [LLM_DOCKER_IMAGE] + prepareLLMBuild( + pipeline, cpu_arch == AARCH64_TRIPLE ? CONFIG_LINUX_AARCH64_NANOBIND : CONFIG_LINUX_X86_64_NANOBIND), ] if (cpu_arch == X86_64_TRIPLE) { diff --git a/jenkins/BuildDockerImage.groovy b/jenkins/BuildDockerImage.groovy index d283f2d5846d..88ab2650374a 100644 --- a/jenkins/BuildDockerImage.groovy +++ b/jenkins/BuildDockerImage.groovy @@ -12,6 +12,7 @@ withCredentials([string(credentialsId: 'default-llm-repo', variable: 'DEFAULT_LL LLM_REPO = env.gitlabSourceRepoHttpUrl ? env.gitlabSourceRepoHttpUrl : "${DEFAULT_LLM_REPO}" } +ARTIFACT_PATH = env.artifactPath ? env.artifactPath : "sw-tensorrt-generic/llm-artifacts/${JOB_NAME}/${BUILD_NUMBER}" UPLOAD_PATH = env.uploadPath ? env.uploadPath : "sw-tensorrt-generic/llm-artifacts/${JOB_NAME}/${BUILD_NUMBER}" LLM_ROOT = "llm" @@ -25,6 +26,8 @@ LLM_SHORT_COMMIT = env.gitlabCommit ? env.gitlabCommit.substring(0, 7) : "undefi LLM_DEFAULT_TAG = env.defaultTag ?: "${LLM_SHORT_COMMIT}-${LLM_BRANCH_TAG}-${BUILD_NUMBER}" +RUN_SANITY_CHECK = params.runSanityCheck ?: false + BUILD_JOBS = "32" BUILD_JOBS_RELEASE_X86_64 = "32" BUILD_JOBS_RELEASE_SBSA = "32" @@ -37,10 +40,13 @@ def GITHUB_PR_API_URL = "github_pr_api_url" def CACHED_CHANGED_FILE_LIST = "cached_changed_file_list" @Field def ACTION_INFO = "action_info" +@Field +def IMAGE_KEY_TO_TAG = "image_key_to_tag" def globalVars = [ (GITHUB_PR_API_URL): null, (CACHED_CHANGED_FILE_LIST): null, (ACTION_INFO): null, + (IMAGE_KEY_TO_TAG): [:], ] @Field @@ -203,15 +209,11 @@ def buildImage(config, imageKeyToTag) def dependentImageWithTag = "${IMAGE_NAME}/${dependent.dockerfileStage}:${dependentTag}" def customImageWithTag = "${IMAGE_NAME}/${dockerfileStage}:${customTag}" - if (target == "ngc-release") { - if (params.triggerType == "post-merge") { - echo "Use NGC artifacts for post merge build" - dependentImageWithTag = "${NGC_IMAGE_NAME}:${dependentTag}" - imageWithTag = "${NGC_IMAGE_NAME}:${tag}" - customImageWithTag = "${NGC_IMAGE_NAME}:${customTag}" - } - imageKeyToTag["NGC Devel Image ${config.arch}"] = dependentImageWithTag - imageKeyToTag["NGC Release Image ${config.arch}"] = imageWithTag + if (target == "ngc-release" && params.triggerType == "post-merge") { + echo "Use NGC artifacts for post merge build" + dependentImageWithTag = "${NGC_IMAGE_NAME}:${dependentTag}" + imageWithTag = "${NGC_IMAGE_NAME}:${tag}" + customImageWithTag = "${NGC_IMAGE_NAME}:${customTag}" } args += " GITHUB_MIRROR=https://urm.nvidia.com/artifactory/github-go-remote" @@ -266,6 +268,9 @@ def buildImage(config, imageKeyToTag) """ } args += " DEVEL_IMAGE=${dependentImageWithTag}" + if (target == "ngc-release") { + imageKeyToTag["NGC Devel Image ${config.arch}"] = dependentImageWithTag + } } } @@ -290,6 +295,9 @@ def buildImage(config, imageKeyToTag) BUILD_WHEEL_OPTS='-j ${build_jobs}' ${args} """ } + if (target == "ngc-release") { + imageKeyToTag["NGC Release Image ${config.arch}"] = imageWithTag + } } if (customTag) { @@ -429,6 +437,17 @@ def launchBuildJobs(pipeline, globalVars, imageKeyToTag) { } +def getCommonParameters() +{ + return [ + 'gitlabSourceRepoHttpUrl': LLM_REPO, + 'gitlabCommit': env.gitlabCommit, + 'artifactPath': ARTIFACT_PATH, + 'uploadPath': UPLOAD_PATH, + ] +} + + pipeline { agent { kubernetes createKubernetesPodConfig("agent") @@ -494,7 +513,100 @@ pipeline { } } } - stage("Register Images for Security Checks") { + stage("Wait for Build Jobs Complete") { + when { + expression { + RUN_SANITY_CHECK + } + } + steps { + script { + container("python3") { + // Install wget + trtllm_utils.llmExecStepWithRetry(this, script: "apt-get update && apt-get -y install wget") + + // Poll for build artifacts + def artifactBaseUrl = "https://urm.nvidia.com/artifactory/${UPLOAD_PATH}/" + def requiredFiles = [ + "TensorRT-LLM-GH200.tar.gz", + "TensorRT-LLM.tar.gz" + ] + def maxWaitMinutes = 60 + def pollIntervalSeconds = 60 + + echo "Waiting for build artifacts..." + echo "Required files: ${requiredFiles}" + + def startTime = System.currentTimeMillis() + def maxWaitMs = maxWaitMinutes * 60 * 1000 + + while ((System.currentTimeMillis() - startTime) < maxWaitMs) { + def missingFiles = [] + + for (file in requiredFiles) { + def fileUrl = "${artifactBaseUrl}${file}" + def exitCode = sh( + script: "wget --spider --quiet --timeout=30 --tries=1 '${fileUrl}'", + returnStatus: true + ) + + if (exitCode != 0) { + missingFiles.add(file) + } + } + + if (missingFiles.isEmpty()) { + echo "All build artifacts are ready!" + return + } + + def elapsedMinutes = (System.currentTimeMillis() - startTime) / (60 * 1000) + echo "Waiting... (${elapsedMinutes.intValue()} minutes elapsed)" + echo "Missing files: ${missingFiles}" + sleep(pollIntervalSeconds) + } + + def elapsedMinutes = (System.currentTimeMillis() - startTime) / (60 * 1000) + error "Timeout waiting for build artifacts (${elapsedMinutes.intValue()} minutes)" + } + } + } + } + stage("Sanity Check for NGC Images") { + when { + expression { + RUN_SANITY_CHECK + } + } + steps { + script { + globalVars[IMAGE_KEY_TO_TAG] = imageKeyToTag + String globalVarsJson = writeJSON returnText: true, json: globalVars + def parameters = getCommonParameters() + parameters += [ + 'enableFailFast': false, + 'globalVars': globalVarsJson, + ] + + echo "Trigger BuildDockerImageSanityTest job, params: ${parameters}" + + def status = "" + def jobName = "/LLM/helpers/BuildDockerImageSanityTest" + def handle = build( + job: jobName, + parameters: trtllm_utils.toBuildParameters(parameters), + propagate: false, + ) + echo "Triggered job: ${handle.absoluteUrl}" + status = handle.result + + if (status != "SUCCESS") { + error "Downstream job did not succeed" + } + } + } + } + stage("Register NGC Images for Security Checks") { when { expression { return params.nspect_id && params.action == "push" diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index ecfdac3a8dc0..3c59e0d77edd 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -105,15 +105,13 @@ def EXTRA_STAGE_LIST = "extra_stage" @Field def MULTI_GPU_FILE_CHANGED = "multi_gpu_file_changed" @Field -def ONLY_PYTORCH_FILE_CHANGED = "only_pytorch_file_changed" +def ONLY_ONE_GROUP_CHANGED = "only_one_group_changed" @Field def AUTO_TRIGGER_TAG_LIST = "auto_trigger_tag_list" @Field def DEBUG_MODE = "debug" @Field def DETAILED_LOG = "detailed_log" -@Field -def ONLY_DOCS_FILE_CHANGED = "only_docs_file_changed" def testFilter = [ (REUSE_STAGE_LIST): trimForStageList(gitlabParamsFromBot.get(REUSE_STAGE_LIST, null)?.tokenize(',')), @@ -127,11 +125,10 @@ def testFilter = [ (DISABLE_MULTI_GPU_TEST): gitlabParamsFromBot.get((DISABLE_MULTI_GPU_TEST), false), (EXTRA_STAGE_LIST): trimForStageList(gitlabParamsFromBot.get((EXTRA_STAGE_LIST), null)?.tokenize(',')), (MULTI_GPU_FILE_CHANGED): false, - (ONLY_PYTORCH_FILE_CHANGED): false, + (ONLY_ONE_GROUP_CHANGED): "", (DEBUG_MODE): gitlabParamsFromBot.get(DEBUG_MODE, false), (AUTO_TRIGGER_TAG_LIST): [], (DETAILED_LOG): gitlabParamsFromBot.get(DETAILED_LOG, false), - (ONLY_DOCS_FILE_CHANGED): false, ] String reuseBuild = gitlabParamsFromBot.get('reuse_build', null) @@ -142,10 +139,16 @@ def GITHUB_PR_API_URL = "github_pr_api_url" def CACHED_CHANGED_FILE_LIST = "cached_changed_file_list" @Field def ACTION_INFO = "action_info" +@Field +def IMAGE_KEY_TO_TAG = "image_key_to_tag" +@Field +def TARGET_BRANCH = "target_branch" def globalVars = [ (GITHUB_PR_API_URL): gitlabParamsFromBot.get('github_pr_api_url', null), (CACHED_CHANGED_FILE_LIST): null, (ACTION_INFO): gitlabParamsFromBot.get('action_info', null), + (IMAGE_KEY_TO_TAG): [:], + (TARGET_BRANCH): gitlabParamsFromBot.get('target_branch', null), ] // If not running all test stages in the L0 pre-merge, we will not update the GitLab status at the end. @@ -300,32 +303,72 @@ def echoNodeAndGpuInfo(pipeline, stageName) } def setupPipelineEnvironment(pipeline, testFilter, globalVars) +{ + sh "env | sort" + updateGitlabCommitStatus name: "${BUILD_STATUS_NAME}", state: 'running' + echo "Using GitLab repo: ${LLM_REPO}." + sh "git config --global --add safe.directory \"*\"" + // NB: getContainerURIs reads files in ${LLM_ROOT}/jenkins/ + if (env.gitlabMergeRequestLastCommit) { + env.gitlabCommit = env.gitlabMergeRequestLastCommit + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) + } else { + branch = env.gitlabBranch ? env.gitlabBranch : "main" + trtllm_utils.checkoutSource(LLM_REPO, branch, LLM_ROOT, true, true) + checkoutCommit = sh (script: "cd ${LLM_ROOT} && git rev-parse HEAD",returnStdout: true).trim() + env.gitlabCommit = checkoutCommit + } + echo "Env.gitlabMergeRequestLastCommit: ${env.gitlabMergeRequestLastCommit}." + echo "Freeze GitLab commit. Branch: ${env.gitlabBranch}. Commit: ${env.gitlabCommit}." + testFilter[(MULTI_GPU_FILE_CHANGED)] = getMultiGpuFileChanged(pipeline, testFilter, globalVars) + testFilter[(ONLY_ONE_GROUP_CHANGED)] = getOnlyOneGroupChanged(pipeline, testFilter, globalVars) + testFilter[(AUTO_TRIGGER_TAG_LIST)] = getAutoTriggerTagList(pipeline, testFilter, globalVars) + getContainerURIs().each { k, v -> + globalVars[k] = v + } +} + +def mergeWaiveList(pipeline, globalVars) +{ + // Get current waive list + sh "git config --global --add safe.directory \"*\"" + sh "cp ${LLM_ROOT}/tests/integration/test_lists/waives.txt ./waives_CUR_${env.gitlabCommit}.txt" + sh "cp ${LLM_ROOT}/jenkins/scripts/mergeWaiveList.py ./" + + // Get TOT waive list + LLM_TOT_ROOT = "llm-tot" + targetBranch = env.gitlabTargetBranch ? env.gitlabTargetBranch : globalVars[TARGET_BRANCH] + echo "Target branch: ${targetBranch}" + trtllm_utils.checkoutSource(LLM_REPO, targetBranch, LLM_TOT_ROOT, true, true) + targetBranchTOTCommit = sh (script: "cd ${LLM_TOT_ROOT} && git rev-parse HEAD", returnStdout: true).trim() + echo "Target branch TOT commit: ${targetBranchTOTCommit}" + sh "cp ${LLM_TOT_ROOT}/tests/integration/test_lists/waives.txt ./waives_TOT_${targetBranchTOTCommit}.txt" + + // Get waive list diff in current MR + def diff = getMergeRequestOneFileChanges(pipeline, globalVars, "tests/integration/test_lists/waives.txt") + + // Merge waive lists + sh """ + python3 mergeWaiveList.py \ + --cur-waive-list=waives_CUR_${env.gitlabCommit}.txt \ + --latest-waive-list=waives_TOT_${targetBranchTOTCommit}.txt \ + --diff='${diff}' \ + --output-file=waives.txt + """ + trtllm_utils.uploadArtifacts("waives*.txt", "${UPLOAD_PATH}/waive_list/") + echo "New merged test waive list: https://urm.nvidia.com/artifactory/${UPLOAD_PATH}/waive_list/waives.txt" +} + +def preparation(pipeline, testFilter, globalVars) { image = "urm.nvidia.com/docker/golang:1.22" setupPipelineSpec = createKubernetesPodConfig(image, "package") trtllm_utils.launchKubernetesPod(pipeline, setupPipelineSpec, "trt-llm", { - sh "env | sort" - updateGitlabCommitStatus name: "${BUILD_STATUS_NAME}", state: 'running' - echo "Using GitLab repo: ${LLM_REPO}." - sh "git config --global --add safe.directory \"*\"" - // NB: getContainerURIs reads files in ${LLM_ROOT}/jenkins/ - if (env.gitlabMergeRequestLastCommit) { - env.gitlabCommit = env.gitlabMergeRequestLastCommit - trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, true, true) - } else { - branch = env.gitlabBranch ? env.gitlabBranch : "main" - trtllm_utils.checkoutSource(LLM_REPO, branch, LLM_ROOT, true, true) - checkoutCommit = sh (script: "cd ${LLM_ROOT} && git rev-parse HEAD",returnStdout: true).trim() - env.gitlabCommit = checkoutCommit + stage("Setup Environment") { + setupPipelineEnvironment(pipeline, testFilter, globalVars) } - echo "Env.gitlabMergeRequestLastCommit: ${env.gitlabMergeRequestLastCommit}." - echo "Freeze GitLab commit. Branch: ${env.gitlabBranch}. Commit: ${env.gitlabCommit}." - testFilter[(MULTI_GPU_FILE_CHANGED)] = getMultiGpuFileChanged(pipeline, testFilter, globalVars) - testFilter[(ONLY_PYTORCH_FILE_CHANGED)] = getOnlyPytorchFileChanged(pipeline, testFilter, globalVars) - testFilter[(AUTO_TRIGGER_TAG_LIST)] = getAutoTriggerTagList(pipeline, testFilter, globalVars) - testFilter[(ONLY_DOCS_FILE_CHANGED)] = getOnlyDocsFileChanged(pipeline, testFilter, globalVars) - getContainerURIs().each { k, v -> - globalVars[k] = v + stage("Merge Test Waive List") { + mergeWaiveList(pipeline, globalVars) } }) } @@ -416,8 +459,8 @@ def launchReleaseCheck(pipeline) }) } -def getMergeRequestChangedFileListGitlab(pipeline) { - def changedFileList = [] +def getGitlabMRChangedFile(pipeline, function, filePath="") { + def result = null def pageId = 0 withCredentials([ usernamePassword( @@ -437,19 +480,34 @@ def getMergeRequestChangedFileListGitlab(pipeline) { returnStdout: true ) def rawDataList = readJSON text: rawDataJson, returnPojo: true - rawDataList.each { rawData -> - changedFileList += [rawData.get("old_path"), rawData.get("new_path")] + if (function == "getOneFileChanges") { + if (result == null) { + result = "" + } + rawDataList.find { rawData -> + if (rawData.get("new_path") == filePath || rawData.get("old_path") == filePath) { + result = rawData.get("diff") + return true + } + return false + } + if (result != "") { break } + } else if (function == "getChangedFileList") { + if (result == null) { + result = [] + } + rawDataList.each { rawData -> + result += [rawData.get("old_path"), rawData.get("new_path")] + } } if (!rawDataList) { break } } } - def changedFileListStr = changedFileList.join(",\n") - pipeline.echo("The changeset of this MR is: ${changedFileListStr}.") - return changedFileList + return result } -def getMergeRequestChangedFileListGithub(pipeline, githubPrApiUrl) { - def changedFileList = [] +def getGithubMRChangedFile(pipeline, githubPrApiUrl, function, filePath="") { + def result = null def pageId = 0 withCredentials([ string( @@ -468,15 +526,30 @@ def getMergeRequestChangedFileListGithub(pipeline, githubPrApiUrl) { ) echo "rawDataJson: ${rawDataJson}" def rawDataList = readJSON text: rawDataJson, returnPojo: true - rawDataList.each { rawData -> - changedFileList += [rawData.get("filename"), rawData.get("previous_filename")].findAll { it } + if (function == "getOneFileChanges") { + if (result == null) { + result = "" + } + rawDataList.find { rawData -> + if (rawData.get("filename") == filePath || rawData.get("previous_filename") == filePath) { + result = rawData.get("patch") + return true + } + return false + } + if (result != "") { break } + } else if (function == "getChangedFileList") { + if (result == null) { + result = [] + } + rawDataList.each { rawData -> + result += [rawData.get("filename"), rawData.get("previous_filename")].findAll { it } + } } if (!rawDataList) { break } } } - def changedFileListStr = changedFileList.join(",\n") - pipeline.echo("The changeset of this PR is: ${changedFileListStr}.") - return changedFileList + return result } def getMergeRequestChangedFileList(pipeline, globalVars) { @@ -492,11 +565,15 @@ def getMergeRequestChangedFileList(pipeline, globalVars) { return globalVars[CACHED_CHANGED_FILE_LIST] } try { + def changedFileList = [] if (githubPrApiUrl != null) { - globalVars[CACHED_CHANGED_FILE_LIST] = getMergeRequestChangedFileListGithub(pipeline, githubPrApiUrl) + changedFileList = getGithubMRChangedFile(pipeline, githubPrApiUrl, "getChangedFileList") } else { - globalVars[CACHED_CHANGED_FILE_LIST] = getMergeRequestChangedFileListGitlab(pipeline) + changedFileList = getGitlabMRChangedFile(pipeline, "getChangedFileList") } + def changedFileListStr = changedFileList.join(",\n") + pipeline.echo("The changeset of this MR is: ${changedFileListStr}.") + globalVars[CACHED_CHANGED_FILE_LIST] = changedFileList return globalVars[CACHED_CHANGED_FILE_LIST] } catch (InterruptedException e) { throw e @@ -507,6 +584,26 @@ def getMergeRequestChangedFileList(pipeline, globalVars) { } } +def getMergeRequestOneFileChanges(pipeline, globalVars, filePath) { + def githubPrApiUrl = globalVars[GITHUB_PR_API_URL] + def diff = "" + + try { + if (githubPrApiUrl != null) { + diff = getGithubMRChangedFile(pipeline, githubPrApiUrl, "getOneFileChanges", filePath) + } else { + diff = getGitlabMRChangedFile(pipeline, "getOneFileChanges", filePath) + } + pipeline.echo("The change of ${filePath} is: ${diff}") + return diff + } catch (InterruptedException e) { + throw e + } catch (Exception e) { + pipeline.echo("Get merge request one changed file diff failed. Error: ${e.toString()}") + return "" + } +} + def getAutoTriggerTagList(pipeline, testFilter, globalVars) { def autoTriggerTagList = [] def isOfficialPostMergeJob = (env.JOB_NAME ==~ /.*PostMerge.*/) @@ -547,66 +644,69 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) } def relatedFileList = [ + "cpp/include/tensorrt_llm/batch_manager/", + "cpp/include/tensorrt_llm/executor/", "cpp/include/tensorrt_llm/runtime/gptJsonConfig.h", - "cpp/include/tensorrt_llm/runtime/worldConfig.h", "cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h", "cpp/include/tensorrt_llm/runtime/utils/multiDeviceUtils.h", - "cpp/tensorrt_llm/runtime/utils/mpiUtils.cpp", - "cpp/tests/runtime/mpiUtilsTest.cpp", - "cpp/tensorrt_llm/batch_manager/trtGptModelFactory.h", - "cpp/tensorrt_llm/runtime/worldConfig.cpp", - "cpp/tensorrt_llm/runtime/ncclCommunicator.cpp", - "cpp/tensorrt_llm/runtime/workerPool.h", - "cpp/tensorrt_llm/executor_worker/executorWorker.cpp", - "cpp/tensorrt_llm/runtime/ipcUtils.cpp", - "cpp/tensorrt_llm/executor/executor.cpp", - "cpp/tensorrt_llm/executor/executorImpl.cpp", - "cpp/tensorrt_llm/executor/executorImpl.h", - "cpp/tensorrt_llm/runtime/ncclCommunicator.cpp", + "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/thop/allreduceOp.cpp", - "cpp/tensorrt_llm/kernels/customAllReduceKernels.h", "cpp/tensorrt_llm/kernels/customAllReduceKernels.cu", - "cpp/tensorrt_llm/kernels/gptKernels.h", + "cpp/tensorrt_llm/kernels/customAllReduceKernels.h", "cpp/tensorrt_llm/kernels/gptKernels.cu", - "cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h", + "cpp/tensorrt_llm/kernels/gptKernels.h", + "cpp/tensorrt_llm/kernels/moe", "cpp/tensorrt_llm/kernels/unfusedAttentionKernels.cu", + "cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h", "cpp/tensorrt_llm/kernels/userbuffers/", - "cpp/tensorrt_llm/kernels/moe", - "cpp/tensorrt_llm/pybind/", - "cpp/tests/kernels/allReduce/", - "cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.h", "cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.cpp", - "cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h", + "cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.h", "cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp", - "cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h", + "cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h", "cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp", - "cpp/tests/runtime/mpiUtilsTest.cpp", + "cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h", "cpp/tensorrt_llm/plugins/ncclPlugin/", - "tensorrt_llm/functional.py", - "tensorrt_llm/mapping.py", - "tensorrt_llm/llmapi/", - "tensorrt_llm/executor/", + "cpp/tensorrt_llm/pybind/", + "cpp/tensorrt_llm/runtime/ipcUtils.cpp", + "cpp/tensorrt_llm/runtime/ncclCommunicator.cpp", + "cpp/tensorrt_llm/runtime/utils/mpiUtils.cpp", + "cpp/tensorrt_llm/runtime/workerPool.h", + "cpp/tensorrt_llm/runtime/worldConfig.cpp", + "cpp/tensorrt_llm/thop/allgatherOp.cpp", + "cpp/tensorrt_llm/thop/allreduceOp.cpp", + "cpp/tensorrt_llm/thop/reducescatterOp.cpp", + "cpp/tests/executor/", + "cpp/tests/kernels/allReduce/", + "cpp/tests/runtime/mpiUtilsTest.cpp", + "jenkins/L0_Test.groovy", "tensorrt_llm/_ipc_utils.py", - "tensorrt_llm/parameter.py", - "tensorrt_llm/models/llama/", "tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py", "tensorrt_llm/_torch/compilation/patterns/ub_allreduce.py", "tensorrt_llm/_torch/custom_ops/userbuffers_custom_ops.py", - "tensorrt_llm/_torch/pyexecutor/model_engine.py", - "tensorrt_llm/_torch/pyexecutor/py_executor.py", - "tensorrt_llm/_torch/pyexecutor/_util.py", "tensorrt_llm/_torch/models/modeling_llama.py", "tensorrt_llm/_torch/modules/fused_moe/", + "tensorrt_llm/_torch/pyexecutor/_util.py", + "tensorrt_llm/_torch/pyexecutor/model_engine.py", + "tensorrt_llm/_torch/pyexecutor/py_executor.py", + "tensorrt_llm/executor/", + "tensorrt_llm/functional.py", + "tensorrt_llm/llmapi/", + "tensorrt_llm/mapping.py", + "tensorrt_llm/models/llama/", + "tensorrt_llm/parameter.py", + "tensorrt_llm/serve/", "tests/integration/defs/cpp/test_multi_gpu.py", "tests/integration/test_lists/test-db/l0_dgx_h100.yml", "tests/integration/test_lists/test-db/l0_dgx_h200.yml", + "tests/unittest/_torch/auto_deploy/unit/multigpu", "tests/unittest/_torch/multi_gpu/", "tests/unittest/_torch/multi_gpu_modeling/", - "tests/unittest/_torch/auto_deploy/unit/multigpu", + "tests/unittest/disaggregated/", "tests/unittest/llmapi/test_llm_multi_gpu.py", "tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py", - "jenkins/L0_Test.groovy", ] def changedFileList = getMergeRequestChangedFileList(pipeline, globalVars) @@ -638,86 +738,62 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) return relatedFileChanged } -def getOnlyPytorchFileChanged(pipeline, testFilter, globalVars) { +def getOnlyOneGroupChanged(pipeline, testFilter, globalVars) { def isOfficialPostMergeJob = (env.JOB_NAME ==~ /.*PostMerge.*/) if (env.alternativeTRT || isOfficialPostMergeJob) { - pipeline.echo("Force set ONLY_PYTORCH_FILE_CHANGED false.") - return false + pipeline.echo("Force set ONLY_ONE_GROUP_CHANGED \"\".") + return "" } - def pytorchOnlyList = [ - "tensorrt_llm/_torch/", - "tensorrt_llm/scaffolding/", - "tests/unittest/_torch/", - "tests/unittest/scaffolding/", - "tests/unittest/llmapi/test_llm_pytorch.py", - "tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py", - "tests/integration/defs/accuracy/test_llm_api_pytorch.py", - "tests/integration/defs/disaggregated/", - "examples/auto_deploy", - "examples/disaggregated", - "examples/pytorch/", - "examples/scaffolding/", - "docs/" + def groupFileMap = [ + "Docs": [ // TODO: Add more docs path to the list, e.g. *.md files in other directories + "docs/", + ], + "PyTorch": [ + "tensorrt_llm/_torch/", + "tensorrt_llm/scaffolding/", + "tests/unittest/_torch/", + "tests/unittest/scaffolding/", + "tests/unittest/llmapi/test_llm_pytorch.py", + "tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py", + "tests/integration/defs/accuracy/test_llm_api_pytorch.py", + "tests/integration/defs/disaggregated/", + "examples/auto_deploy", + "examples/disaggregated", + "examples/pytorch/", + "examples/scaffolding/", + "docs/", + ], + "Triton": [ + "tests/integration/defs/triton_server/", + "triton_backend/", + ], ] def changedFileList = getMergeRequestChangedFileList(pipeline, globalVars) - if (!changedFileList || changedFileList.isEmpty()) { - return false + return "" } - def result = true - for (file in changedFileList) { - def isPytorchFile = false - for (prefix in pytorchOnlyList) { - if (file.startsWith(prefix)) { - isPytorchFile = true - break - } - } - if (!isPytorchFile) { - pipeline.echo("Found non-PyTorch file: ${file}") - result = false - break + for (group in groupFileMap.keySet()) { + def groupPrefixes = groupFileMap[group] + def allFilesInGroup = changedFileList.every { file -> + groupPrefixes.any { prefix -> file.startsWith(prefix) } } - } - - pipeline.echo("Only PyTorch files changed: ${result}") - return result -} - -def getOnlyDocsFileChanged(pipeline, testFilter, globalVars) { - def isOfficialPostMergeJob = (env.JOB_NAME ==~ /.*PostMerge.*/) - if (env.alternativeTRT || isOfficialPostMergeJob) { - pipeline.echo("Force set ONLY_DOCS_FILE_CHANGED false.") - return false - } - - // TODO: Add more docs path to the list, e.g. *.md files in other directories - def docsFileList = [ - "docs/", - ] - def changedFileList = getMergeRequestChangedFileList(pipeline, globalVars) - if (!changedFileList || changedFileList.isEmpty()) { - return false - } - - for (file in changedFileList) { - def isDocsFile = false - for (prefix in docsFileList) { - if (file.startsWith(prefix)) { - isDocsFile = true - break + if (allFilesInGroup) { + pipeline.echo("Only ${group} files changed.") + return group + } else { + def nonGroupFile = changedFileList.find { file -> + !groupPrefixes.any { prefix -> file.startsWith(prefix) } + } + if (nonGroupFile != null) { + pipeline.echo("Found non-${group} file: ${nonGroupFile}") } - } - if (!isDocsFile) { - pipeline.echo("Found non-docs file: ${file}") - return false } } - pipeline.echo("Only docs files changed.") - return true + + return "" } def collectTestResults(pipeline, testFilter) @@ -878,6 +954,45 @@ def triggerJob(jobName, parameters, jenkinsUrl = "", credentials = "") return status } +def launchJob(jobName, reuseBuild, enableFailFast, globalVars, platform="x86_64", additionalParameters = [:]) { + def parameters = getCommonParameters() + String globalVarsJson = writeJSON returnText: true, json: globalVars + parameters += [ + 'enableFailFast': enableFailFast, + 'globalVars': globalVarsJson, + ] + additionalParameters + + if (env.alternativeTRT && platform == "x86_64") { + parameters += [ + 'alternativeTRT': env.alternativeTRT, + ] + } + + if (env.alternativeTrtSBSA && platform == "SBSA") { + parameters += [ + 'alternativeTRT': env.alternativeTrtSBSA, + ] + } + + if (env.testPhase2StageName) { + parameters += [ + 'testPhase2StageName': env.testPhase2StageName, + ] + } + + if (reuseBuild) { + parameters['reuseArtifactPath'] = "sw-tensorrt-generic/llm-artifacts/${JOB_NAME}/${reuseBuild}" + } + + echo "Trigger ${jobName} job, params: ${parameters}" + + def status = triggerJob(jobName, parameters) + if (status != "SUCCESS") { + error "Downstream job did not succeed" + } + return status +} + def launchStages(pipeline, reuseBuild, testFilter, enableFailFast, globalVars) { stages = [ @@ -889,78 +1004,87 @@ def launchStages(pipeline, reuseBuild, testFilter, enableFailFast, globalVars) "x86_64-linux": { script { stage("Build") { - def parameters = getCommonParameters() - String globalVarsJson = writeJSON returnText: true, json: globalVars - parameters += [ - 'enableFailFast': enableFailFast, + def additionalParameters = [ 'dockerImage': globalVars["LLM_DOCKER_IMAGE"], 'wheelDockerImagePy310': globalVars["LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE"], 'wheelDockerImagePy312': globalVars["LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE"], - 'globalVars': globalVarsJson, ] - - if (env.alternativeTRT) { - parameters += [ - 'alternativeTRT': env.alternativeTRT, - ] - } - - if (reuseBuild) { - parameters['reuseArtifactPath'] = "sw-tensorrt-generic/llm-artifacts/${JOB_NAME}/${reuseBuild}" - } - - echo "trigger x86_64 build job, params: ${parameters}" - - def status = triggerJob("/LLM/helpers/Build-x86_64", parameters) - if (status != "SUCCESS") { - error "Downstream job did not succeed" - } - - } - def testStageName = "[Test-x86_64] Run" - if (env.localJobCredentials) { - testStageName = "[Test-x86_64] Remote Run" + launchJob("/LLM/helpers/Build-x86_64", reuseBuild, enableFailFast, globalVars, "x86_64", additionalParameters) } + def testStageName = "[Test-x86_64-Single-GPU] ${env.localJobCredentials ? "Remote Run" : "Run"}" + def singleGpuTestFailed = false stage(testStageName) { if (X86_TEST_CHOICE == STAGE_CHOICE_SKIP) { echo "x86_64 test job is skipped due to Jenkins configuration" return } try { - parameters = getCommonParameters() String testFilterJson = writeJSON returnText: true, json: testFilter - String globalVarsJson = writeJSON returnText: true, json: globalVars - parameters += [ - 'enableFailFast': enableFailFast, + def additionalParameters = [ 'testFilter': testFilterJson, 'dockerImage': globalVars["LLM_DOCKER_IMAGE"], 'wheelDockerImagePy310': globalVars["LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE"], 'wheelDockerImagePy312': globalVars["LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE"], - 'globalVars': globalVarsJson, ] - if (env.alternativeTRT) { - parameters += [ - 'alternativeTRT': env.alternativeTRT, - ] + launchJob("L0_Test-x86_64-Single-GPU", false, enableFailFast, globalVars, "x86_64", additionalParameters) + } catch (InterruptedException e) { + throw e + } catch (Exception e) { + if (X86_TEST_CHOICE == STAGE_CHOICE_IGNORE) { + catchError( + buildResult: 'SUCCESS', + stageResult: 'FAILURE') { + error "x86_64 test failed but ignored due to Jenkins configuration" + } + } else { + catchError( + buildResult: 'FAILURE', + stageResult: 'FAILURE') { + error "x86_64 single-GPU test failed" + } + singleGpuTestFailed = true } + } + } + + def requireMultiGpuTesting = currentBuild.description?.contains("Require Multi-GPU Testing") ?: false + echo "requireMultiGpuTesting: ${requireMultiGpuTesting}" + if (!requireMultiGpuTesting) { + if (singleGpuTestFailed) { + error "Single-GPU test failed" + } + return + } - if (env.testPhase2StageName) { - parameters += [ - 'testPhase2StageName': env.testPhase2StageName, - ] + if (singleGpuTestFailed) { + if (env.JOB_NAME ==~ /.*PostMerge.*/) { + echo "In the official post-merge pipeline, single-GPU test failed, whereas multi-GPU test is still kept running." + } else { + stage("[Test-x86_64-Multi-GPU] Blocked") { + error "This pipeline requires running multi-GPU test, but single-GPU test has failed." } + return + } + } - echo "trigger x86_64 test job, params: ${parameters}" + testStageName = "[Test-x86_64-Multi-GPU] ${env.localJobCredentials ? "Remote Run" : "Run"}" + stage(testStageName) { + if (X86_TEST_CHOICE == STAGE_CHOICE_SKIP) { + echo "x86_64 test job is skipped due to Jenkins configuration" + return + } + try { + def testFilterJson = writeJSON returnText: true, json: testFilter + def additionalParameters = [ + 'testFilter': testFilterJson, + 'dockerImage': globalVars["LLM_DOCKER_IMAGE"], + 'wheelDockerImagePy310': globalVars["LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE"], + 'wheelDockerImagePy312': globalVars["LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE"], + ] - def status = triggerJob( - "L0_Test-x86_64", - parameters, - ) + launchJob("L0_Test-x86_64-Multi-GPU", false, enableFailFast, globalVars, "x86_64", additionalParameters) - if (status != "SUCCESS") { - error "Downstream job did not succeed" - } } catch (InterruptedException e) { throw e } catch (Exception e) { @@ -986,43 +1110,16 @@ def launchStages(pipeline, reuseBuild, testFilter, enableFailFast, globalVars) testStageName = "[Test-SBSA] Remote Run" } - if (testFilter[(ONLY_DOCS_FILE_CHANGED)]) { + if (testFilter[(ONLY_ONE_GROUP_CHANGED)] == "Docs") { echo "SBSA build job is skipped due to Jenkins configuration or conditional pipeline run" return } - def stageName = "Build" - stage(stageName) { - def parameters = getCommonParameters() - String globalVarsJson = writeJSON returnText: true, json: globalVars - parameters += [ - 'enableFailFast': enableFailFast, + stage("Build") { + def additionalParameters = [ "dockerImage": globalVars["LLM_SBSA_DOCKER_IMAGE"], - 'globalVars': globalVarsJson, ] - - if (env.alternativeTrtSBSA) { - parameters += [ - "alternativeTRT": env.alternativeTrtSBSA, - ] - } - - if (reuseBuild) { - parameters['reuseArtifactPath'] = "sw-tensorrt-generic/llm-artifacts/${JOB_NAME}/${reuseBuild}" - } - - echo "trigger SBSA build job, params: ${parameters}" - - def status = triggerJob( - "/LLM/helpers/Build-SBSA", - parameters, - jenkinsUrl, - credentials, - ) - - if (status != "SUCCESS") { - error "Downstream job did not succeed" - } + launchJob("/LLM/helpers/Build-SBSA", reuseBuild, enableFailFast, globalVars, "SBSA", additionalParameters) } stage(testStageName) { if (SBSA_TEST_CHOICE == STAGE_CHOICE_SKIP) { @@ -1030,40 +1127,14 @@ def launchStages(pipeline, reuseBuild, testFilter, enableFailFast, globalVars) return } try { - def parameters = getCommonParameters() - String testFilterJson = writeJSON returnText: true, json: testFilter - String globalVarsJson = writeJSON returnText: true, json: globalVars - parameters += [ - 'enableFailFast': enableFailFast, + def testFilterJson = writeJSON returnText: true, json: testFilter + def additionalParameters = [ 'testFilter': testFilterJson, "dockerImage": globalVars["LLM_SBSA_DOCKER_IMAGE"], - 'globalVars': globalVarsJson, ] - if (env.alternativeTrtSBSA) { - parameters += [ - "alternativeTRT": env.alternativeTrtSBSA, - ] - } - - if (env.testPhase2StageName) { - parameters += [ - 'testPhase2StageName': env.testPhase2StageName, - ] - } - - echo "trigger SBSA test job, params: ${parameters}" - - def status = triggerJob( - "L0_Test-SBSA", - parameters, - jenkinsUrl, - credentials, - ) + launchJob("L0_Test-SBSA", false, enableFailFast, globalVars, "SBSA", additionalParameters) - if (status != "SUCCESS") { - error "Downstream job did not succeed" - } } catch (InterruptedException e) { throw e } catch (Exception e) { @@ -1085,31 +1156,24 @@ def launchStages(pipeline, reuseBuild, testFilter, enableFailFast, globalVars) "Build-Docker-Images": { script { stage("[Build-Docker-Images] Remote Run") { - def parameters = getCommonParameters() - String globalVarsJson = writeJSON returnText: true, json: globalVars def branch = env.gitlabBranch ? env.gitlabBranch : "main" if (globalVars[GITHUB_PR_API_URL]) { branch = "github-pr-" + globalVars[GITHUB_PR_API_URL].split('/').last() } - parameters += [ - 'enableFailFast': enableFailFast, + def additionalParameters = [ 'branch': branch, 'action': "push", 'triggerType': env.JOB_NAME ==~ /.*PostMerge.*/ ? "post-merge" : "pre-merge", - 'globalVars': globalVarsJson, + 'runSanityCheck': true, ] - echo "trigger BuildDockerImages job, params: ${parameters}" - - def status = triggerJob("/LLM/helpers/BuildDockerImages", parameters) - if (status != "SUCCESS") { - error "Downstream job did not succeed" - } + launchJob("/LLM/helpers/BuildDockerImages", false, enableFailFast, globalVars, "x86_64", additionalParameters) } } } ] + if (env.JOB_NAME ==~ /.*PostMerge.*/) { stages += dockerBuildJob } @@ -1170,12 +1234,12 @@ pipeline { } } stages { - stage("Setup environment") + stage("Preparation") { steps { script { - setupPipelineEnvironment(this, testFilter, globalVars) + preparation(this, testFilter, globalVars) println globalVars globalVars[ACTION_INFO] = trtllm_utils.setupPipelineDescription(this, globalVars[ACTION_INFO]) echo "enableFailFast is: ${enableFailFast}" diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 7dec81f7fde2..ffa1fa3267c1 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -64,6 +64,9 @@ def LLVM_CONFIG = "LLVM" @Field LINUX_AARCH64_CONFIG = "linux_aarch64" +@Field +def NANOBIND_CONFIG = "Nanobind" + @Field def BUILD_CONFIGS = [ // Vanilla TARNAME is used for packaging in runLLMPackage @@ -71,6 +74,7 @@ def BUILD_CONFIGS = [ (SINGLE_DEVICE_CONFIG) : [(TARNAME) : "single-device-TensorRT-LLM.tar.gz"], (LLVM_CONFIG) : [(TARNAME) : "llvm-TensorRT-LLM.tar.gz"], (LINUX_AARCH64_CONFIG) : [(TARNAME) : "TensorRT-LLM-GH200.tar.gz"], + (NANOBIND_CONFIG) : [(TARNAME) : "nanobind-TensorRT-LLM.tar.gz"], ] // TODO: Move common variables to an unified location @@ -91,6 +95,10 @@ TESTER_MEMORY = "96Gi" CCACHE_DIR="/mnt/sw-tensorrt-pvc/scratch.trt_ccache/llm_ccache" MODEL_CACHE_DIR="/scratch.trt_llm_data/llm-models" +// ENABLE_NGC_DEVEL_IMAGE_TEST is currently disabled in the Jenkins BuildDockerImageSanityTest job config +ENABLE_NGC_DEVEL_IMAGE_TEST = params.enableNgcDevelImageTest ?: false +ENABLE_NGC_RELEASE_IMAGE_TEST = params.enableNgcReleaseImageTest ?: false + def uploadResults(def pipeline, SlurmCluster cluster, String nodeName, String stageName){ withCredentials([usernamePassword(credentialsId: 'svc_tensorrt', usernameVariable: 'USERNAME', passwordVariable: 'PASSWORD')]) { def remote = [ @@ -254,6 +262,10 @@ def runLLMTestlistOnSlurm(pipeline, platform, testList, config=VANILLA_CONFIG, p if (CloudManager.isNodeOnline(nodeName)) { def dockerArgs = "--gpus ${gpuCount} --cap-add=SYS_ADMIN --ipc=host --security-opt seccomp=unconfined -u root:root -v /home/scratch.trt_llm_data:/scratch.trt_llm_data:ro -v /tmp/ccache:${CCACHE_DIR}:rw -v /tmp/pipcache/http-v2:/root/.cache/pip/http-v2:rw --cap-add syslog" + + if (partition.clusterName == "dlcluster") { + dockerArgs += " -e NVIDIA_IMEX_CHANNELS=0" + } slurmRunner = runInDockerOnNodeMultiStage(LLM_DOCKER_IMAGE, nodeName, dockerArgs, false) executeLLMTestOnSlurm(pipeline, platform, testList, config, perfMode, stageName, splitId, splits, skipInstallWheel, cpver, slurmRunner) } else { @@ -309,6 +321,7 @@ def runLLMTestlistOnSlurm_MultiNodes(pipeline, platform, testList, config=VANILL def llmSrcLocal = "${llmPath}/TensorRT-LLM/src" def scriptRunNode = "${jobWorkspace}/slurm_run.sh" def testListPathNode = "${jobWorkspace}/${testList}.txt" + def waivesListPathNode = "${jobWorkspace}/waives.txt" def isAarch64 = config.contains("aarch64") def pytestTestTimeout = "7200" @@ -325,6 +338,10 @@ def runLLMTestlistOnSlurm_MultiNodes(pipeline, platform, testList, config=VANILL Utils.exec(pipeline, script: "chmod +x ${scriptRunLocalPath}", returnStdout: true) Utils.exec(pipeline, script: "sshpass -p '${remote.passwd}' scp -r -p -oStrictHostKeyChecking=no ${scriptRunLocalPath} ${remote.user}@${remote.host}:${scriptRunNode}",) + // Upload waives.txt to Frontend node + def waivesListLocalPath = "${llmSrcLocal}/tests/integration/test_lists/waives.txt" + Utils.exec(pipeline, script: "sshpass -p '${remote.passwd}' scp -r -p -oStrictHostKeyChecking=no ${waivesListLocalPath} ${remote.user}@${remote.host}:${waivesListPathNode}",) + // Generate Test List and Upload to Frontend Node def makoArgs = getMakoArgsFromStageName(stageName, true) // TODO: currently the options will only be processed if the first @@ -349,6 +366,7 @@ def runLLMTestlistOnSlurm_MultiNodes(pipeline, platform, testList, config=VANILL "--container-image=${container}", "--container-workdir=/home/svc_tensorrt/bloom/scripts", "--container-mounts=${mounts}", + "--container-env=NVIDIA_IMEX_CHANNELS" ].join(" ") def scriptLaunch = "/home/svc_tensorrt/bloom/scripts/${jobUID}/slurm_launch.sh" @@ -362,12 +380,14 @@ def runLLMTestlistOnSlurm_MultiNodes(pipeline, platform, testList, config=VANILL export stageName=$stageName export testList=$testList export testListPathNode=$testListPathNode + export waivesListPathNode=$waivesListPathNode export pytestTestTimeout=$pytestTestTimeout export splits=$splits export splitId=$splitId export perfMode=$perfMode export resourcePathNode=$resourcePathNode export MODEL_CACHE_DIR=$MODEL_CACHE_DIR + export NVIDIA_IMEX_CHANNELS=0 chmod +x ${scriptRunNode} ${srunCmd} """.stripIndent() @@ -429,7 +449,7 @@ def EXTRA_STAGE_LIST = "extra_stage" @Field def MULTI_GPU_FILE_CHANGED = "multi_gpu_file_changed" @Field -def ONLY_PYTORCH_FILE_CHANGED = "only_pytorch_file_changed" +def ONLY_ONE_GROUP_CHANGED = "only_one_group_changed" @Field def AUTO_TRIGGER_TAG_LIST = "auto_trigger_tag_list" @Field @@ -437,8 +457,6 @@ def DEBUG_MODE = "debug" @Field def DETAILED_LOG = "detailed_log" @Field -def ONLY_DOCS_FILE_CHANGED = "only_docs_file_changed" -@Field def testFilter = [ (REUSE_STAGE_LIST): null, (ENABLE_SKIP_TEST): false, @@ -451,11 +469,10 @@ def testFilter = [ (DISABLE_MULTI_GPU_TEST): false, (EXTRA_STAGE_LIST): null, (MULTI_GPU_FILE_CHANGED): false, - (ONLY_PYTORCH_FILE_CHANGED): false, + (ONLY_ONE_GROUP_CHANGED): "", (DEBUG_MODE): false, (AUTO_TRIGGER_TAG_LIST): [], (DETAILED_LOG): false, - (ONLY_DOCS_FILE_CHANGED): false, ] @Field @@ -464,10 +481,13 @@ def GITHUB_PR_API_URL = "github_pr_api_url" def CACHED_CHANGED_FILE_LIST = "cached_changed_file_list" @Field def ACTION_INFO = "action_info" +@Field +def IMAGE_KEY_TO_TAG = "image_key_to_tag" def globalVars = [ (GITHUB_PR_API_URL): null, (CACHED_CHANGED_FILE_LIST): null, (ACTION_INFO): null, + (IMAGE_KEY_TO_TAG): [:], ] String getShortenedJobName(String path) @@ -480,6 +500,7 @@ String getShortenedJobName(String path) "L1_Custom": "l1-cus", "L1_Nightly": "l1-nt", "L1_Stable": "l1-stb", + "BuildDockerImageSanityTest": "img-check", ] def parts = path.split('/') // Apply nameMapping to the last part (jobName) @@ -1269,6 +1290,22 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO trtllm_utils.llmExecStepWithRetry(pipeline, script: "cd ${llmPath} && wget -nv ${llmTarfile}") sh "cd ${llmPath} && tar -zxf ${tarName}" + // Download the new merged waives.txt + def waivesTxt = "https://urm.nvidia.com/artifactory/${ARTIFACT_PATH}/waive_list/waives.txt" + try { + trtllm_utils.llmExecStepWithRetry(pipeline, script: "wget -nv ${waivesTxt}") + if (!fileExists("waives.txt")) { + error "There is no merged waives.txt file, use the default waives.txt." + } + sh "rm ${llmSrc}/tests/integration/test_lists/waives.txt" + sh "mv waives.txt ${llmSrc}/tests/integration/test_lists/waives.txt" + echo "Download merged waives.txt successfully" + } catch (InterruptedException e) { + throw e + } catch (Exception e) { + echo "Failed to download merged waives.txt, use the default waives.txt. Error: ${e.message}" + } + // install python package if (env.alternativeTRT) { sh "cd ${llmSrc} && sed -i 's#tensorrt~=.*\$#tensorrt#g' requirements.txt && cat requirements.txt" @@ -1704,6 +1741,24 @@ def runInKubernetes(pipeline, podSpec, containerName) def launchTestJobs(pipeline, testFilter, dockerNode=null) { def dockerArgs = "-v /mnt/scratch.trt_llm_data:/scratch.trt_llm_data:ro -v /tmp/ccache:${CCACHE_DIR}:rw -v /tmp/pipcache/http-v2:/root/.cache/pip/http-v2:rw --cap-add syslog" + + // IMPORTANT: Stage Configuration Syntax Requirement + // + // The test_to_stage_mapping.py script expects stage definitions in the following format: + // "Stage-Name": ["platform", "yaml_file", split_id, split_count, gpu_count] + // + // Where: + // - Stage-Name: Must be quoted string, used to identify the Jenkins stage + // - platform: Hardware platform identifier (e.g., "a10", "h100-cr") + // - yaml_file: Test database YAML filename without .yml extension (e.g., "l0_a10") + // - split_id: Current split number (1-based) + // - split_count: Total number of splits + // - gpu_count: Number of GPUs required (optional, defaults to 1) + // + // This format is parsed by scripts/test_to_stage_mapping.py to provide bidirectional + // mapping between test names and Jenkins stage names. Any changes to this syntax + // may break the mapping functionality. + x86TestConfigs = [ "DGX_H100-4_GPUs-PyTorch-DeepSeek-1": ["dgx-h100-x4", "l0_dgx_h100", 1, 2, 4], "DGX_H100-4_GPUs-PyTorch-DeepSeek-2": ["dgx-h100-x4", "l0_dgx_h100", 2, 2, 4], @@ -1718,6 +1773,7 @@ def launchTestJobs(pipeline, testFilter, dockerNode=null) "A10-TensorRT-4": ["a10", "l0_a10", 4, 6], "A10-TensorRT-5": ["a10", "l0_a10", 5, 6], "A10-TensorRT-6": ["a10", "l0_a10", 6, 6], + "A10-Nanobind": ["a10", "l0_a10_nanobind", 1, 1], "A30-Triton-1": ["a30", "l0_a30", 1, 1], "A30-PyTorch-1": ["a30", "l0_a30", 1, 2], "A30-PyTorch-2": ["a30", "l0_a30", 2, 2], @@ -1794,6 +1850,9 @@ def launchTestJobs(pipeline, testFilter, dockerNode=null) if (key.contains("llvm")) { config = LLVM_CONFIG } + if (key.contains("Nanobind")) { + config = NANOBIND_CONFIG + } runLLMTestlistOnPlatform(pipeline, values[0], values[1], config, key.contains("Perf"), key, values[2], values[3]) }]]} fullSet = parallelJobs.keySet() @@ -2163,22 +2222,28 @@ def launchTestJobs(pipeline, testFilter, dockerNode=null) println parallelJobsFiltered.keySet() } - if (testFilter[(ONLY_PYTORCH_FILE_CHANGED)]) { + if (testFilter[(ONLY_ONE_GROUP_CHANGED)] == "Docs") { + echo "Only docs files are changed, run doc build stage only." + parallelJobsFiltered = docBuildJobs + println parallelJobsFiltered.keySet() + } else if (testFilter[(ONLY_ONE_GROUP_CHANGED)] != "") { if (testFilter[(TEST_BACKEND)] != null) { - echo "Force disable ONLY_PYTORCH_FILE_CHANGED mode. Backend mode set by flag: ${testFilter[(TEST_BACKEND)]}." + echo "Force disable ONLY_ONE_GROUP_CHANGED mode. Backend mode set by flag: ${testFilter[(TEST_BACKEND)]}." } else { - echo "ONLY_PYTORCH_FILE_CHANGED mode is true." - parallelJobsFiltered = parallelJobsFiltered.findAll { !it.key.contains("-CPP-") && !it.key.contains("-TensorRT-") } + echo "ONLY_ONE_GROUP_CHANGED mode is true. The group is: ${testFilter[(ONLY_ONE_GROUP_CHANGED)]}." + def excludedBackends = new HashMap() + excludedBackends["PyTorch"] = ["-CPP-", "-TensorRT-", "-Triton-"] + excludedBackends["Triton"] = ["-PyTorch-", "-CPP-", "-TensorRT-"] + def group = testFilter[(ONLY_ONE_GROUP_CHANGED)] + if (excludedBackends.containsKey(group)) { + parallelJobsFiltered = parallelJobsFiltered.findAll { key, value -> + !excludedBackends[group].any { backend -> key.contains(backend) } + } + } println parallelJobsFiltered.keySet() } } - if (testFilter[(ONLY_DOCS_FILE_CHANGED)]) { - echo "Only docs files are changed, run doc build stage only." - parallelJobsFiltered = docBuildJobs - println parallelJobsFiltered.keySet() - } - // Check --stage-list, only run the stages in stage-list. if (testFilter[TEST_STAGE_LIST] != null) { echo "Use TEST_STAGE_LIST for filtering. Stages: ${testFilter[(TEST_STAGE_LIST)]}." @@ -2232,6 +2297,90 @@ def launchTestJobs(pipeline, testFilter, dockerNode=null) return parallelJobsFiltered } + + +def launchTestJobsForImagesSanityCheck(pipeline, globalVars) { + def testConfigs = [ + "NGC Devel Image amd64": [ + name: "NGC-Devel-Image-amd64-Sanity-Test", + k8sArch: "amd64", + wheelInstalled: false, + config: VANILLA_CONFIG, + ], + "NGC Devel Image arm64": [ + name: "NGC-Devel-Image-arm64-Sanity-Test", + k8sArch: "arm64", + wheelInstalled: false, + config: LINUX_AARCH64_CONFIG, + ], + "NGC Release Image amd64": [ + name: "NGC-Release-Image-amd64-Sanity-Test-A10", + gpuType: "a10", + k8sArch: "amd64", + wheelInstalled: true, + config: VANILLA_CONFIG, + ], + "NGC Release Image arm64": [ + name: "NGC-Release-Image-arm64-Sanity-Test-GH200", + gpuType: "gh200", + k8sArch: "arm64", + wheelInstalled: true, + config: LINUX_AARCH64_CONFIG, + ], + ] + if (!ENABLE_NGC_DEVEL_IMAGE_TEST) { + ["NGC Devel Image amd64", "NGC Devel Image arm64"].each { key -> + testConfigs.remove(key) + } + echo "NGC Devel Image test is disabled." + } + if (!ENABLE_NGC_RELEASE_IMAGE_TEST) { + ["NGC Release Image amd64", "NGC Release Image arm64"].each { key -> + testConfigs.remove(key) + } + echo "NGC Release Image test is disabled." + } + // Update testConfigs image field using the map from globalVars + testConfigs.each { key, config -> + if (globalVars[IMAGE_KEY_TO_TAG] && globalVars[IMAGE_KEY_TO_TAG][key]) { + config.image = globalVars[IMAGE_KEY_TO_TAG][key] + } + } + // Filter out all configs that don't have image set + testConfigs = testConfigs.findAll { key, config -> + return config.image != null + } + + echo "Filtered test configs with images:" + println testConfigs + + def testJobs = testConfigs.collectEntries { key, values -> [values.name, { + if (values.wheelInstalled) { + stage(values.name) { + echo "Run ${values.name} sanity test." + imageSanitySpec = createKubernetesPodConfig(values.image, values.gpuType, values.k8sArch) + trtllm_utils.launchKubernetesPod(pipeline, imageSanitySpec, "trt-llm", { + sh "env | sort" + trtllm_utils.llmExecStepWithRetry(pipeline, script: "apt-get update && apt-get install -y git rsync curl") + runLLMTestlistOnPlatform(pipeline, values.gpuType, "l0_sanity_check", values.config, false, values.name , 1, 1, true, null) + }) + } + } else { + stage(values.name) { + imageSanitySpec = createKubernetesPodConfig(values.image, "build", values.k8sArch) + trtllm_utils.launchKubernetesPod(pipeline, imageSanitySpec, "trt-llm", { + sh "env | sort" + def cpuArch = values.k8sArch == "amd64" ? X86_64_TRIPLE : AARCH64_TRIPLE + runLLMBuild(pipeline, cpuArch, false, "imageTest/") + }) + } + } + }]} + + return testJobs +} + + pipeline { agent { kubernetes createKubernetesPodConfig("", "agent") @@ -2274,7 +2423,10 @@ pipeline { when { expression { // Only run the test list validation when necessary - env.targetArch == X86_64_TRIPLE && testFilter[ONLY_DOCS_FILE_CHANGED] == false + env.targetArch == X86_64_TRIPLE && + testFilter[ONLY_ONE_GROUP_CHANGED] != "Docs" && + !(env.JOB_NAME ==~ /.*Multi-GPU.*/) && + !(env.JOB_NAME ==~ /.*BuildDockerImageSanityTest.*/) } } steps @@ -2287,7 +2439,11 @@ pipeline { stage("Test") { steps { script { - parallelJobs = launchTestJobs(this, testFilter) + if (env.JOB_NAME ==~ /.*BuildDockerImageSanityTest.*/) { + parallelJobs = launchTestJobsForImagesSanityCheck(this, globalVars) + } else { + parallelJobs = launchTestJobs(this, testFilter) + } singleGpuJobs = parallelJobs dgxJobs = [:] @@ -2299,17 +2455,47 @@ pipeline { dgxJobs = parallelJobs.findAll{dgxSigns.any{sign -> it.key.contains(sign)}} } - if (singleGpuJobs.size() > 0) { - singleGpuJobs.failFast = params.enableFailFast - parallel singleGpuJobs - } else { - echo "Skip single-GPU testing. No test to run." - } - - if (dgxJobs.size() > 0) { - stage(testPhase2StageName) { + if (env.JOB_NAME ==~ /.*Single-GPU.*/) { + echo "Only run single-GPU tests." + if (dgxJobs.size() > 0) { + if (globalVars[ACTION_INFO]['parents'].size() > 0) { + // We add a special marker to the parent job's description. + // This will be used to decide whether to run multi-GPU test stage. + def parentJob = globalVars[ACTION_INFO]['parents'][-2] + trtllm_utils.appendBuildDescription(this, parentJob['name'], parentJob['build_number'], "====Require Multi-GPU Testing====
") + } else { + echo "No parent job found to add the special marker for executing multi-GPU test stage." + } + } else { + echo "Skip multi-GPU testing. No test to run." + } + if (singleGpuJobs.size() > 0) { + singleGpuJobs.failFast = params.enableFailFast + parallel singleGpuJobs + } else { + echo "Skip single-GPU testing. No test to run." + } + } else if (env.JOB_NAME ==~ /.*Multi-GPU.*/) { + echo "Only run multi-GPU tests." + if (dgxJobs.size() > 0) { dgxJobs.failFast = params.enableFailFast parallel dgxJobs + } else { + error "Skip multi-GPU testing. No test to run." + } + } else { + if (singleGpuJobs.size() > 0) { + singleGpuJobs.failFast = params.enableFailFast + parallel singleGpuJobs + } else { + echo "Skip single-GPU testing. No test to run." + } + + if (dgxJobs.size() > 0) { + stage(testPhase2StageName) { + dgxJobs.failFast = params.enableFailFast + parallel dgxJobs + } } } } diff --git a/jenkins/current_image_tags.properties b/jenkins/current_image_tags.properties index 11552a7f131e..6e4863a11edf 100644 --- a/jenkins/current_image_tags.properties +++ b/jenkins/current_image_tags.properties @@ -8,7 +8,10 @@ # NB: Although string interpolation is supported, redundant substrings are # kept in the variables below for interoperability with # scripts/rename_docker_images.py -LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.05-py3-x86_64-ubuntu24.04-trt10.11.0.33-skip-tritondevel-202507101530-5434 -LLM_SBSA_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.05-py3-aarch64-ubuntu24.04-trt10.11.0.33-skip-tritondevel-202507101530-5434 -LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py310-trt10.11.0.33-skip-tritondevel-202507101530-5434 -LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py312-trt10.11.0.33-skip-tritondevel-202507101530-5434 +# +# NB: Typically, the suffix indicates the PR whose CI pipeline generated the images. In case that +# images are adopted from PostMerge pipelines, the abbreviated commit hash is used instead. +LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.05-py3-x86_64-ubuntu24.04-trt10.11.0.33-skip-tritondevel-202507162011-ec3ebae +LLM_SBSA_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.05-py3-aarch64-ubuntu24.04-trt10.11.0.33-skip-tritondevel-202507162011-ec3ebae +LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py310-trt10.11.0.33-skip-tritondevel-202507162011-ec3ebae +LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py312-trt10.11.0.33-skip-tritondevel-202507162011-ec3ebae diff --git a/jenkins/scripts/mergeWaiveList.py b/jenkins/scripts/mergeWaiveList.py new file mode 100644 index 000000000000..418a603b9af7 --- /dev/null +++ b/jenkins/scripts/mergeWaiveList.py @@ -0,0 +1,56 @@ +import argparse +import sys + +# Generate the merged waive list: +# 1. Parse the current MR waive list, and get the removed lines from the diff +# 2. Parse the TOT waive list +# 3. Merge the current MR waive list and TOT waive list, and remove the removed lines from the step 1 + + +def get_remove_lines_from_diff(diff): + lines = diff.split('\n') + remove_lines = [ + line[1:] + '\n' for line in lines + if len(line) > 1 and line.startswith('-') + ] + return remove_lines + + +def parse_waive_txt(waive_txt): + with open(waive_txt, 'r') as f: + lines = f.readlines() + waive_list = [line for line in lines if line.strip()] + return waive_list + + +def write_waive_list(waive_list, output_file): + with open(output_file, 'w') as f: + for line in waive_list: + f.write(line) + + +def merge_waive_list(cur_list, main_list, remove_lines, output_file): + merged = list(dict.fromkeys(cur_list + main_list)) + for line in reversed(remove_lines): + for i in range(len(merged) - 1, -1, -1): + if merged[i] == line: + merged.pop(i) + break + write_waive_list(merged, output_file) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--cur-waive-list', + required=True, + help='Current waive list') + parser.add_argument('--latest-waive-list', + required=True, + help='Latest waive list') + parser.add_argument('--diff', required=True, help='Diff of the waive list') + parser.add_argument('--output-file', required=True, help='Output file') + args = parser.parse_args(sys.argv[1:]) + cur_list = parse_waive_txt(args.cur_waive_list) + main_list = parse_waive_txt(args.latest_waive_list) + remove_lines = get_remove_lines_from_diff(args.diff) + merge_waive_list(cur_list, main_list, remove_lines, args.output_file) diff --git a/jenkins/scripts/slurm_run.sh b/jenkins/scripts/slurm_run.sh index 9c055d8cd34e..4b6337fca5de 100755 --- a/jenkins/scripts/slurm_run.sh +++ b/jenkins/scripts/slurm_run.sh @@ -45,6 +45,7 @@ testCmdLines=( "-v" "--timeout=$pytestTestTimeout" "--test-list=$testListPathNode" + "--waives-file=$waivesListPathNode" "--rootdir $llmSrcNode/tests/integration/defs" "--test-prefix=$stageName" "--splits $splits" diff --git a/requirements.txt b/requirements.txt index 88bf8552f7ba..16c1e4b5f8ca 100644 --- a/requirements.txt +++ b/requirements.txt @@ -25,12 +25,13 @@ tensorrt~=10.11.0 # https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-25-05.html#rel-25-05 uses 2.8.0a0. torch>=2.7.1,<=2.8.0a0 torchvision -nvidia-modelopt[torch]~=0.31.0 +nvidia-modelopt[torch]~=0.33.0 nvidia-nccl-cu12 nvidia-cuda-nvrtc-cu12 transformers==4.53.1 pydantic>=2.9.1 -pydantic-settings +pydantic-settings[yaml] +omegaconf pillow==10.3.0 wheel<=0.45.1 optimum @@ -59,3 +60,4 @@ ninja etcd3 blake3 llguidance==0.7.29 +soundfile diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index 06a22d93ff8f..3fdaa93febb2 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -144,7 +144,7 @@ def setup_venv(project_dir: Path, requirements_file: Path, no_venv: bool): # Ensure PyPI PyTorch is not installed in the venv purelib_dir = Path(scheme["purelib"]) pytorch_package_dir = purelib_dir / "torch" - if venv_prefix != sys.base_prefix and pytorch_package_dir.exists(): + if str(venv_prefix) != sys.base_prefix and pytorch_package_dir.exists(): warnings.warn( f"Using the NVIDIA PyTorch container with PyPI distributed PyTorch may lead to compatibility issues.\n" f"If you encounter any problems, please delete the environment at `{venv_prefix}` so that " @@ -298,7 +298,6 @@ def main(*, install: bool = False, skip_building_wheel: bool = False, linking_install_binary: bool = False, - python_bindings: bool = True, binding_type: str = "pybind", benchmarks: bool = False, micro_benchmarks: bool = False, @@ -860,11 +859,6 @@ def add_arguments(parser: ArgumentParser): "--linking_install_binary", action="store_true", help="Install the built binary by symbolic linking instead of copying.") - parser.add_argument( - "--python_bindings", - "-p", - action="store_true", - help="(deprecated) Build the python bindings for the C++ runtime.") parser.add_argument("--binding_type", choices=["pybind", "nanobind"], default="pybind", diff --git a/scripts/dco_check.py b/scripts/dco_check.py index dedd1a0b9c97..1fbe509ccc58 100755 --- a/scripts/dco_check.py +++ b/scripts/dco_check.py @@ -22,7 +22,7 @@ def commit_message_has_signoff(message): def main(): if len(sys.argv) != 2: - print("Usage: python commit-msg.py ") + print("Usage: python dco_check.py ") sys.exit(1) # Read the commit message from the file passed as an argument by Git diff --git a/scripts/test_to_stage_mapping.py b/scripts/test_to_stage_mapping.py new file mode 100644 index 000000000000..d51623a80c9d --- /dev/null +++ b/scripts/test_to_stage_mapping.py @@ -0,0 +1,266 @@ +"""Lookup Jenkins stage names for integration tests and vice versa. + +This helper parses ``jenkins/L0_Test.groovy`` and the YAML files under +``tests/integration/test_lists/test-db`` to provide a bidirectional mapping +between test names and Jenkins stage names. When ``--tests`` or ``--test-list`` +options are used, each value is treated as a substring pattern. Any test whose +fully qualified name contains the pattern will be matched. If the pattern +corresponds exactly to a test name, it naturally matches that test as well. + +Example usage:: + + python scripts/test_to_stage_mapping.py --tests \\ + "triton_server/test_triton.py::test_gpt_ib_ptuning[gpt-ib-ptuning]" + python scripts/test_to_stage_mapping.py --tests gpt_ib_ptuning + python scripts/test_to_stage_mapping.py --stages \\ + A100X-Triton-Post-Merge-1 + +Tests can also be provided via ``--test-list`` pointing to either a plain text +file or a YAML list file. Quote individual test names on the command line so +the shell does not interpret ``[`` and ``]`` characters. +""" + +import argparse +import os +import re +from collections import defaultdict +from glob import glob +from typing import List + +import yaml + + +def _load_tests_file(path: str) -> List[str]: + tests: List[str] = [] + yaml_mode = path.endswith('.yml') or path.endswith('.yaml') + with open(path, 'r') as f: + for line in f: + line = line.strip() + if not line or line.startswith('#'): + continue + if yaml_mode: + if line.startswith('- '): + tests.append(line[2:].strip()) + else: + tests.append(line) + return tests + + +# Regex to parse Jenkins stage configurations from Groovy files +# Matches patterns like: "Stage-Name": ["platform", "yaml_file", split_id, split_count, gpu_count] +# +# Pattern breakdown: +# "(?P[^"]+)" - Captures stage name in quotes (group 'stage') +# \s*:\s* - Matches colon with optional whitespace +# \[ - Matches opening bracket +# "[^"]+" - Matches platform string in quotes (ignored) +# ,\s* - Matches comma with optional whitespace +# "(?P[^"]+)" - Captures yaml filename in quotes (group 'yml') +# (?:,\s*\d+)* - Matches zero or more comma-separated numbers (split_id, split_count, gpu_count) +# \s*\] - Matches closing bracket with optional whitespace +_STAGE_RE = re.compile( + r'"(?P[^"]+)"\s*:\s*\["[^"]+",\s*"(?P[^"]+)"(?:,\s*\d+)*\s*\]') + + +def _extract_terms(entry): + """Extract terms from either direct 'terms' or 'condition.terms'.""" + terms = entry.get('terms', {}) + if not terms: + terms = entry.get('condition', {}).get('terms', {}) + return terms + + +class StageQuery: + + def __init__(self, groovy_path: str, test_db_dir: str): + self.stage_to_yaml, self.yaml_to_stages = self._parse_stage_mapping( + groovy_path) + self.test_map, self.yaml_stage_tests = self._parse_tests(test_db_dir) + # Build dynamic backend mapping from discovered data + self._backend_keywords = self._discover_backend_keywords() + + @staticmethod + def _parse_stage_mapping(path): + stage_to_yaml = {} + yaml_to_stages = defaultdict(list) + with open(path, 'r') as f: + for line in f: + m = _STAGE_RE.search(line) + if m: + stage = m.group('stage') + yml = m.group('yml') + '.yml' + stage_to_yaml[stage] = yml + yaml_to_stages[yml].append(stage) + return stage_to_yaml, yaml_to_stages + + def _parse_tests(self, db_dir): + """Parse tests from YAML files, supporting both .yml and .yaml.""" + test_map = defaultdict(list) + yaml_stage_tests = defaultdict(lambda: defaultdict(list)) + + yaml_files = (glob(os.path.join(db_dir, '*.yml')) + + glob(os.path.join(db_dir, '*.yaml'))) + + for path in yaml_files: + with open(path, 'r') as f: + data = yaml.safe_load(f) + for key, entries in data.items(): + if key == 'version' or entries is None: + continue + for entry in entries: + terms = _extract_terms(entry) + + stage = terms.get('stage') + if stage is None: + continue + + backend = terms.get('backend', '') # Default to empty + + tests = entry.get('tests', []) + yml = os.path.basename(path) + for t in tests: + test_map[t].append((yml, stage, backend)) + yaml_stage_tests[yml][stage].append(t) + return test_map, yaml_stage_tests + + def _discover_backend_keywords(self): + """Discover backend keywords from existing data dynamically.""" + backend_keywords = {} + + # Collect all backends from test data + all_backends = set() + for mappings in self.test_map.values(): + for yml, stage_type, backend in mappings: + if backend and backend.strip(): + all_backends.add(backend.strip().lower()) + + # Map backends to their likely stage name keywords + for backend in all_backends: + backend_keywords[backend] = backend.upper() + + # Add common variations/aliases + aliases = { + 'tensorrt': ['TENSORRT', 'TRT'], + 'pytorch': ['PYTORCH', 'TORCH'], + 'cpp': ['CPP', 'C++'], + 'triton': ['TRITON'] + } + + for backend, keywords in aliases.items(): + if backend in backend_keywords: + backend_keywords[backend] = keywords + + return backend_keywords + + def search_tests(self, pattern: str): + parts = pattern.split() + result = [] + for test in self.test_map: + name = test.lower() + if all(p.lower() in name for p in parts): + result.append(test) + return result + + def tests_to_stages(self, tests): + result = set() + for t in tests: + for yml, stage_type, backend in self.test_map.get(t, []): + for s in self.yaml_to_stages.get(yml, []): + if stage_type == 'post_merge' and 'Post-Merge' not in s: + continue + if stage_type == 'pre_merge' and 'Post-Merge' in s: + continue + + # Filter by backend if specified + if backend and backend != '': + backend_keywords = self._backend_keywords.get( + backend.lower(), [backend.upper()]) + if isinstance(backend_keywords, str): + backend_keywords = [backend_keywords] + + if not any(keyword in s.upper() + for keyword in backend_keywords): + continue + + result.add(s) + return sorted(result) + + def stages_to_tests(self, stages): + result = set() + for s in stages: + yml = self.stage_to_yaml.get(s) + if not yml: + continue + stage_type = 'post_merge' if 'Post-Merge' in s else 'pre_merge' + + # Determine expected backend dynamically from stage name + expected_backend = None + stage_upper = s.upper() + for backend, keywords in self._backend_keywords.items(): + if isinstance(keywords, str): + keywords = [keywords] + if any(keyword in stage_upper for keyword in keywords): + expected_backend = backend + break + + # Get all tests for yml/stage_type, then filter by backend + all_tests = self.yaml_stage_tests.get(yml, {}).get(stage_type, []) + for test in all_tests: + # Check if test's backend matches stage's expected backend + test_mappings = self.test_map.get(test, []) + for test_yml, test_stage, test_backend in test_mappings: + if (test_yml == yml and test_stage == stage_type + and (expected_backend is None + or test_backend == expected_backend)): + result.add(test) + break + return sorted(result) + + +def main(): + parser = argparse.ArgumentParser( + description='Map Jenkins stages to tests and vice versa.') + group = parser.add_mutually_exclusive_group(required=True) + group.add_argument( + '--tests', + nargs='+', + help='One or more test name patterns to resolve to Jenkins stages') + group.add_argument( + '--test-list', + help=('File with test name patterns, either newline separated ' + 'or a YAML list')) + group.add_argument('--stages', + nargs='+', + help='List of stage names to look up') + parser.add_argument('--repo-root', + default=os.path.dirname(os.path.dirname(__file__)), + help='Path to repository root') + args = parser.parse_args() + + groovy = os.path.join(args.repo_root, 'jenkins', 'L0_Test.groovy') + db_dir = os.path.join(args.repo_root, 'tests', 'integration', 'test_lists', + 'test-db') + query = StageQuery(groovy, db_dir) + + if args.tests or args.test_list: + patterns = [] + if args.tests: + patterns.extend(args.tests) + if args.test_list: + patterns.extend(_load_tests_file(args.test_list)) + + collected = [] + for pat in patterns: + collected.extend(query.search_tests(pat)) + tests = sorted(set(collected)) + stages = query.tests_to_stages(tests) + for s in stages: + print(s) + else: + tests = query.stages_to_tests(args.stages) + for t in tests: + print(t) + + +if __name__ == '__main__': + main() diff --git a/setup.py b/setup.py index 38c24c13bb19..c436dfd834bc 100644 --- a/setup.py +++ b/setup.py @@ -115,6 +115,7 @@ def has_ext_modules(self): 'tools/plugin_gen/templates/*', 'bench/build/benchmark_config.yml', 'evaluate/lm_eval_tasks/**/*', + "_torch/auto_deploy/config/*.yaml", ] @@ -185,7 +186,7 @@ def extract_from_precompiled(precompiled_location: str, package_data: List[str], with zipfile.ZipFile(wheel_path) as wheel: for file in wheel.filelist: - if file.filename.endswith(".py"): + if file.filename.endswith((".py", ".yaml")): continue for filename_pattern in package_data: if fnmatch.fnmatchcase(file.filename, diff --git a/tensorrt_llm/_torch/__init__.py b/tensorrt_llm/_torch/__init__.py index 7d2de6d643ce..23257d915047 100644 --- a/tensorrt_llm/_torch/__init__.py +++ b/tensorrt_llm/_torch/__init__.py @@ -1,4 +1,5 @@ from .llm import LLM from .model_config import MoeLoadBalancerConfig +from .models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader -__all__ = ["LLM", "MoeLoadBalancerConfig"] +__all__ = ["LLM", "MoeLoadBalancerConfig", "BaseCheckpointLoader"] diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index c62fa0e15579..463658bde633 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -297,10 +297,16 @@ def prepare(self) -> None: self._positions[:positions.size(0)].copy_(positions, non_blocking=True) - for plan_params in self._plan_params_to_wrappers: - # Re-plan the cached wrappers for a new set of requests. - self._plan_params_to_wrappers[plan_params].is_planned = False - self._plan_with_params(plan_params) + # Generally, plan_params with non-trivial attention_mask_data are relevant only the + # corresponding forward pass. So, flush them out here as they won't be relevant for + # subsequent forward calls. + for plan_params in list(self._plan_params_to_wrappers.keys()): + if plan_params.attention_mask_data is None: + # Re-plan the cached wrappers for a new set of requests. + self._plan_params_to_wrappers[plan_params].is_planned = False + self._plan_with_params(plan_params) + else: + del self._plan_params_to_wrappers[plan_params] if self.cross is not None and self.cross is not self: self.cross.prepare() @@ -426,7 +432,7 @@ def decode_plan(): kv_data_type=plan_params.kv_dtype, ) - # Must sync after append_paged_kv_cache and before plan + # Must sync after append_paged_kv_cache and before plan. torch.cuda.current_stream().synchronize() if self.num_contexts > 0: diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 3a3410e8a9c0..a50d475681b9 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -135,6 +135,9 @@ class AttentionMetadata: _num_ctx_tokens: int = field(init=False, default=0, repr=False) _num_tokens: int = field(init=False, default=0, repr=False) + # This buffer is currently only used for TrtllmAttentionMetadata. + cache_indirection: Optional[torch.Tensor] = None + def __post_init__(self) -> None: if self.is_cross: assert self.cross is None or self.cross is self, "Cross attention metadata should not have sub metadata" @@ -351,6 +354,8 @@ class RopeParams: beta_slow: int = 1 mscale: float = 1.0 mscale_all_dim: float = 0.0 + short_factor: Optional[Tuple[float]] = None + long_factor: Optional[Tuple[float]] = None @staticmethod def from_config(config) -> "RopeParams": @@ -386,12 +391,18 @@ def from_config(config) -> "RopeParams": "low_freq_factor", 1.0) rope_params.high_freq_factor = rope_scaling.get( "high_freq_factor", 4.0) - rope_params.original_max_positions = rope_scaling.get( - "original_max_position_embeddings", 1024) + rope_params.original_max_positions = getattr( + config, + "original_max_position_embeddings", None) or rope_scaling.get( + "original_max_position_embeddings", None) or 1024 rope_params.beta_fast = rope_scaling.get("beta_fast", 32) rope_params.beta_slow = rope_scaling.get("beta_slow", 1) rope_params.mscale = rope_scaling.get("mscale", 1.0) rope_params.mscale_all_dim = rope_scaling.get("mscale_all_dim", 0.0) + if "short_factor" in rope_scaling: + rope_params.short_factor = tuple(rope_scaling["short_factor"]) + if "long_factor" in rope_scaling: + rope_params.long_factor = tuple(rope_scaling["long_factor"]) # Workaround for DeepSeek V3 Lite since its rope_scaling is null in config.json. elif config.model_type == "deepseek_v3": rope_params.scale_type = RotaryScalingType.yarn @@ -428,7 +439,14 @@ def create_rope_const_params(self, interleave: bool = True): self.mscale_all_dim, ) elif self.scale_type == RotaryScalingType.longrope: - raise NotImplementedError("Long RoPE is not supported.") + rope_inv_freq, rope_cos_sin = RopeEmbeddingUtils.create_sinusoidal_positions_long_rope_for_attention_plugin( + num_pos=self.max_positions, + dim=self.dim, + theta=self.theta, + original_max_pos=self.original_max_positions, + short_factor=self.short_factor, + long_factor=self.long_factor, + ) else: rope_inv_freq, rope_cos_sin = RopeEmbeddingUtils.create_sinusoidal_positions_for_attention_plugin( self.max_positions, diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index b23ed0a84ff4..143fae88d62e 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -517,10 +517,9 @@ def is_nvfp4_output_kernel_available( class TrtllmAttentionMetadata(AttentionMetadata): workspace: Optional[torch.Tensor] = None - # TrtllmAttention needs to know the beam width and access to the cache indirection buffer, + # TrtllmAttention needs to know the beam width to access to the cache indirection buffer, # when beam search is enabled. beam_width: int = 1 - cache_indirection: Optional[torch.Tensor] = None # TrtllmAttention needs to know the max sequence length. # Implemented as a property to support no cache mode. diff --git a/tensorrt_llm/_torch/auto_deploy/__init__.py b/tensorrt_llm/_torch/auto_deploy/__init__.py index 3043228f98d5..7650b2dde698 100644 --- a/tensorrt_llm/_torch/auto_deploy/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/__init__.py @@ -1,5 +1,5 @@ # import submodules that require registration process -from . import compile, custom_ops, models, shim # noqa: F401 +from . import compile, custom_ops, export, models, shim # noqa: F401 # import AutoDeploy LLM and LlmArgs from .llm import * diff --git a/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py b/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py index 71bc5d44fdb2..0b309ae2bf89 100644 --- a/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py +++ b/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py @@ -35,10 +35,11 @@ def __init__( self._out_buffer_flat: List[torch.Tensor] = None self._args_hash: Optional[Tuple[int, ...]] = None self.cuda_graph_batch_sizes = ( - cuda_graph_batch_sizes + sorted(cuda_graph_batch_sizes, reverse=True) if cuda_graph_batch_sizes is not None else self._get_graph_batch_sizes(self.max_batch_size) ) + self._cuda_graph_mem_pool = None def _get_hash(self, flat_args: List[Any]) -> Tuple[int, ...]: return tuple(hash(a) for a in flat_args) @@ -64,7 +65,7 @@ def _capture_one_graph(self, *args, **kwargs) -> torch.cuda.CUDAGraph: # capture graph now torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): + with torch.cuda.graph(graph, pool=self._cuda_graph_mem_pool): # compute output out = self.model(*args, **kwargs) # write out into output buffer up to out batch size @@ -73,7 +74,7 @@ def _capture_one_graph(self, *args, **kwargs) -> torch.cuda.CUDAGraph: for o_buffer, o in zip(self._out_buffer_flat, out_flat): o_buffer[: o.shape[0]] = o torch.cuda.synchronize() - + self._cuda_graph_mem_pool = self._cuda_graph_mem_pool or graph.pool() return graph @staticmethod @@ -88,7 +89,7 @@ def _get_graph_batch_sizes( batch_sizes.update(range(multiplier, max_bs + 1, multiplier)) # return as sorted list - return sorted(batch_sizes) + return sorted(batch_sizes, reverse=True) def capture_graph(self, *args, **kwargs): """Capture and pre-fetch the graph for variable batch size.""" @@ -118,6 +119,7 @@ def capture_graph(self, *args, **kwargs): # capture output once with max batch size to capture output buffers with CudaGraphWarmUpPhase(): + ad_logger.info(f"Warm up with {self.max_batch_size=} before graph capture") out = self.model(*args, **kwargs) self._out_buffer_flat, out_spec = tree_flatten(out) assert out_spec == self._out_spec, "Output spec mismatch." diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml new file mode 100644 index 000000000000..5908c1271e42 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -0,0 +1,21 @@ +# Additional default args for AutoDeployConfig/LlmArgs in _torch/auto_deploy/llm_args.py +transforms: + build_model: + stage: factory + device: meta + # nothing to clean up + run_graph_cleanup: false + requires_clean_graph: false + export_to_gm: + stage: export + clone_state_dict: false + strict: false + # nothing to clean up + run_graph_cleanup: false + requires_clean_graph: false + cleanup_noop_slice: + stage: post_export + cleanup_noop_add: + stage: post_export + cleanup_input_constraints: + stage: post_export diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/__init__.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/__init__.py index f80d1e5ca918..23a80b94d743 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/__init__.py @@ -7,7 +7,9 @@ from .linear import * from .mla import * from .quant import * +from .rms_norm import * from .torch_attention import * +from .torch_backend_attention import * from .torch_moe import * from .torch_rope import * from .triton_attention import * diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/_triton_attention_internal.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/_triton_attention_internal.py index 18452d3b4175..f1d6e61932e4 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/_triton_attention_internal.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/_triton_attention_internal.py @@ -100,6 +100,8 @@ def _paged_generate_mha( n_heads, d_head, SEQ_BLOCK_SIZE, + False, + None, ) @@ -338,6 +340,7 @@ def _generate_mha_rope_fusion( d_head, SEQ_BLOCK_SIZE, HEAD_BLOCK_SIZE, + -1, ) attention_kv_stage2[(b, n_heads, 1)]( stage1_output_values, @@ -348,6 +351,8 @@ def _generate_mha_rope_fusion( n_heads, d_head, SEQ_BLOCK_SIZE, + False, + None, ) @@ -414,7 +419,9 @@ def _flattened_context_mha_rope_fusion( d_head, SEQ_BLOCK, max_cache_seq_len, - num_stages=2, + -1, + False, + None, ) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index c9a964eaec0b..13c91652bff4 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -117,14 +117,20 @@ def __post_init__(self): # if the provided max_num_tokens is less than the max_batch_size * max_seq_len, # we use the provided max_num_tokens to calculate the number of pages total_tokens = min(self.max_num_tokens, self.max_batch_size * max_seq_len_adjusted) - self._num_pages = (total_tokens) // self.page_size + (total_tokens % self.page_size > 0) + # Num pages can not be less than max_batch_size. + self._num_pages = max( + self.max_batch_size, + (total_tokens) // self.page_size + (total_tokens % self.page_size > 0), + ) self.input_ids = torch.ones(self.max_batch_size, 1, dtype=torch.int) self.position_ids = torch.zeros(self.max_batch_size, 1, dtype=torch.long) self.seq_len = torch.empty(self.max_batch_size, dtype=torch.int) self.input_pos = torch.empty_like(self.seq_len) self.cache_loc = torch.empty(self.num_pages, dtype=torch.int) self.pages_per_seq = torch.empty_like(self.seq_len) - + assert self.num_pages >= self.max_batch_size, ( + "num_pages must be greater than max_batch_size" + ) # dynamic shape descriptors for tensor args self._dynamic_shapes: Optional[Tuple[Dict[str, Dim]]] = None @@ -378,10 +384,11 @@ def set_generate_only_batch(self) -> None: def _update_position_ids(self) -> None: # set new position_ids as new tensor from input_pos and seq_len via torch.arange position_ids_list = [ - torch.arange(in_pos, in_pos + seq_len, dtype=torch.long) + num for in_pos, seq_len in zip(self.input_positions, self.sequence_lengths) + for num in range(in_pos, in_pos + seq_len) ] - self.position_ids = torch.cat(position_ids_list, dim=0).to(self.device) + self.position_ids = torch.tensor(position_ids_list, dtype=torch.long).to(self.device) # use [b,1] shape to indicate generate-only batch, otherwise use [1,total_len] if self.is_generate: @@ -398,13 +405,15 @@ def nest_sequences(self, input_ids: Sequence[Sequence[int]]) -> None: seq_lens = [len(ids) for ids in input_ids] self.seq_len.zero_() self.seq_len[: len(seq_lens)].copy_(torch.tensor(seq_lens), non_blocking=True) - + # We'll preserve the dtype of the input_ids tensor if it is a tensor, otherwise we'll use int + dtype = input_ids.dtype if isinstance(input_ids, torch.Tensor) else torch.int # set new input_ids as new tensor from flattened input_ids - ids_tnsr_list = [ - lst.detach() if isinstance(lst, torch.Tensor) else torch.tensor(lst, dtype=torch.int) + ids_list = [ + val for lst in input_ids + for val in (lst.detach().tolist() if isinstance(lst, torch.Tensor) else lst) ] - self.input_ids = torch.cat(ids_tnsr_list, dim=0).to(self.device) + self.input_ids = torch.tensor(ids_list, dtype=dtype).to(self.device) # set derivative properties self._sequence_lengths = seq_lens diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py new file mode 100644 index 000000000000..cd23ce7519b4 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py @@ -0,0 +1,82 @@ +"""Custom operator for FlashInfer and Triton RMSNorm implementation.""" + +import flashinfer +import torch + +from .triton_kernels.rms_norm import rms_norm + + +@torch.library.custom_op("auto_deploy::flashinfer_rms_norm", mutates_args=()) +def flashinfer_rmsnorm(input: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """Custom operator for FlashInfer RMSNorm implementation. + + Args: + input: Input tensor to normalize. + weight: Scaling weights for the normalized output. + eps: Small constant for numerical stability. + + Returns: + Normalized and scaled tensor using FlashInfer implementation. + """ + # Flashinfer rmsnorm expects a 2D input + input_flat = input.reshape(-1, input.shape[-1]) + rmsnorm_flat = flashinfer.norm.rmsnorm(input_flat, weight, eps) + return rmsnorm_flat.reshape(input.shape) + + +@flashinfer_rmsnorm.register_fake +def _(input: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """Fake implementation for the custom operator during tracing. + + Args: + input: Input tensor to normalize. + weight: Scaling weights for the normalized output. + eps: Small constant for numerical stability. + + Returns: + Empty tensor with same shape as input. + """ + return torch.empty_like(input) + + +@torch.library.custom_op("auto_deploy::triton_rms_norm", mutates_args=()) +def triton_rmsnorm(input: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """Custom operator for Triton RMSNorm implementation. + + Args: + input: Input tensor to normalize. + weight: Scaling weights for the normalized output. + eps: Small constant for numerical stability. + + Returns: + Normalized and scaled tensor using Triton implementation. + """ + return rms_norm(input, weight, eps) + + +@triton_rmsnorm.register_fake +def _(input: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """Fake implementation for the custom operator during tracing.""" + return torch.empty_like(input) + + +@torch.library.custom_op("auto_deploy::torch_rmsnorm", mutates_args=()) +def torch_rmsnorm(input: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """Custom operator for Torch RMSNorm implementation. + + Args: + input: Input tensor to normalize. + weight: Scaling weights for the normalized output. + eps: Small constant for numerical stability. + """ + input_dtype = input.dtype + input = input.to(torch.float32) + variance = input.pow(2).mean(-1, keepdim=True) + input = input * torch.rsqrt(variance + eps) + return weight * input.to(input_dtype) + + +@torch_rmsnorm.register_fake +def _(input: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """Fake implementation for the custom operator during tracing.""" + return torch.empty_like(input) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_attention.py index 6764ca3d91e2..68175233f91f 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_attention.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_attention.py @@ -7,6 +7,8 @@ import torch.nn as nn import torch.nn.functional as F +# TODO (nvchenghaoz): Remove related kernels once we have a backend-specific implementation for attention. + @torch.library.custom_op("auto_deploy::torch_attention_repeat_kv", mutates_args=()) def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: @@ -113,6 +115,9 @@ def bsnd_grouped_sdpa( dropout_p: float = 0.0, is_causal: bool = False, scale: Optional[float] = None, + sinks: Optional[torch.Tensor] = None, + sliding_window: Optional[int] = None, + logit_cap: Optional[float] = None, ) -> torch.Tensor: """Attention that assumes the input layout is bsnd. @@ -132,7 +137,16 @@ def bsnd_grouped_sdpa( @bsnd_grouped_sdpa.register_fake def bsnd_grouped_sdpa_fake( - query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + scale=None, + sinks=None, + sliding_window=None, + logit_cap=None, ): """Fake implementation of bnsd grouped SDPA.""" return query.new_empty(*query.shape[:-1], value.shape[-1]).contiguous() diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.py new file mode 100644 index 000000000000..9eccd0c83a9e --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.py @@ -0,0 +1,495 @@ +"""Torch backend attention using pure PyTorch reference implementations.""" + +import math +from typing import List, Optional, Tuple + +import torch +from torch._ops import OpOverloadPacket +from torch._subclasses import FakeTensor +from torch.fx import Node + +from ..utils.logger import ad_logger +from ..utils.node_utils import extract_op_args +from .attention_interface import ( + AttentionDescriptor, + AttentionLayout, + AttentionRegistry, + BufferInitializerDict, + CacheConfig, + CacheInitializerDict, + Constant, + MHACallable, + PrepareMetadataCallable, + SequenceInfo, +) +from .torch_attention import repeat_kv, update_kv_cache + + +def _apply_logit_softcapping(attn_scores: torch.Tensor, logit_cap: Optional[float]) -> torch.Tensor: + """Apply logit softcapping using the formula: logit_cap * tanh(logits / logit_cap)""" + if logit_cap is not None and logit_cap > 0.0: + return logit_cap * torch.tanh(attn_scores / logit_cap) + return attn_scores + + +def _torch_generate_mha( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_loc: torch.Tensor, + input_pos: torch.Tensor, + scale: float, + out: torch.Tensor, + logit_cap: Optional[float] = None, + sliding_window_size: Optional[int] = None, + sinks: Optional[torch.Tensor] = None, +): + """Generate-only attention (single token per sequence) using manual computation with existing update_kv_cache.""" + b, s, n_heads, head_dim = q.shape # q has shape (b, 1, n_heads, head_dim) in generate phase + assert s == 1, f"Expected sequence length 1 for generate phase, got {s}" + n_kv_heads = k.shape[2] # k has shape (b, 1, n_kv_heads, head_dim) + + # Update KV cache for single token + for i in range(b): + cache_idx = cache_loc[i].item() + pos = input_pos[i].item() + k_cache[cache_idx, pos] = k[i, 0] # Remove sequence dim + v_cache[cache_idx, pos] = v[i, 0] # Remove sequence dim + + # Compute attention for each sequence using manual computation + for i in range(b): + cache_idx = cache_loc[i].item() + pos = input_pos[i].item() + + # Get query, key, value for this sequence + q_i = q[i, 0] # [n_heads, head_dim] + + # Apply sliding window: limit the range of keys/values we attend to + if sliding_window_size is not None and sliding_window_size > 0: + # Sliding window: attend to [max(0, pos - sliding_window_size + 1), pos] + start_pos = max(0, pos - sliding_window_size + 1) + k_i = k_cache[cache_idx, start_pos : pos + 1] # [window_len, n_kv_heads, head_dim] + v_i = v_cache[cache_idx, start_pos : pos + 1] # [window_len, n_kv_heads, v_head_dim] + else: + # No sliding window: attend to all previous tokens [0, pos] + k_i = k_cache[cache_idx, : pos + 1] # [seq_len, n_kv_heads, head_dim] + v_i = v_cache[cache_idx, : pos + 1] # [seq_len, n_kv_heads, v_head_dim] + + # Transpose for attention: [n_heads, 1, head_dim] and [n_kv_heads, seq_len, head_dim] + q_i = q_i.unsqueeze(1) # [n_heads, 1, head_dim] + k_i = k_i.transpose(0, 1) # [n_kv_heads, seq_len, head_dim] + v_i = v_i.transpose(0, 1) # [n_kv_heads, seq_len, v_head_dim] + + # Handle GQA using existing repeat_kv function if needed + if n_heads != n_kv_heads: + n_rep = n_heads // n_kv_heads + # Reshape to [batch, num_kv_heads, seq_len, head_dim] for repeat_kv + # k_i is currently [n_kv_heads, seq_len, head_dim] + k_i_batch = k_i.unsqueeze(0) # [1, n_kv_heads, seq_len, head_dim] + v_i_batch = v_i.unsqueeze(0) # [1, n_kv_heads, seq_len, v_head_dim] + k_i_expanded = repeat_kv(k_i_batch, n_rep) # [1, n_heads, seq_len, head_dim] + v_i_expanded = repeat_kv(v_i_batch, n_rep) # [1, n_heads, seq_len, v_head_dim] + k_i = k_i_expanded[0] # [n_heads, seq_len, head_dim] + v_i = v_i_expanded[0] # [n_heads, seq_len, v_head_dim] + + # Compute attention scores + attn_scores = torch.matmul(q_i, k_i.transpose(-2, -1)) * scale # [n_heads, 1, seq_len] + + # Apply logit softcapping if enabled + attn_scores = _apply_logit_softcapping(attn_scores, logit_cap) + + # Apply sinks if provided (following the model file pattern) + if sinks is not None: + # Concatenate sinks to attention scores + sinks = sinks.reshape(-1, 1, 1).expand(-1, attn_scores.shape[-2], -1) + attn_weights = torch.cat([attn_scores, sinks], dim=-1) + attn_weights = torch.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype) + # Use only the non-sink portion for computing output (ignore sinks) + attn_out = torch.matmul( + attn_weights[..., : -sinks.size(-1)], v_i + ) # [n_heads, 1, v_head_dim] + else: + attn_weights = torch.softmax(attn_scores, dim=-1, dtype=torch.float32).to(q.dtype) + attn_out = torch.matmul(attn_weights, v_i) # [n_heads, 1, v_head_dim] + + # Store result: remove sequence dimension + out[i] = attn_out.squeeze(1) # [n_heads, v_head_dim] + + +def _torch_context_mha( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + input_pos: torch.Tensor, + cache_loc: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + seq_len: torch.Tensor, + seq_start: torch.Tensor, + scale: float, + out: torch.Tensor, + logit_cap: Optional[float] = None, + sliding_window_size: Optional[int] = None, + sinks: Optional[torch.Tensor] = None, +) -> None: + """Context attention (multiple tokens, potentially multiple sequences) using existing torch functions.""" + # Update KV cache first using existing function + update_kv_cache(k, v, k_cache, v_cache, seq_len, input_pos, cache_loc, seq_start) + + # Compute attention for each sequence + attn_outputs = [] + for idx in range(seq_len.shape[0]): + seq_len_i = seq_len[idx].item() + input_pos_i = input_pos[idx].item() + cache_loc_i = cache_loc[idx].item() + seq_start_i = seq_start[idx].item() + + # Skip sequences with zero length + if seq_len_i == 0: + continue + + # Get query for this sequence + q_seq = q[seq_start_i : seq_start_i + seq_len_i] # [seq_len_i, n_heads, head_dim] + + # Get keys and values from cache + kv_seq_len = input_pos_i + seq_len_i + k_seq = k_cache[cache_loc_i, :kv_seq_len] # [kv_seq_len, n_kv_heads, head_dim] + v_seq = v_cache[cache_loc_i, :kv_seq_len] # [kv_seq_len, n_kv_heads, head_dim] + + # Manual attention computation (shared path for both softcapping and non-softcapping) + n_heads = q_seq.shape[1] + n_kv_heads = k_seq.shape[1] + + # Transpose to [batch, num_heads, seq_len, head_dim] format + q_seq_t = q_seq.transpose(0, 1).unsqueeze(0) # [1, n_heads, seq_len_i, head_dim] + k_seq_t = k_seq.transpose(0, 1).unsqueeze(0) # [1, n_kv_heads, kv_seq_len, head_dim] + v_seq_t = v_seq.transpose(0, 1).unsqueeze(0) # [1, n_kv_heads, kv_seq_len, head_dim] + + # Handle GQA by repeating KV if needed + if n_heads != n_kv_heads: + n_rep = n_heads // n_kv_heads + k_seq_t = repeat_kv(k_seq_t, n_rep) # [1, n_heads, kv_seq_len, head_dim] + v_seq_t = repeat_kv(v_seq_t, n_rep) # [1, n_heads, kv_seq_len, head_dim] + + # Compute attention scores: Q @ K^T + attn_scores = ( + torch.matmul(q_seq_t, k_seq_t.transpose(-2, -1)) * scale + ) # [1, n_heads, seq_len_i, kv_seq_len] + + # Apply causal mask + causal_mask = torch.triu( + torch.ones(seq_len_i, kv_seq_len, device=q.device, dtype=torch.bool), + diagonal=kv_seq_len - seq_len_i + 1, + ) + + # Apply sliding window mask if specified + if sliding_window_size is not None and sliding_window_size > 0: + # Create sliding window mask: each query position i can only attend to keys in [i-window_size+1, i] + # For context phase, we need to account for the offset between query and key positions + + # Query positions are [input_pos_i, input_pos_i + seq_len_i) + # Key positions are [0, input_pos_i + seq_len_i) + query_positions = torch.arange( + input_pos_i, input_pos_i + seq_len_i, device=q.device + ) # [seq_len_i] + key_positions = torch.arange(0, kv_seq_len, device=q.device) # [kv_seq_len] + + # Create position difference matrix: query_pos - key_pos + pos_diff = query_positions.unsqueeze(1) - key_positions.unsqueeze( + 0 + ) # [seq_len_i, kv_seq_len] + + # Sliding window mask: allow attention only if 0 <= pos_diff < sliding_window_size + sliding_window_mask = (pos_diff < 0) | ( + pos_diff >= sliding_window_size + ) # [seq_len_i, kv_seq_len] + + # Combine causal and sliding window masks + combined_mask = causal_mask | sliding_window_mask + else: + combined_mask = causal_mask + + attn_scores.masked_fill_(combined_mask.unsqueeze(0).unsqueeze(0), float("-inf")) + + # Apply logit softcapping if enabled + attn_scores = _apply_logit_softcapping(attn_scores, logit_cap) + + # Apply sinks if provided (following the model file pattern) + if sinks is not None: + # Concatenate sinks to attention scores + sinks = sinks.reshape(1, -1, 1, 1).expand( + attn_scores.shape[0], -1, attn_scores.shape[-2], -1 + ) + attn_weights = torch.cat([attn_scores, sinks], dim=-1) + attn_weights = torch.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype) + # Use only the non-sink portion for computing output (ignore sinks) + attn_out = torch.matmul( + attn_weights[..., : -sinks.size(-1)], v_seq_t + ) # [1, n_heads, seq_len_i, v_head_dim] + else: + attn_weights = torch.softmax(attn_scores, dim=-1, dtype=torch.float32).to(q.dtype) + attn_out = torch.matmul(attn_weights, v_seq_t) # [1, n_heads, seq_len_i, v_head_dim] + + # Remove batch dimension and transpose back to [seq_len_i, n_heads, v_head_dim] + attn_out = attn_out[0].transpose(0, 1) + + attn_outputs.append(attn_out) + + # Concatenate all outputs + if len(attn_outputs) == 0: + # No sequences to process - this shouldn't happen but handle gracefully + out.zero_() + elif len(attn_outputs) == 1: + # Single sequence + out.copy_(attn_outputs[0]) + else: + # Multiple sequences or context phase + out.copy_(torch.cat(attn_outputs, dim=0)) + + +@torch.library.custom_op("auto_deploy::torch_cached_attention_with_cache", mutates_args=()) +def torch_backend_mha_with_cache( + # Q, K, V + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + # METADATA + seq_len: torch.Tensor, + input_pos: torch.Tensor, + cache_loc: torch.Tensor, + seq_start: torch.Tensor, + # CACHES + k_cache: torch.Tensor, + v_cache: torch.Tensor, + # BUFFERS + # + # CONSTANTS + scale: Optional[float], + sinks: Optional[torch.Tensor] = None, + sliding_window_size: Optional[int] = None, + logit_cap: Optional[float] = None, +) -> torch.Tensor: + """Torch backend MHA with cache that takes q, k, v in BSND layout.""" + # Get dimensions + num_kv_heads, qk_head_dim = k_cache.shape[-2:] + v_head_dim = v_cache.shape[-1] + b, s = q.shape[:2] + + # check for num_heads + num_heads = q.shape[2] // qk_head_dim if q.ndim == 3 else q.shape[2] + + # Define output shape + output_shape = (b, s, num_heads * v_head_dim) if q.ndim == 3 else (b, s, num_heads, v_head_dim) + + # Reshape to standard layout + if s == 1: + bs_view = (b, s) + else: + bs_view = (b * s,) + + q = q.contiguous().view(*bs_view, num_heads, qk_head_dim) + k = k.contiguous().view(*bs_view, num_kv_heads, qk_head_dim) + v = v.contiguous().view(*bs_view, num_kv_heads, v_head_dim) + + scale = 1.0 / math.sqrt(qk_head_dim) if scale is None else scale + + # Create output tensor + y = q.new_empty(*bs_view, num_heads, v_head_dim).contiguous() + + # Compute attention + if s == 1: + # Generate-only phase + _torch_generate_mha( + q, + k, + v, + k_cache, + v_cache, + cache_loc, + input_pos, + scale, + y, + logit_cap, + sliding_window_size, + sinks, + ) + else: + # Context phase + _torch_context_mha( + q, + k, + v, + input_pos, + cache_loc, + k_cache, + v_cache, + seq_len, + seq_start, + scale, + y, + logit_cap, + sliding_window_size, + sinks, + ) + + return y.view(*output_shape) + + +@torch_backend_mha_with_cache.register_fake +def torch_backend_mha_with_cache_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + seq_len: torch.Tensor, + input_pos: torch.Tensor, + cache_loc: torch.Tensor, + seq_start: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + scale: Optional[float], + sinks: Optional[torch.Tensor] = None, + sliding_window_size: Optional[int] = None, + logit_cap: Optional[float] = None, +): + return q.new_empty(*q.shape[:-1], v.shape[-1]).contiguous() + + +@torch.library.custom_op("auto_deploy::torch_cached_attention_prepare_metadata", mutates_args=()) +def torch_backend_prepare_metadata( + input_ids: torch.Tensor, + position_ids: torch.Tensor, + seq_len: torch.Tensor, + input_pos: torch.Tensor, + cache_loc: torch.Tensor, + pages_per_seq: torch.Tensor, + page_size: int, +) -> List[torch.Tensor]: + """Prepare metadata for torch backend attention (similar to triton backend).""" + num_seq = SequenceInfo._get_sanitized_num_sequences(input_ids, seq_len) + seq_start = torch.zeros_like(seq_len[:num_seq]) + seq_start[1:] = torch.cumsum(seq_len[: num_seq - 1], 0) + return ( + seq_len[:num_seq].clone(), + input_pos[:num_seq].clone(), + cache_loc[:num_seq].clone(), + seq_start, + ) + + +@torch_backend_prepare_metadata.register_fake +def torch_backend_prepare_metadata_fake( + input_ids, position_ids, seq_len, input_pos, cache_loc, pages_per_seq, page_size +): + num_seq = SequenceInfo._get_sanitized_num_sequences(input_ids, seq_len) + return ( + torch.empty_like(seq_len[:num_seq]), + torch.empty_like(input_pos[:num_seq]), + torch.empty_like(cache_loc[:num_seq]), + torch.empty_like(seq_len[:num_seq]), + ) + + +@AttentionRegistry.register("torch") +class TorchBackendAttention(AttentionDescriptor): + @classmethod + def is_paged(cls) -> bool: + """Return if the attention op is paged or not.""" + return False + + @classmethod + def get_attention_layout(cls) -> AttentionLayout: + """Get the attention layout expected by the source op and the cached attention op.""" + return "bsnd" + + @classmethod + def get_num_qkv_args(cls) -> int: + """Get the number of qkv arguments expected by the source op.""" + return 3 + + @classmethod + def get_source_attention_op(cls) -> OpOverloadPacket: + return torch.ops.auto_deploy.torch_attention_bsnd_grouped_sdpa + + @classmethod + def get_cached_attention_op(cls) -> MHACallable: + return torch.ops.auto_deploy.torch_cached_attention_with_cache + + @classmethod + def get_prepare_metadata_op(cls) -> Tuple[PrepareMetadataCallable, int]: + return torch.ops.auto_deploy.torch_cached_attention_prepare_metadata, 4 + + @classmethod + def get_cache_initializers( + cls, source_attn_node: Node, cache_config: CacheConfig + ) -> CacheInitializerDict: + # source op is [bsnd] layout already + k_fake: FakeTensor = source_attn_node.args[1].meta["val"] + v_fake: FakeTensor = source_attn_node.args[2].meta["val"] + num_kv_heads = k_fake.shape[2] + k_head_dim = k_fake.shape[3] + v_head_dim = v_fake.shape[3] + + def _get_k_cache(si: SequenceInfo): + assert not si.is_paged, "Paged cache not supported for torch backend" + return torch.empty( + si.num_pages, + si.page_size, + num_kv_heads, + k_head_dim, + device=si.device, + dtype=cache_config.dtype or k_fake.dtype, + ) + + def _get_v_cache(si: SequenceInfo): + assert not si.is_paged, "Paged cache not supported for torch backend" + return torch.empty( + si.num_pages, + si.page_size, + num_kv_heads, + v_head_dim, + device=si.device, + dtype=cache_config.dtype or v_fake.dtype, + ) + + return {"k_cache": _get_k_cache, "v_cache": _get_v_cache} + + @classmethod + def get_global_buffer_initializers(cls, source_attn_node: Node) -> BufferInitializerDict: + return {} + + @classmethod + def get_constants(cls, source_attn_node: Node) -> List[Constant]: + # Check other arguments + attn_mask, dropout_p, is_causal = extract_op_args( + source_attn_node, "attn_mask", "dropout_p", "is_causal" + ) + if attn_mask is not None or dropout_p != 0.0 or not is_causal: + ad_logger.debug( + "Unsupported attention arguments for " + f"{source_attn_node=}: {attn_mask=}, {dropout_p=}, {is_causal=}" + ) + + # Get scale from args or kwargs + if len(source_attn_node.args) > 6: + scale = source_attn_node.args[6] + else: + scale = source_attn_node.kwargs.get("scale", None) + + # Validate scale + if not isinstance(scale, float): + ad_logger.warning("Provided scale is not a float. Using default scale instead.") + scale = None + + # Get sinks, sliding_window, and logit_cap from args or kwargs + sinks = extract_op_args(source_attn_node, "sinks")[0] + sliding_window = extract_op_args(source_attn_node, "sliding_window")[0] + logit_cap = extract_op_args(source_attn_node, "logit_cap")[0] + + return [ + scale, # softmax scale + sinks, # sinks parameter + sliding_window, # sliding window parameter + logit_cap, # logit cap parameter + ] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_moe.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_moe.py index f5e7373c47a3..5b7131f12963 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_moe.py @@ -1,9 +1,45 @@ -from typing import List +from typing import Callable, List import torch import torch.nn.functional as F +def _template_moe( + x: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + mlps: List[Callable[[torch.Tensor], torch.Tensor]], +) -> torch.Tensor: + """Mixtral-style generic MoE template, dispatching tokens to expert MLPs based on routing info.""" + x_shape = x.shape + hidden_dim = x_shape[-1] + x = x.view(-1, hidden_dim) + num_experts = len(mlps) + + final_hidden_states = torch.zeros_like(x) + valid_mask = (selected_experts >= 0) & (selected_experts < num_experts) + # For out-of-range indices, set them to num_experts + selected_experts_fixed = torch.where( + valid_mask, selected_experts, torch.full_like(selected_experts, num_experts) + ) + # Create one-hot encoding with an extra class. + one_hot = F.one_hot(selected_experts_fixed, num_classes=num_experts + 1) + expert_mask = one_hot[..., :num_experts].permute(2, 1, 0) + + for expert_idx in range(num_experts): + idx, top_x = torch.where(expert_mask[expert_idx]) + tokens_for_this_expert = x[None, top_x].reshape(-1, hidden_dim) + if not tokens_for_this_expert.shape[0]: + continue # input of shape [0, hidden_dim] breaks fp4 kernel + + expert_out = mlps[expert_idx](tokens_for_this_expert) + current_hidden_states = expert_out * routing_weights[top_x, idx, None] + final_hidden_states.index_add_( + 0, top_x, current_hidden_states.to(final_hidden_states.dtype) + ) + return final_hidden_states.view(x_shape) + + @torch.library.custom_op("auto_deploy::torch_moe", mutates_args=()) def torch_moe( x: torch.Tensor, @@ -33,41 +69,17 @@ def torch_moe( torch.Tensor: Output tensor with the same shape as the input x. """ - x_shape = x.shape - hidden_dim = x_shape[-1] - x = x.view(-1, hidden_dim) - num_experts = len(w1_weight) - - final_hidden_states = torch.zeros_like(x) - valid_mask = (selected_experts >= 0) & (selected_experts < num_experts) - # For out-of-range indices, set them to num_experts - selected_experts_fixed = torch.where( - valid_mask, selected_experts, torch.full_like(selected_experts, num_experts) - ) - # Create one-hot encoding with an extra class. - one_hot = torch.nn.functional.one_hot(selected_experts_fixed, num_classes=num_experts + 1) - expert_mask = one_hot[..., :num_experts].permute(2, 1, 0) - - for expert_idx in range(num_experts): - idx, top_x = torch.where(expert_mask[expert_idx]) - tokens_for_this_expert = x[None, top_x].reshape(-1, hidden_dim) - - gate_out = F.linear(tokens_for_this_expert, w1_weight[expert_idx]) - up_out = F.linear(tokens_for_this_expert, w3_weight[expert_idx]) - activated = F.silu(gate_out) - prod = activated * up_out - expert_out = F.linear(prod, w2_weight[expert_idx]) - - current_hidden_states = expert_out * routing_weights[top_x, idx, None] - final_hidden_states.index_add_( - 0, top_x, current_hidden_states.to(final_hidden_states.dtype) + def make_mlp(i): + return lambda inp: F.linear( + F.silu(F.linear(inp, w1_weight[i])) * F.linear(inp, w3_weight[i]), w2_weight[i] ) - return final_hidden_states.view(x_shape) + mlps = [make_mlp(i) for i in range(len(w1_weight))] + return _template_moe(x, selected_experts, routing_weights, mlps) @torch_moe.register_fake -def torch_moe( +def torch_moe_fake( x: torch.Tensor, selected_experts: torch.Tensor, routing_weights: torch.Tensor, @@ -133,7 +145,7 @@ def torch_fused_moe( @torch_fused_moe.register_fake -def torch_fused_moe( +def torch_fused_moe_fake( x: torch.Tensor, selected_experts: torch.Tensor, routing_weights: torch.Tensor, @@ -141,3 +153,174 @@ def torch_fused_moe( w2_stacked_weight: torch.Tensor, ) -> torch.Tensor: return torch.empty_like(x) + + +@torch.library.custom_op("auto_deploy::torch_quant_fp8_moe", mutates_args=()) +def torch_quant_fp8_moe( + x: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + w1_weight: List[torch.Tensor], + w2_weight: List[torch.Tensor], + w3_weight: List[torch.Tensor], + w1_input_scale: List[torch.Tensor], + w2_input_scale: List[torch.Tensor], + w3_input_scale: List[torch.Tensor], + w1_weight_scale: List[torch.Tensor], + w2_weight_scale: List[torch.Tensor], + w3_weight_scale: List[torch.Tensor], +) -> torch.Tensor: + """ + FP8 MoE op using quantized linear operations. + + Computes a Mixture-of-Experts layer similar to the reference auto_deploy::torch_moe op, but uses the + quantized FP8 linear op for expert computations. + + Args: + x: Input tensor of shape (B, H) or (B, S, H). + selected_experts: Tensor (B, TOP_K) or (B*S, TOP_K) containing expert indices. + routing_weights: Tensor of normalized routing weights. + w1_weight, w2_weight, w3_weight: Lists of pre-quantized weight tensors for the three linear ops. + w1_input_scale, w2_input_scale, w3_input_scale: Lists of input scale tensors for the corresponding ops. + w1_weight_scale, w2_weight_scale, w3_weight_scale: Lists of weight scale tensors for the corresponding ops. + + """ + + def make_fp8_mlp(i): + def mlp(inp): + gate_out = torch.ops.auto_deploy.torch_quant_fp8_linear( + inp, + w1_weight[i], + bias=None, + input_scale=w1_input_scale[i], + weight_scale=w1_weight_scale[i], + ) + up_out = torch.ops.auto_deploy.torch_quant_fp8_linear( + inp, + w3_weight[i], + bias=None, + input_scale=w3_input_scale[i], + weight_scale=w3_weight_scale[i], + ) + prod = F.silu(gate_out) * up_out + return torch.ops.auto_deploy.torch_quant_fp8_linear( + prod, + w2_weight[i], + bias=None, + input_scale=w2_input_scale[i], + weight_scale=w2_weight_scale[i], + ) + + return mlp + + mlps = [make_fp8_mlp(i) for i in range(len(w1_weight))] + return _template_moe(x, selected_experts, routing_weights, mlps) + + +@torch_quant_fp8_moe.register_fake +def torch_quant_fp8_moe_fake( + x: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + w1_weight: List[torch.Tensor], + w2_weight: List[torch.Tensor], + w3_weight: List[torch.Tensor], + w1_input_scale: List[torch.Tensor], + w2_input_scale: List[torch.Tensor], + w3_input_scale: List[torch.Tensor], + w1_weight_scale: List[torch.Tensor], + w2_weight_scale: List[torch.Tensor], + w3_weight_scale: List[torch.Tensor], +) -> torch.Tensor: + return torch.empty_like(x) + + +@torch.library.custom_op("auto_deploy::torch_quant_fp4_moe", mutates_args=()) +def torch_quant_fp4_moe( + x: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + w1_weight: List[torch.Tensor], + w2_weight: List[torch.Tensor], + w3_weight: List[torch.Tensor], + w1_input_scale: List[torch.Tensor], + w2_input_scale: List[torch.Tensor], + w3_input_scale: List[torch.Tensor], + w1_weight_scale: List[torch.Tensor], + w2_weight_scale: List[torch.Tensor], + w3_weight_scale: List[torch.Tensor], + w1_alpha: List[torch.Tensor], + w2_alpha: List[torch.Tensor], + w3_alpha: List[torch.Tensor], +) -> torch.Tensor: + """ + FP4 MoE op using quantized linear operations. + + Computes a Mixture-of-Experts layer similar to the reference auto_deploy::torch_moe op, + but uses the NVFP4 quantized linear op for expert computations. + + Args: + x: Input tensor of shape (B, H) or (B, S, H). + selected_experts: Tensor (B, TOP_K) or (B*S, TOP_K) containing expert indices. + routing_weights: Tensor of normalized routing weights. + w1_weight, w2_weight, w3_weight: Lists of pre-quantized weight tensors for the three linear ops. + w1_input_scale, w2_input_scale, w3_input_scale: Lists of input scale tensors. + w1_weight_scale, w2_weight_scale, w3_weight_scale: Lists of weight scale tensors. + w1_alpha, w2_alpha, w3_alpha: Lists of alpha scale tensors for FP4 quantization. + """ + + def make_fp4_mlp(i): + def mlp(inp): + if inp.shape[0] == 0: + return torch.zeros_like(inp) + gate_out = torch.ops.auto_deploy.torch_quant_fp4_linear( + inp, + w1_weight[i], + bias=None, + input_scale=w1_input_scale[i], + weight_scale=w1_weight_scale[i], + alpha=w1_alpha[i], + ) + up_out = torch.ops.auto_deploy.torch_quant_fp4_linear( + inp, + w3_weight[i], + bias=None, + input_scale=w3_input_scale[i], + weight_scale=w3_weight_scale[i], + alpha=w3_alpha[i], + ) + prod = F.silu(gate_out) * up_out + return torch.ops.auto_deploy.torch_quant_fp4_linear( + prod, + w2_weight[i], + bias=None, + input_scale=w2_input_scale[i], + weight_scale=w2_weight_scale[i], + alpha=w2_alpha[i], + ) + + return mlp + + mlps = [make_fp4_mlp(i) for i in range(len(w1_weight))] + return _template_moe(x, selected_experts, routing_weights, mlps) + + +@torch_quant_fp4_moe.register_fake +def torch_quant_fp4_moe_fake( + x: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + w1_weight: List[torch.Tensor], + w2_weight: List[torch.Tensor], + w3_weight: List[torch.Tensor], + w1_input_scale: List[torch.Tensor], + w2_input_scale: List[torch.Tensor], + w3_input_scale: List[torch.Tensor], + w1_weight_scale: List[torch.Tensor], + w2_weight_scale: List[torch.Tensor], + w3_weight_scale: List[torch.Tensor], + w1_alpha: List[torch.Tensor], + w2_alpha: List[torch.Tensor], + w3_alpha: List[torch.Tensor], +) -> torch.Tensor: + return torch.empty_like(x) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.py index b5c7780be121..e6bac2aeb812 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.py @@ -41,6 +41,8 @@ def _generate_mha( input_pos: torch.Tensor, scale: float, out: torch.Tensor, + sinks: Optional[torch.Tensor] = None, + sliding_window: Optional[int] = None, ): b, (n_heads, q_d_head) = q.shape[0], q.shape[-2:] max_seq_len, n_kv_heads = k_cache.shape[1:3] @@ -97,7 +99,10 @@ def _generate_mha( v_d_head, SEQ_BLOCK_SIZE, HEAD_BLOCK_SIZE, + sliding_window if sliding_window is not None else -1, ) + has_sinks = sinks is not None + attention_kv_stage2[(b, n_heads, 1)]( stage1_output_values, stage1_output_logsumexp, @@ -107,6 +112,8 @@ def _generate_mha( n_heads, v_d_head, SEQ_BLOCK_SIZE, + has_sinks, + sinks, ) @@ -122,6 +129,8 @@ def _flattened_context_mha( seq_start: torch.Tensor, scale: float, out: torch.Tensor, + sinks: Optional[torch.Tensor] = None, + sliding_window: Optional[int] = None, ) -> None: # NOTE: s_total == sum(seq_len) s_total, n_heads, q_d_head = q.shape @@ -149,6 +158,8 @@ def _flattened_context_mha( # TODO: use input_pos to get the correct cache locations grid = (BATCH_SIZE, n_heads, (max(seq_len) + SEQ_BLOCK - 1) // SEQ_BLOCK) + has_sinks = sinks is not None + context_attention_kv_flattened[grid]( q, seq_len, @@ -165,7 +176,9 @@ def _flattened_context_mha( v_d_head, SEQ_BLOCK, max_cache_seq_len, - num_stages=2, + sliding_window if sliding_window is not None else -1, + has_sinks, + sinks, ) @@ -187,6 +200,8 @@ def flattened_mha_with_cache( # # CONSTANTS scale: Optional[float], + sinks: Optional[torch.Tensor] = None, + sliding_window: Optional[int] = None, ) -> torch.Tensor: """Flattened MHA with cache that takes q, k, v in BSND layout. @@ -223,7 +238,9 @@ def flattened_mha_with_cache( y = q.new_empty(*bs_view, num_heads, v_head_dim).contiguous() if s == 1: # generate-only phase - _generate_mha(q, k, v, k_cache, v_cache, cache_loc, input_pos, scale, y) + _generate_mha( + q, k, v, k_cache, v_cache, cache_loc, input_pos, scale, y, sinks, sliding_window + ) else: # mixed context + generate phase _flattened_context_mha( @@ -238,6 +255,8 @@ def flattened_mha_with_cache( seq_start, scale, y, + sinks, + sliding_window, ) return y.view(*output_shape) @@ -255,6 +274,8 @@ def flattened_mha_fake( k_cache: torch.Tensor, v_cache: torch.Tensor, scale: Optional[float], + sinks: Optional[torch.Tensor] = None, + sliding_window: Optional[int] = None, ): return q.new_empty(*q.shape[:-1], v.shape[-1]).contiguous() @@ -388,7 +409,11 @@ def get_constants(cls, source_attn_node: Node) -> List[Constant]: if not isinstance(scale, float): ad_logger.warning("Provided scale is not a float, Using default scale instead.") scale = None - + # Get sinks and sliding_window from args or kwargs + sinks = extract_op_args(source_attn_node, "sinks")[0] + sliding_window = extract_op_args(source_attn_node, "sliding_window")[0] return [ scale, # softmax scale + sinks, + sliding_window, ] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/triton_kernels/attention_with_kv_cache.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/triton_kernels/attention_with_kv_cache.py index 9a59a363dc44..ac1c43f0c913 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/triton_kernels/attention_with_kv_cache.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/triton_kernels/attention_with_kv_cache.py @@ -112,6 +112,7 @@ def gqa_attention_kv_stage1( V_D_HEAD: tl.constexpr, # Dimension of each key/value head SEQ_BLOCK_SIZE: tl.constexpr, # Block size used for tiling the sequence dim. HEAD_BLOCK_SIZE: tl.constexpr, # pad to 16 if HEAD_RATIO is < 16 to invoke tensor cores. + SLIDING_WINDOW: tl.constexpr, ): """Attention kernel to be used for generate-only batches. @@ -122,7 +123,7 @@ def gqa_attention_kv_stage1( Supports non-power-of-2 D_HEAD Uses flash decoding. - KV-cache layout is assumed to be [Batch,Seq, Head, Dim] + KV-cache layout is assumed to be [Batch, Seq, Head, Dim] 1. Fetch the K-cache from 0 to input_pos 2. Fetch the V-cache from 0 to input_pos 3. A = Q*K^T [1,D_HEAD] * [1,seq_len,D_HEAD] -> [1, seq_len] @@ -145,10 +146,20 @@ def gqa_attention_kv_stage1( # The number of Q heads that map to each KV head. HEAD_RATIO: tl.constexpr = N_HEADS // N_KV_HEADS # This needs to be a power-of-2 - if seq_start_pos > kv_position: - return - seq_offsets = seq_start_pos + tl.arange(0, SEQ_BLOCK_SIZE) - seq_mask = seq_offsets <= kv_position + + # Apply sliding window constraints + if SLIDING_WINDOW > 0: + # For sliding window, limit the sequence range + sliding_start = tl.maximum(0, kv_position - SLIDING_WINDOW + 1) + if seq_start_pos + SEQ_BLOCK_SIZE <= sliding_start or seq_start_pos > kv_position: + return + seq_offsets = seq_start_pos + tl.arange(0, SEQ_BLOCK_SIZE) + seq_mask = (seq_offsets <= kv_position) & (seq_offsets >= sliding_start) + else: + if seq_start_pos > kv_position: + return + seq_offsets = seq_start_pos + tl.arange(0, SEQ_BLOCK_SIZE) + seq_mask = seq_offsets <= kv_position # Need to pad the head dim to 16 if HEAD_RATIO is < 16 so that tensor cores can be invoked # @@ -358,6 +369,8 @@ def attention_kv_stage2( N_HEADS: tl.constexpr, D_HEAD: tl.constexpr, SEQ_BLOCK_SIZE: tl.constexpr, # Nearest power of 2 for num_blocks + HAS_SINKS: tl.constexpr, + sinks_ptr, ): # There are batch * N_HEADS programs batch_id = tl.program_id(axis=0) @@ -382,6 +395,11 @@ def attention_kv_stage2( sumexp = tl.exp(logsumexp - max_logsumexp) # [NUM_BLOCKS_POW2] aggregate_sumexp = tl.sum(sumexp, axis=0) + # Add sinks contribution to the softmax denominator + if HAS_SINKS: + sinks_val = tl.load(sinks_ptr + batch_id * N_HEADS + head_id) + sinks_exp = tl.exp(sinks_val - max_logsumexp) + aggregate_sumexp += sinks_exp values_offsets = block_offsets[:, None] * D_HEAD + dhead_offsets[None, :] values_mask = block_mask[:, None] * dhead_mask[None, :] @@ -573,6 +591,9 @@ def context_attention_kv_flattened( V_D_HEAD: tl.constexpr, # Dimension of each value head. SEQ_BLOCK: tl.constexpr, MAX_SEQ_LENGTH: tl.constexpr, + SLIDING_WINDOW: tl.constexpr, # Sliding window size, -1 means no sliding window + HAS_SINKS: tl.constexpr, + sinks_ptr, ): """Kernel for context phase. @@ -623,7 +644,15 @@ def context_attention_kv_flattened( # input_pos_ptr stores the location at which kv must be written back for the given batch. kv_position = tl.load(input_pos_ptr + batch_id) num_blocks = (kv_position + seq_len + SEQ_BLOCK - 1) // SEQ_BLOCK - for s in range(0, num_blocks + 1, 1): + start = 0 + if SLIDING_WINDOW > 0: + # Use the LAST query in this block for more conservative start calculation + last_q_pos = ( + (seq_block_id + 1) * SEQ_BLOCK - 1 + kv_position + ) # Last query's absolute position + earliest_kv_pos = max(0, last_q_pos - SLIDING_WINDOW + 1) + start = max(0, earliest_kv_pos // SEQ_BLOCK) + for s in range(start, num_blocks + 1): kv_seq_offsets = s * SEQ_BLOCK + tl.arange(0, SEQ_BLOCK) kv_seq_mask = kv_seq_offsets < (kv_position + seq_len) @@ -637,9 +666,17 @@ def context_attention_kv_flattened( ) qk = tl.zeros([SEQ_BLOCK, SEQ_BLOCK], dtype=tl.float32) qk += tl.dot(q, k.trans()) - qk = tl.where( - (seq_offsets[:, None] + kv_position) >= kv_seq_offsets[None, :], qk, float("-inf") - ) + # Apply causal mask + causal_mask = (seq_offsets[:, None] + kv_position) >= kv_seq_offsets[None, :] + # Apply sliding window mask if enabled + if SLIDING_WINDOW > 0: + sliding_window_mask = kv_seq_offsets[None, :] >= ( + seq_offsets[:, None] + kv_position - SLIDING_WINDOW + 1 + ) + combined_mask = sliding_window_mask & causal_mask + else: + combined_mask = causal_mask + qk = tl.where(combined_mask, qk, float("-inf")) qk *= SCALE # rowmax m_ij = tl.maximum(tl.max(qk, 1), lse_i) @@ -662,6 +699,16 @@ def context_attention_kv_flattened( l_i_new = tl.exp(lse_i - m_ij) + l_ij lse_i = m_ij + tl.log(l_i_new) + # Add sinks contribution to the final softmax calculation + if HAS_SINKS: + sinks_val = tl.load(sinks_ptr + batch_id * N_HEADS + head_id) + m_sinks = tl.maximum(m_i, sinks_val) + acc_scale = tl.exp(m_i - m_sinks) + acc = acc * acc_scale[:, None] + l_sinks = tl.exp(lse_i - m_sinks) + tl.exp(sinks_val - m_sinks) + lse_i = m_sinks + tl.log(l_sinks) + m_i = m_sinks + o_scale = tl.exp(m_i - lse_i) acc = acc * o_scale[:, None] diff --git a/tensorrt_llm/_torch/auto_deploy/export/__init__.py b/tensorrt_llm/_torch/auto_deploy/export/__init__.py new file mode 100644 index 000000000000..f655c5043cc9 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/export/__init__.py @@ -0,0 +1,5 @@ +"""AutoDeploy's modular export patch system.""" + +from . import library # ensure all patches are registered +from .export import * +from .interface import * diff --git a/tensorrt_llm/_torch/auto_deploy/export/export.py b/tensorrt_llm/_torch/auto_deploy/export/export.py new file mode 100644 index 000000000000..475017a28401 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/export/export.py @@ -0,0 +1,284 @@ +"""Main export functionality with utilities for torch.export.""" + +from collections import defaultdict +from contextlib import nullcontext +from functools import partial +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.export as te +import torch.nn as nn +from torch import fx + +from ..transformations._graph import ( + canonicalize_graph, + lift_to_meta, + load_buffers_and_params, + tree_to, +) +from ..utils.logger import ad_logger +from ..utils.node_utils import is_op +from .interface import ExportPatchRegistry, apply_export_patches + +try: + from modelopt.torch.quantization.utils import export_torch_mode as torch_export_context +except ImportError: + torch_export_context = nullcontext + + +def _clean_up_device_info(gm: fx.GraphModule) -> None: + """Correct device information in the graph.""" + devices = {t.device for _, t in gm.named_parameters()} + if len(devices) == 0: + return + elif len(devices) > 1: + raise AssertionError("All parameters should be on the same device.") + device = devices.pop() + meta_device = torch.device("meta") + + for node in gm.graph.nodes: + if any(a == meta_device for a in node.args): + new_args = list(node.args) + new_args = [a if a != meta_device else device for a in new_args] + node.args = tuple(new_args) + if any(a == meta_device for a in node.kwargs.values()): + new_kwargs = dict(node.kwargs) + new_kwargs = {k: v if v != meta_device else device for k, v in new_kwargs.items()} + node.kwargs = new_kwargs + + canonicalize_graph(gm) + + +def _load_hook_for_deduplication( + state_dict, prefix, *args, param_key_remaining: str, param_key_removed: str +): + """Check for removed param key and and put it into the key that is remaining.""" + ad_logger.debug(f"Loading hook for deduplication: {param_key_remaining} <- {param_key_removed}") + k_remaining = prefix + param_key_remaining + k_removed = prefix + param_key_removed + if k_removed in state_dict: + state_dict[k_remaining] = state_dict.pop(k_removed) + + +def _deduplicate_params_and_buffers(gm: fx.GraphModule) -> None: + """This will de-duplicate params and buffers that share the same tensor.""" + # get all get_attr nodes + get_attr_nodes = [n for n in gm.graph.nodes if n.op == "get_attr"] + + # sort by id of target + targets: Dict[int, List[fx.Node]] = defaultdict(list) + for n in get_attr_nodes: + submod, _, name = n.target.rpartition(".") + t_target = getattr(gm.get_submodule(submod), name) + targets[id(t_target)].append(n) + # now replace all instances of the same tensor with the same get_attr node (idx 0 in the list) + for nodes in targets.values(): + node_kept = nodes[0] + for n in nodes[1:]: + n.replace_all_uses_with(node_kept) + gm.graph.erase_node(n) + + # remove the param/buffer from the submodule + submod, _, name = n.target.rpartition(".") + delattr(gm.get_submodule(submod), name) + + # add load hooks to also load the weights correctly + gm._register_load_state_dict_pre_hook( + partial( + _load_hook_for_deduplication, + param_key_remaining=str(node_kept.target), + param_key_removed=str(n.target), + ) + ) + + ad_logger.debug(f"Deduplicated: {n.target} --> {node_kept.target}") + + canonicalize_graph(gm) + + +def _add_missing_load_hooks(gm: fx.GraphModule, model: nn.Module) -> None: + """Adds back the state dict load hooks stripped away during export.""" + hooks = { + k: mod._load_state_dict_pre_hooks + for k, mod in model.named_modules() + if mod._load_state_dict_pre_hooks + } + + for mod_name, mod in gm.named_modules(): + if mod_name in hooks: + for hook in hooks.pop(mod_name).values(): + mod._register_load_state_dict_pre_hook(hook.hook, with_module=hook.with_module) + assert not (bool(hooks)), f"""Mismatch in names of exported and source modules with hooks. + The following module names were not found in exported module {list(hooks.keys())}""" + + +def _add_load_hook_for_aliased_params(gm: fx.GraphModule, model: nn.Module) -> None: + """ + Add a load hook to handle aliased parameters in the model. + + When parameters are aliased (multiple parameter names point to the same tensor), + we need to ensure all aliases get the same value during loading. This hook: + 1. Identifies groups of aliased parameters + 2. For each group, finds a valid parameter value from the state dict + 3. Applies that value to all aliases in the group + + Args: + gm: The graph module to add the hook to + model: The source model containing the original parameter aliases + """ + + def find_valid_param_value( + state_dict: Dict[str, torch.Tensor], param_names: List[str] + ) -> Optional[torch.Tensor]: + """Find a valid parameter value from state dict for a group of aliased parameters. + + Args: + state_dict: The state dict being loaded + param_names: List of parameter names that are aliases of each other + + Returns: + A valid tensor value if found, None otherwise + """ + # First try to find a non-meta tensor value + value = None + for name in param_names: + if name in state_dict: + value = state_dict[name] + if value.device.type != "meta": + return value + + return value + + def aliasing_load_pre_hook(state_dict: Dict[str, torch.Tensor], prefix: str, *args, **kwargs): + """Load hook that ensures aliased parameters get the same value.""" + for group in aliased_groups: + # Find a valid value for this group of aliases + value = find_valid_param_value(state_dict, group) + + if value is not None: + # Apply the value to all aliases + for name in group: + state_dict[name] = value + + ad_logger.debug(f"Applied value from {group[0]} to aliased parameters: {group}") + + # Find all parameter aliases in the source model + param_to_names = defaultdict(list) + for name, param in model.named_parameters(remove_duplicate=False): + param_to_names[id(param)].append(name) + + # Filter to only groups with multiple aliases + aliased_groups = [names for names in param_to_names.values() if len(names) > 1] + + if not aliased_groups: + return + + # Register the hook + gm._register_load_state_dict_pre_hook(aliasing_load_pre_hook) + + +def _clean_up_assertions(gm: fx.GraphModule): + """This transformations removes shape checks and assertions from the graph.""" + check_ops = { + torch.ops.aten._assert_scalar, + torch.ops.aten.sym_constrain_range, + torch.ops.aten.sym_constrain_range_for_size, + torch.ops.aten._assert_tensor_metadata, + # torch.ops.aten._functional_sym_constrain_range, + # torch.ops.aten._functional_sym_constrain_range_for_size + } + graph: fx.Graph = gm.graph + for node in reversed(graph.nodes): + if len(node.users) > 0 or not is_op(node, check_ops): + continue + graph.erase_node(node) + canonicalize_graph(gm) + + +def torch_export_to_gm( + model: nn.Module, + args: Tuple[Any, ...], + kwargs: Optional[Dict[str, Any]] = None, + clone: bool = False, # clone or don't clone the model state_dict + *, + dynamic_shapes: Optional[Union[dict[str, Any], tuple[Any], list[Any]]] = None, + strict: bool = False, + patch_configs: Optional[Dict[str, Union[dict, Any]]] = None, + patch_list: Optional[List[str]] = None, +) -> fx.GraphModule: + """torch's export with wrapping into GraphModule + useful additions to the resulting module. + + This utility improves over stock torch.export.export in the following aspects: + + 1. Provide patches for certain corner cases that torch.export does not support. + 2. Standardize the export process to strictly run on the meta device. + 3. Automatically extract the GraphModule from the exported program. + 4. Retain load hooks for state_dict loading from the original module. + 5. Manage parameter aliasing in the model. + 6. Remove assertions from the graph. + + Args: + model: The model to export + args: Arguments for the model + kwargs: Keyword arguments for the model + clone: Whether to clone the model state_dict + dynamic_shapes: Dynamic shapes for the export + strict: Whether to use strict mode for export + patch_configs: Optional patch configurations. If None, all registered patches + will be applied with default settings. + patch_list: Optional list of patch names to apply with default settings. + Cannot be used together with patch_configs. + """ + # Validate that both patch_configs and patch_list are not provided simultaneously + if patch_configs is not None and patch_list is not None: + raise ValueError("Cannot specify both patch_configs and patch_list. Use only one.") + + # Handle patch configuration + if patch_list is not None: + # Convert patch_list to patch_configs format + patch_configs = {patch_name: {} for patch_name in patch_list} + elif patch_configs is None: + # Default patch configurations - apply all registered patches with default settings + patch_configs = {patch_name: {} for patch_name in ExportPatchRegistry.list_patches()} + + # run export with patches and lifted to meta + with apply_export_patches(patch_configs), lift_to_meta(model) as state_dict: + # clean up args, kwargs and move to correct device + args, kwargs = tree_to((args, kwargs or {}), device="meta") + + # NOTE (lucaslie): export is VERY sensitive to the location of the inference_mode + # context manager. Do NOT move it unless absolutely necessary. + with torch.inference_mode(): + ep = te.export(model, args, kwargs, dynamic_shapes=dynamic_shapes, strict=strict) + egm = ep.module() + assert isinstance(egm, fx.GraphModule) + + # load state_dict into egm + # NOTE: export might have removed unused params/buffers (hence we allow unexpected keys) + load_buffers_and_params( + egm, state_dict, strict_missing=True, strict_unexpected=False, clone=clone + ) + + # Export strips away all methods not traced during forward. The model could have + # load hooks that contain logic for correct state_dict loading. We need to add those + # hooks back to the exported graph module. + _add_missing_load_hooks(egm, model) + + # Add load hook to correctly load parameters that are aliased in the source model. + # deduplicate params and buffers + # TODO (lucaslie, suyoggupta): seems there is some overlap here. I believe we should just have + # the deduplicate function and extend it to handle reading from state dict for any name. + _add_load_hook_for_aliased_params(egm, model) + _deduplicate_params_and_buffers(egm) + + # clean up devices in the graph + # This is a consequence of lifting to meta during export. + _clean_up_device_info(egm) + + # clean up checks --> generally the sanity checks are overly conservative and we can remove them + _clean_up_assertions(egm) + + # show exported graph + ad_logger.debug("exported graph: " + str(egm)) + + return egm diff --git a/tensorrt_llm/_torch/auto_deploy/export/interface.py b/tensorrt_llm/_torch/auto_deploy/export/interface.py new file mode 100644 index 000000000000..c97b056a00d6 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/export/interface.py @@ -0,0 +1,249 @@ +"""The interface for all export patches. + +This module defines the base classes and interfaces for all export patches. +""" + +from abc import ABC, abstractmethod +from contextlib import contextmanager +from typing import Any, Callable, Dict, List, Type, Union, final + +from pydantic import BaseModel, Field + +from ..utils.logger import ad_logger + + +class ExportPatchError(Exception): + """An exception raised when an export patch fails.""" + + pass + + +class ExportPatchConfig(BaseModel): + """Base configuration class for export patches.""" + + model_config = { + "extra": "allow", # Allow subclasses to add more fields + } + + enabled: bool = Field( + default=True, + description="Whether to enable this patch.", + ) + skip_on_error: bool = Field( + default=False, + description="Whether to skip the patch if an error occurs during application.", + ) + + +class BaseExportPatch(ABC): + """Base class for all export patches. + + Export patches are context managers that apply temporary modifications + to the global state during torch.export, then revert them afterwards. + """ + + config: ExportPatchConfig + _patch_key: str # Set by ExportPatchRegistry.register() decorator + + @classmethod + def get_patch_key(cls) -> str: + """Get the short name of the patch.""" + if hasattr(cls, "_patch_key"): + return cls._patch_key + raise NotImplementedError( + f"Patch class {cls.__name__} must be registered with ExportPatchRegistry.register() " + "or manually implement get_patch_key()" + ) + + @classmethod + def get_config_class(cls) -> Type[ExportPatchConfig]: + """Get the configuration class for the patch.""" + return ExportPatchConfig + + @final + def __init__(self, config: ExportPatchConfig): + """Initialize the patch. + + Args: + config: The configuration for the patch. + """ + if not isinstance(config, self.get_config_class()): + config = self.get_config_class()(**config.model_dump()) + self.config = config + self.original_values = {} + self._post_init() + + def _post_init(self): + """Post-initialization hook that can be overridden by subclasses.""" + pass + + @final + @classmethod + def from_kwargs(cls, **kwargs) -> "BaseExportPatch": + """Create a patch from kwargs.""" + config = cls.get_config_class()(**kwargs) + return cls(config=config) + + @final + def __enter__(self): + """Enter the context manager and apply the patch.""" + if not self.config.enabled: + ad_logger.debug(f"Patch {self.get_patch_key()} is disabled, skipping") + return self + + try: + ad_logger.debug(f"Applying patch: {self.get_patch_key()}") + self._apply_patch() + except Exception as e: + error_msg = f"Patch {self.get_patch_key()} failed to apply" + if self.config.skip_on_error: + ad_logger.warning(f"{error_msg}: {e}") + else: + raise ExportPatchError(error_msg) from e + + return self + + @final + def __exit__(self, exc_type, exc_val, exc_tb): + """Exit the context manager and revert the patch.""" + if not self.config.enabled: + return + + try: + ad_logger.debug(f"Reverting patch: {self.get_patch_key()}") + self._revert_patch() + except Exception as e: + error_msg = f"Patch {self.get_patch_key()} failed to revert" + if self.config.skip_on_error: + ad_logger.warning(f"{error_msg}: {e}") + else: + raise ExportPatchError(error_msg) from e + + @abstractmethod + def _apply_patch(self): + """Apply the patch. Should store original values in self.original_values.""" + pass + + @abstractmethod + def _revert_patch(self): + """Revert the patch using stored original values.""" + pass + + +class ContextManagerPatch(BaseExportPatch): + """A patch that wraps an existing context manager. + + This allows easy registration of context managers as patches without + having to implement the full BaseExportPatch interface. + + Subclasses must implement `init_context_manager()` to return the context manager. + """ + + def _post_init(self): + self.context_manager: Any = None + + @abstractmethod + def init_context_manager(self) -> Any: + """Initialize and return the context manager. + + Returns: + A context manager that will be used during export. + """ + pass + + def _apply_patch(self): + """Apply the patch by entering the context manager.""" + self.context_manager = self.init_context_manager() + self.context_manager.__enter__() + + def _revert_patch(self): + """Revert the patch by exiting the context manager.""" + if self.context_manager is not None: + self.context_manager.__exit__(None, None, None) + self.context_manager = None + + +class ExportPatchRegistry: + """Registry for export patches.""" + + _registry: Dict[str, Type[BaseExportPatch]] = {} + + @classmethod + def register(cls, name: str) -> Callable[[Type[BaseExportPatch]], Type[BaseExportPatch]]: + """Register a patch class with the given name.""" + + def inner(patch_cls: Type[BaseExportPatch]) -> Type[BaseExportPatch]: + cls._registry[name] = patch_cls + # Auto-store the patch key as a class attribute + patch_cls._patch_key = name + return patch_cls + + return inner + + @classmethod + def get(cls, name: str) -> Type[BaseExportPatch]: + """Get a patch class by name.""" + return cls._registry[name] + + @classmethod + def get_config_class(cls, name: str) -> Type[ExportPatchConfig]: + """Get the configuration class for a patch by name.""" + return cls.get(name).get_config_class() + + @classmethod + def has(cls, name: str) -> bool: + """Check if a patch is registered.""" + return name in cls._registry + + @classmethod + def create_patch( + cls, name: str, config: Union[ExportPatchConfig, Dict[str, Any]] + ) -> BaseExportPatch: + """Create a patch instance by name.""" + patch_cls = cls.get(name) + if isinstance(config, dict): + config = patch_cls.get_config_class()(**config) + return patch_cls(config) + + @classmethod + def list_patches(cls) -> List[str]: + """List all registered patch names.""" + return list(cls._registry.keys()) + + +@contextmanager +def apply_export_patches(patch_configs: Dict[str, Union[ExportPatchConfig, Dict[str, Any]]]): + """Context manager to apply multiple patches. + + Args: + patch_configs: Dict mapping patch names to their configurations. + """ + patches = [] + + # Create patch instances + for name, config in patch_configs.items(): + if not ExportPatchRegistry.has(name): + raise ValueError(f"Unknown patch: {name}") + patch = ExportPatchRegistry.create_patch(name, config) + patches.append(patch) + + # Apply patches using nested context managers + if not patches: + yield + return + + def _apply_patches(remaining_patches): + if not remaining_patches: + yield + return + + patch = remaining_patches[0] + with patch: + yield from _apply_patches(remaining_patches[1:]) + + # log applied patches + ad_logger.debug( + f"applying export patches: {', '.join([patch.get_patch_key() for patch in patches])}" + ) + + yield from _apply_patches(patches) diff --git a/tensorrt_llm/_torch/auto_deploy/export/library/__init__.py b/tensorrt_llm/_torch/auto_deploy/export/library/__init__.py new file mode 100644 index 000000000000..fcc425ad26d1 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/export/library/__init__.py @@ -0,0 +1,16 @@ +"""AutoDeploy's library of export patches. + +This file ensures that all publicly listed files/patches in the library folder are auto-imported +and the corresponding patches are registered. +""" + +import importlib +import pkgutil + +__all__ = [] + +for _, module_name, is_pkg in pkgutil.iter_modules(__path__): + if module_name.startswith("_"): + continue + __all__.append(module_name) + importlib.import_module(f"{__name__}.{module_name}") diff --git a/tensorrt_llm/_torch/auto_deploy/export/library/autocast_noop.py b/tensorrt_llm/_torch/auto_deploy/export/library/autocast_noop.py new file mode 100644 index 000000000000..4392b6ba3715 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/export/library/autocast_noop.py @@ -0,0 +1,28 @@ +"""Patch to make torch.autocast a no-op during export.""" + +from contextlib import nullcontext + +import torch + +from ..interface import BaseExportPatch, ExportPatchRegistry + + +@ExportPatchRegistry.register("autocast_noop") +class AutocastNoopPatch(BaseExportPatch): + """Patch torch.autocast to be a no-op during export. + + This patch replaces torch.autocast with a null context manager + that can interfere with export. + """ + + def _apply_patch(self): + """Apply the autocast no-op patch.""" + # Store original function + self.original_values["torch.autocast"] = torch.autocast + + # Apply patch + torch.autocast = lambda *args, **kwargs: nullcontext() + + def _revert_patch(self): + """Revert the autocast no-op patch.""" + torch.autocast = self.original_values["torch.autocast"] diff --git a/tensorrt_llm/_torch/auto_deploy/export/library/linear.py b/tensorrt_llm/_torch/auto_deploy/export/library/linear.py new file mode 100644 index 000000000000..b8304671250d --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/export/library/linear.py @@ -0,0 +1,35 @@ +"""Patch for F.linear to use simpler implementation during export.""" + +from typing import Optional + +import torch +import torch.nn.functional as F + +from ..interface import BaseExportPatch, ExportPatchRegistry + + +@ExportPatchRegistry.register("linear") +class LinearPatch(BaseExportPatch): + """Patch F.linear to use a simpler implementation for export. + + This patch replaces F.linear with a version that avoids exporting + view operations used to flatten/unflatten multiple batch dimensions. + """ + + def _apply_patch(self): + """Apply the linear patch.""" + # Store original function + self.original_values["F.linear"] = F.linear + + # Create patched function + def _torch_linear_patch( + input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None + ) -> torch.Tensor: + return torch.ops.auto_deploy.torch_linear_simple(input, weight, bias) + + # Apply patch + F.linear = _torch_linear_patch + + def _revert_patch(self): + """Revert the linear patch.""" + F.linear = self.original_values["F.linear"] diff --git a/tensorrt_llm/_torch/auto_deploy/export/library/modelopt_context.py b/tensorrt_llm/_torch/auto_deploy/export/library/modelopt_context.py new file mode 100644 index 000000000000..d6f27cd31906 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/export/library/modelopt_context.py @@ -0,0 +1,23 @@ +"""Patch for modelopt's torch_export_context.""" + +from contextlib import nullcontext + +from ..interface import ContextManagerPatch, ExportPatchRegistry + + +@ExportPatchRegistry.register("modelopt_context") +class ModeloptContextPatch(ContextManagerPatch): + """Patch to apply modelopt's torch_export_context during export. + + This patch applies the modelopt quantization context manager around + the export process when available, otherwise uses a null context. + """ + + def init_context_manager(self): + """Initialize and return the modelopt context manager or nullcontext if not available.""" + try: + from modelopt.torch.quantization.utils import export_torch_mode as torch_export_context + + return torch_export_context() + except ImportError: + return nullcontext() diff --git a/tensorrt_llm/_torch/auto_deploy/export/library/sdpa.py b/tensorrt_llm/_torch/auto_deploy/export/library/sdpa.py new file mode 100644 index 000000000000..475b0c71b2aa --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/export/library/sdpa.py @@ -0,0 +1,27 @@ +"""Patch for F.scaled_dot_product_attention to use custom op.""" + +import torch +import torch.nn.functional as F + +from ..interface import BaseExportPatch, ExportPatchRegistry + + +@ExportPatchRegistry.register("sdpa") +class SdpaPatch(BaseExportPatch): + """Patch F.scaled_dot_product_attention to use custom op during export. + + This patch ensures that scaled_dot_product_attention is represented consistently + in the exported graph by using a custom operation. + """ + + def _apply_patch(self): + """Apply the SDPA patch.""" + # Store original function + self.original_values["F.scaled_dot_product_attention"] = F.scaled_dot_product_attention + + # Apply patch + F.scaled_dot_product_attention = torch.ops.auto_deploy.torch_attention_sdpa + + def _revert_patch(self): + """Revert the SDPA patch.""" + F.scaled_dot_product_attention = self.original_values["F.scaled_dot_product_attention"] diff --git a/tensorrt_llm/_torch/auto_deploy/export/library/sdpa_kernel_noop.py b/tensorrt_llm/_torch/auto_deploy/export/library/sdpa_kernel_noop.py new file mode 100644 index 000000000000..52dec06cd971 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/export/library/sdpa_kernel_noop.py @@ -0,0 +1,28 @@ +"""Patch to make torch.nn.attention.sdpa_kernel a no-op during export.""" + +from contextlib import nullcontext + +import torch + +from ..interface import BaseExportPatch, ExportPatchRegistry + + +@ExportPatchRegistry.register("sdpa_kernel_noop") +class SdpaKernelNoopPatch(BaseExportPatch): + """Patch torch.nn.attention.sdpa_kernel to be a no-op during export. + + This patch replaces torch.nn.attention.sdpa_kernel with a null context manager + that can interfere with export. + """ + + def _apply_patch(self): + """Apply the sdpa_kernel no-op patch.""" + # Store original function + self.original_values["torch.nn.attention.sdpa_kernel"] = torch.nn.attention.sdpa_kernel + + # Apply patch + torch.nn.attention.sdpa_kernel = lambda *args, **kwargs: nullcontext() + + def _revert_patch(self): + """Revert the sdpa_kernel no-op patch.""" + torch.nn.attention.sdpa_kernel = self.original_values["torch.nn.attention.sdpa_kernel"] diff --git a/tensorrt_llm/_torch/auto_deploy/export/library/tensor_meta_device.py b/tensorrt_llm/_torch/auto_deploy/export/library/tensor_meta_device.py new file mode 100644 index 000000000000..45879897496f --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/export/library/tensor_meta_device.py @@ -0,0 +1,33 @@ +"""Patch for torch.tensor to handle 0.0 on meta device.""" + +import torch + +from ..interface import BaseExportPatch, ExportPatchRegistry + + +@ExportPatchRegistry.register("tensor_meta_device") +class TensorMetaDevicePatch(BaseExportPatch): + """Patch torch.tensor to handle 0.0 on meta device. + + This patch addresses an issue where torch.tensor(0.0, device="meta") + doesn't work and needs to be replaced with torch.zeros((), device="meta"). + """ + + def _apply_patch(self): + """Apply the tensor meta device patch.""" + # Store original function + self.original_values["torch.tensor"] = torch.tensor + + # Create patched function + def _torch_tensor_patch(data, **kwargs): + device = kwargs.get("device", None) + if data == 0.0 and device is not None and torch.device(device) == torch.device("meta"): + return torch.zeros((), **kwargs) + return self.original_values["torch.tensor"](data, **kwargs) + + # Apply patch + torch.tensor = _torch_tensor_patch + + def _revert_patch(self): + """Revert the tensor meta device patch.""" + torch.tensor = self.original_values["torch.tensor"] diff --git a/tensorrt_llm/_torch/auto_deploy/export/library/torch_modulelist_getitem.py b/tensorrt_llm/_torch/auto_deploy/export/library/torch_modulelist_getitem.py new file mode 100644 index 000000000000..e97670146bc2 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/export/library/torch_modulelist_getitem.py @@ -0,0 +1,43 @@ +"""Patch for nn.ModuleList.__getitem__ to handle slicing during export.""" + +import torch.nn as nn + +from ..interface import BaseExportPatch, ExportPatchRegistry + + +@ExportPatchRegistry.register("torch_modulelist_getitem") +class TorchModuleListGetitemPatch(BaseExportPatch): + """Patch nn.ModuleList.__getitem__ to handle slicing during export. + + This patch addresses a PyTorch issue where nn.ModuleList.__getitem__ with slice + indexing doesn't work correctly during export. The workaround returns a simple + list for slice operations. + + Reference: https://github.com/pytorch/pytorch/issues/142439 + """ + + def _apply_patch(self): + """Apply the ModuleList getitem patch.""" + # Store original function + self.original_values["nn.ModuleList.__getitem__"] = nn.ModuleList.__getitem__ + + # Capture the original function for use in closure + original_getitem = nn.ModuleList.__getitem__ + + # Create patched function + def _torch_modulelist_getitem_patch(self: nn.ModuleList, idx): + if isinstance(idx, slice): + # return a simple list. + # NOTE: this obviously only works for any use case where we access the sliced module list + # like a regular list like a for-loop. For most other things, this hack will not work. + return list(self._modules.values())[idx] + else: + # Call the original function + return original_getitem(self, idx) + + # Apply patch (type ignore needed as return type differs for slice case) + nn.ModuleList.__getitem__ = _torch_modulelist_getitem_patch # type: ignore + + def _revert_patch(self): + """Revert the ModuleList getitem patch.""" + nn.ModuleList.__getitem__ = self.original_values["nn.ModuleList.__getitem__"] diff --git a/tensorrt_llm/_torch/auto_deploy/export/library/torch_where.py b/tensorrt_llm/_torch/auto_deploy/export/library/torch_where.py new file mode 100644 index 000000000000..071eff221bd2 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/export/library/torch_where.py @@ -0,0 +1,33 @@ +"""Patch for torch.where to handle case where only condition is provided.""" + +import torch + +from ..interface import BaseExportPatch, ExportPatchRegistry + + +@ExportPatchRegistry.register("torch_where") +class TorchWherePatch(BaseExportPatch): + """Patch torch.where to handle the case where only condition is provided. + + This patch addresses the issue where torch.where(condition) should return + torch.nonzero(condition, as_tuple=True) but the export process doesn't + handle this correctly. + """ + + def _apply_patch(self): + """Apply the torch.where patch.""" + # Store original function + self.original_values["torch.where"] = torch.where + + # Create patched function + def _torch_where_patch(condition: torch.Tensor, *args, **kwargs): + if len(args) == 0 and len(kwargs) == 0: + return torch.nonzero(condition, as_tuple=True) + return self.original_values["torch.where"](condition, *args, **kwargs) + + # Apply patch + torch.where = _torch_where_patch + + def _revert_patch(self): + """Revert the torch.where patch.""" + torch.where = self.original_values["torch.where"] diff --git a/tensorrt_llm/_torch/auto_deploy/export/library/transformers_sdpa_mask.py b/tensorrt_llm/_torch/auto_deploy/export/library/transformers_sdpa_mask.py new file mode 100644 index 000000000000..fd21604d1b61 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/export/library/transformers_sdpa_mask.py @@ -0,0 +1,78 @@ +"""Patch for transformers SDPA mask to be export-compatible.""" + +import importlib.metadata + +from packaging import version + +from ..interface import BaseExportPatch, ExportPatchRegistry + + +def _transformers_version() -> str: + """Get the version of transformers.""" + return version.parse(importlib.metadata.version("transformers")).base_version + + +@ExportPatchRegistry.register("transformers_sdpa_mask") +class TransformersSdpaMaskPatch(BaseExportPatch): + """Patch transformers.masking_utils.sdpa_mask to be export-compatible. + + This patch replaces the transformers SDPA mask implementation with an + export-compatible version for transformers >= 4.53.0. + """ + + def _apply_patch(self): + """Apply the transformers SDPA mask patch.""" + # this patch is only needed+compatible for transformers >= 4.53.0 + if version.parse(_transformers_version()) < version.parse("4.53.0"): + return # Skip patch for older versions + + try: + # imports only after version check + from transformers import masking_utils + from transformers.integrations.executorch import sdpa_mask_without_vmap + + # recall original implementation + self.original_values["masking_utils.sdpa_mask"] = masking_utils.sdpa_mask + + # patch function and mask attention interface + masking_utils.sdpa_mask = sdpa_mask_without_vmap + + if "sdpa" in masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._local_mapping: + self.original_values["sdpa_local_original"] = ( + masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._local_mapping["sdpa"] + ) + else: + self.original_values["sdpa_local_original"] = None + + masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"] = sdpa_mask_without_vmap + + except ImportError: + # If transformers is not available or doesn't have required modules, skip patch + pass + + def _revert_patch(self): + """Revert the transformers SDPA mask patch.""" + # this patch is only needed+compatible for transformers >= 4.53.0 + if version.parse(_transformers_version()) < version.parse("4.53.0"): + return # Skip revert for older versions + + try: + # imports only after version check + from transformers import masking_utils + + # revert patches + if "masking_utils.sdpa_mask" in self.original_values: + masking_utils.sdpa_mask = self.original_values["masking_utils.sdpa_mask"] + + if "sdpa_local_original" in self.original_values: + if self.original_values["sdpa_local_original"] is None: + if "sdpa" in masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._local_mapping: + del masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"] + else: + masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"] = self.original_values[ + "sdpa_local_original" + ] + + except ImportError: + # If transformers is not available, skip revert + pass diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py index ba6ad81595bb..61337ae3f420 100644 --- a/tensorrt_llm/_torch/auto_deploy/llm_args.py +++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py @@ -1,35 +1,60 @@ -import json +from importlib.resources import files from pathlib import Path -from typing import Any, Dict, List, Literal, Optional, Union +from typing import Any, Dict, List, Literal, Optional, Type, Union import torch -from pydantic import Field, field_validator, model_validator +from pydantic import Field, ValidationInfo, field_validator, model_validator +from pydantic_settings import BaseSettings, SettingsConfigDict from ...llmapi.llm_args import BaseLlmArgs, BuildConfig, _ParallelConfig from ...llmapi.utils import get_type_repr from .models import ModelFactory, ModelFactoryRegistry +from .transform.interface import TransformConfig +from .utils._config import DynamicYamlMixInForSettings +PathLike = Union[str, Path] -def _try_decode_dict_with_str_values(value: Dict[str, Any]) -> Dict[str, Any]: - """Try to parse string values as JSON to convert to native types if possible.""" - for k, v in value.items(): - if isinstance(v, str): - try: - value[k] = json.loads(v) - except json.JSONDecodeError: - pass + +def _get_config_dict() -> SettingsConfigDict: + return SettingsConfigDict( + arbitrary_types_allowed=True, + extra="forbid", + yaml_file=str(files("tensorrt_llm._torch.auto_deploy.config") / "default.yaml"), + nested_model_default_partial_update=True, + ) + + +def _check_for_default_value_only( + cls: Type[BaseSettings], value: Any, info: ValidationInfo, msg: str +) -> Any: + """Check if the value is the default value for the field. + + If the value is not the default value, raise a ValueError. + """ + field_name = info.field_name + assert field_name is not None, "field_name should be set for validated field." + if value != cls.model_fields[field_name].get_default(call_default_factory=True): + raise ValueError(msg) return value -class LlmArgs(BaseLlmArgs): - """LLM arguments specifically for AutoDeploy backend. +class AutoDeployConfig(DynamicYamlMixInForSettings, BaseSettings): + """An argument class stripped down to AutoDeploy-specific configurations. + + This class be used as a drop-in replacement to simplify configuring the AutoDeploy backend and + should be used in place of LlmArgs unless more advanced features are needed. - This class extends BaseLlmArgs with AutoDeploy-specific configuration options. - AutoDeploy provides automatic deployment and optimization of language models - with various attention backends and optimization strategies. + It is compatible with AutoDeploy's LLM API (``tensorrt_llm._torch.auto_deploy.llm.LLM``) and + exposes the full set of parameters used in AutoDeploy's ``InferenceOptimizer``. """ + model_config = _get_config_dict() + ### MODEL AND TOKENIZER FACTORY ################################################################ + model: PathLike = Field( + description="The path to the model checkpoint or the model name from the Hugging Face Hub." + ) + model_factory: Literal["AutoModelForCausalLM", "AutoModelForImageTextToText"] = Field( default="AutoModelForCausalLM", description="The model factory to use for loading the model.", @@ -56,7 +81,7 @@ class LlmArgs(BaseLlmArgs): "Defaults to the same device as the rest of the pipeline.", ) - tokenizer: Optional[Union[str, Path]] = Field( + tokenizer: Optional[PathLike] = Field( description="The tokenizer", default=None, repr=False, @@ -70,13 +95,14 @@ class LlmArgs(BaseLlmArgs): "https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/tokenization_llama_fast.py#L127.", ) + skip_tokenizer_init: bool = Field( + default=False, description="Whether to skip the tokenizer initialization." + ) + ### RUNTIME FEATURES ########################################################################### disable_overlap_scheduler: bool = Field( - default=True, - description="Disable the overlap scheduler. This is a temporary field until the overlap " - "scheduler is supported (https://github.com/NVIDIA/TensorRT-LLM/issues/4364).", - frozen=True, - repr=False, + default=False, + description="Disable the overlap scheduler in trtllm runtime", ) enable_mixed_sampler: bool = Field( @@ -102,8 +128,14 @@ class LlmArgs(BaseLlmArgs): "supported in AutoDeploy.", ) - # INFERENCE OPTIMIZER CONFIG ################################################################### - attn_backend: Literal["flashinfer", "triton"] = Field( + max_beam_width: int = Field( + default=1, + description="The maximum beam width. >1 is not supported by AutoDeploy.", + frozen=True, + ) + + ### INFERENCE OPTIMIZER CONFIG ################################################################# + attn_backend: Literal["flashinfer", "triton", "torch"] = Field( default="flashinfer", description="Attention backend to use." ) @@ -138,18 +170,75 @@ class LlmArgs(BaseLlmArgs): visualize: bool = Field(default=False, description="Whether to visualize the model graph.") + ### NEW INFERENCE OPTIMIZER CONFIG ############################################################# + transforms: Dict[str, TransformConfig] = Field( + default_factory=dict, + description="A dictionary of transform configurations. The key is the transform name and " + "the value is the transform configuration.", + ) + ### SEQUENCE INTERFACE CONFIG ################################################################## + max_input_len: int = Field(default=1024, description="The maximum input length.") + max_num_tokens: Optional[int] = Field(default=None, description="The maximum number of tokens.") max_seq_len: int = Field(default=512, ge=1, description="The maximum sequence length.") max_batch_size: int = Field(default=8, ge=1, description="The maximum batch size.") attn_page_size: int = Field( default=64, ge=1, - description="Page size for attention (tokens_per_block). For triton " - "backend, this should equal max_seq_len. Temporary field until tokens_per_block gets " + description="Page size for attention (tokens_per_block). For triton and torch " + "backends, this should equal max_seq_len. Temporary field until tokens_per_block gets " "properly passed through.", ) - ### !!! DO NOT USE !!! ######################################################################### + ### VALIDATION ################################################################################# + @model_validator(mode="after") + def update_attn_page_size(self): + # NOTE force attn_page_size to equal max_seq_len for triton backend + if self.attn_backend == "triton" or self.attn_backend == "torch": + self.attn_page_size = self.max_seq_len + return self + + ### UTILITY METHODS ############################################################################ + def create_factory(self) -> ModelFactory: + """Create a model factory from the arguments.""" + + # TODO (lucaslie): consider supporting Path objects in the model factory + return ModelFactoryRegistry.get(self.model_factory)( + model=str(self.model), + model_kwargs=self.model_kwargs, + tokenizer=None if self.tokenizer is None else str(self.tokenizer), + tokenizer_kwargs=self.tokenizer_kwargs, + skip_loading_weights=self.skip_loading_weights, + max_seq_len=self.max_seq_len, + ) + + def to_dict(self) -> Dict[str, Any]: + """Convert the arguments to a dictionary.""" + return self.model_dump() + + def to_llm_args(self) -> "LlmArgs": + """Convert the arguments to a LlmArgs instance that is used for the LLM API.""" + return LlmArgs(**self.to_dict()) + + +class LlmArgs(AutoDeployConfig, BaseLlmArgs, BaseSettings): + """LlmArgs config class for providing full expert configurability of the AutoDeploy backend. + + Specifically, this class extends AutoDeployConfig with all the fields from BaseLlmArgs for + providing configurability beyond what is provided by AutoDeployConfig. + + Just like AutoDeployConfig, this class is compatible with AutoDeploy's LLM API + (``tensorrt_llm._torch.auto_deploy.llm.LLM``) but provides greater configurability. + + NOTE: this class should only be used directly for advanced use cases. For most use cases, + AutoDeployConfig should be used instead. + + NOTE: this class may expose redundant fields from BaseLlmArgs or fields that are ignored or + have overlapping functionality with AutoDeployConfig. Please be careful when using this class. + """ + + model_config = _get_config_dict() + build_config: Optional[object] = Field( default_factory=lambda: BuildConfig(), description="!!! DO NOT USE !!! Internal only; needed for BaseLlmArgs compatibility.", @@ -173,16 +262,25 @@ class LlmArgs(BaseLlmArgs): ### VALIDATION ################################################################################# @field_validator("build_config", mode="before") @classmethod - def ensure_no_build_config(cls, value: Any) -> Any: - if value is not None: - raise ValueError("build_config is not used") - return value - - @field_validator("model_kwargs", "tokenizer_kwargs", mode="after") + 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", + "context_parallel_size", + "moe_cluster_parallel_size", + "moe_tensor_parallel_size", + "moe_expert_parallel_size", + "enable_attention_dp", + "cp_config", + mode="before", + ) @classmethod - def validate_model_kwargs(cls, value: Dict[str, Any]) -> Dict[str, Any]: - """Try to parse string values as JSON to convert to native types if possible.""" - return _try_decode_dict_with_str_values(value) + def ensure_no_custom_parallel_config(cls, value: Any, info: ValidationInfo) -> Any: + msg = "AutoDeploy only supports parallelization via the `world_size` argument." + return _check_for_default_value_only(cls, value, info, msg) @model_validator(mode="after") def validate_parallel_config(self): @@ -192,7 +290,6 @@ def validate_parallel_config(self): rank to automatically shard the model. This is just to ensure that other objects in the runtime that may read parallel_config can do so. """ - # setup parallel config self._parallel_config = _ParallelConfig( auto_parallel=True, gpus_per_node=self.gpus_per_node ) @@ -204,26 +301,7 @@ def validate_and_init_tokenizer(self): """Skip tokenizer initialization in config. We do this in the AutoDeploy LLM class.""" return self - @model_validator(mode="after") - def update_attn_page_size(self): - # NOTE force attn_page_size to equal max_seq_len for triton backend - if self.attn_backend == "triton": - self.attn_page_size = self.max_seq_len - return self - ### UTILITY METHODS ############################################################################ - def create_factory(self) -> ModelFactory: - """Create a model factory from the arguments.""" - - return ModelFactoryRegistry.get(self.model_factory)( - model=self.model, - model_kwargs=self.model_kwargs, - tokenizer=self.tokenizer, - tokenizer_kwargs=self.tokenizer_kwargs, - skip_loading_weights=self.skip_loading_weights, - max_seq_len=self.max_seq_len, - ) - # TODO: Remove this after the PyTorch backend is fully migrated to LlmArgs from ExecutorConfig def get_pytorch_backend_config(self) -> "LlmArgs": """Return the LlmArgs (self) object.""" diff --git a/tensorrt_llm/_torch/auto_deploy/models/__init__.py b/tensorrt_llm/_torch/auto_deploy/models/__init__.py index 8e1fd728bba1..a004f7a8b134 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/models/__init__.py @@ -1,7 +1,2 @@ -from . import hf -from .decilm import * -from .deepseek import * +from . import hf, patches from .factory import * -from .mixtral import * -from .phi import * -from .qwen3 import * diff --git a/tensorrt_llm/_torch/auto_deploy/models/factory.py b/tensorrt_llm/_torch/auto_deploy/models/factory.py index 1f0617706a9c..42a304025370 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/factory.py +++ b/tensorrt_llm/_torch/auto_deploy/models/factory.py @@ -211,9 +211,7 @@ class ModelFactoryRegistry: _registry: Dict[str, Type[ModelFactory]] = {} @classmethod - def register( - cls: Type[ModelFactory], name: str - ) -> Callable[[Type[ModelFactory]], Type[ModelFactory]]: + def register(cls, name: str) -> Callable[[Type[ModelFactory]], Type[ModelFactory]]: def inner(fn: Type[ModelFactory]) -> Type[ModelFactory]: cls._registry[name] = fn return fn diff --git a/tensorrt_llm/_torch/auto_deploy/models/hf.py b/tensorrt_llm/_torch/auto_deploy/models/hf.py index 6295f291e90e..f407a0425383 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/hf.py +++ b/tensorrt_llm/_torch/auto_deploy/models/hf.py @@ -28,6 +28,7 @@ ) from ..custom_ops.attention_interface import CacheConfig +from ..utils._config import deep_merge_dicts from ..utils.logger import ad_logger from .factory import ModelFactory, ModelFactoryRegistry @@ -62,25 +63,27 @@ def load_state_dict_with_device(checkpoint_file, device_map=None): @ModelFactoryRegistry.register("AutoModelForCausalLM") class AutoModelForCausalLMFactory(ModelFactory): + _tokenizer_defaults = { + "legacy": False, + "padding_side": "left", + "truncation_side": "left", + "trust_remote_code": True, + "use_fast": True, + } + + _model_defaults = { + "use_cache": False, + "max_position_embeddings": 1024, + } + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._quant_config: Optional[Dict] = None - # Relevant default tokenizer kwargs for HF-style tokenizer - defaults = { - "legacy": False, - "padding_side": "left", - "truncation_side": "left", - "trust_remote_code": True, - "use_fast": True, - } - self.tokenizer_kwargs = {**defaults, **self.tokenizer_kwargs} - - # NEVER use cache - self.model_kwargs["use_cache"] = False - # Ensure max_seq_len is propagated to model_kwargs - self.model_kwargs["max_position_embeddings"] = self.max_seq_len + # Ingest defaults for tokenizer and model kwargs + self.tokenizer_kwargs = deep_merge_dicts(self._tokenizer_defaults, self.tokenizer_kwargs) + self.model_kwargs = deep_merge_dicts(self._model_defaults, self.model_kwargs) # special handling for torch_dtype in model_kwargs since HF does not correctly update # torch_dtype string to an actual torch.dtype object (only with default) @@ -114,7 +117,7 @@ def _simple_forward(model: nn.Module, input_ids: torch.Tensor, position_ids: tor def _recursive_update_config(self, config: PretrainedConfig, update_dict: Dict[str, Any]): """ - Recursively update a PretrainedConfig object with values from update_dict. + Deep-merge a PretrainedConfig object with values from update_dict. Args: config: PretrainedConfig object to update @@ -302,7 +305,13 @@ def _load_checkpoint(self, model: nn.Module, device: DeviceLikeType): ckpt_file = self._get_checkpoint_file(self.model) # reuse the load checkpoint utility from accelerate with hf_load_state_dict_with_device(device): - load_checkpoint_in_model(model, checkpoint=ckpt_file) + # Set `full_state_dict=False` to skip Accelerate's FSDP weight sync logic. + # Internally, load_checkpoint_in_model → set_model_state_dict → _load_model_state_dict, + # which collects local model params, syncs weights from checkpoint, and applies them via + # model.load_state_dict. + # This sync step can interfere with load_hooks by mixing raw checkpoint weights and + # model-transformed weights,leading to unexpected key mismatches or format issues. + load_checkpoint_in_model(model, checkpoint=ckpt_file, full_state_dict=False) def _load_quantization_config(self): """Load the quantization config from the model directory if not done already.""" @@ -326,21 +335,14 @@ def _load_quantization_config(self): @ModelFactoryRegistry.register("AutoModelForImageTextToText") class AutoModelForImageTextToTextFactory(AutoModelForCausalLMFactory): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - # additional heuristic to propagate "important keys" - # TODO (lucaslie): WAR until we have better support on dashboard to control model_kwargs - keys_to_propagate = [ - "num_hidden_layers", - "max_position_embeddings", - "use_cache", - "torch_dtype", - ] - self.model_kwargs["text_config"] = self.model_kwargs.get("text_config", {}) - for key in keys_to_propagate: - if key in self.model_kwargs: - self.model_kwargs["text_config"][key] = self.model_kwargs[key] + _model_defaults = { + "use_cache": False, + "max_position_embeddings": 1024, + "text_config": { + "max_position_embeddings": 1024, + "use_cache": False, + }, + } @property def automodel_from_config(self): diff --git a/tensorrt_llm/_torch/auto_deploy/models/patches/__init__.py b/tensorrt_llm/_torch/auto_deploy/models/patches/__init__.py new file mode 100644 index 000000000000..e98cf311b383 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/models/patches/__init__.py @@ -0,0 +1,16 @@ +"""AutoDeploy's library of export patches for models. + +This file ensures that all publicly listed files/patches in the library folder are auto-imported +and the corresponding patches are registered. +""" + +import importlib +import pkgutil + +__all__ = [] + +for _, module_name, is_pkg in pkgutil.iter_modules(__path__): + if module_name.startswith("_"): + continue + __all__.append(module_name) + importlib.import_module(f"{__name__}.{module_name}") diff --git a/tensorrt_llm/_torch/auto_deploy/models/decilm.py b/tensorrt_llm/_torch/auto_deploy/models/patches/decilm.py similarity index 86% rename from tensorrt_llm/_torch/auto_deploy/models/decilm.py rename to tensorrt_llm/_torch/auto_deploy/models/patches/decilm.py index 1a9f7368a646..c8989d62cc6b 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/decilm.py +++ b/tensorrt_llm/_torch/auto_deploy/models/patches/decilm.py @@ -12,4 +12,5 @@ def _from_pretrained_patched(pretrained_model_name_or_path, **kwargs): return _orig_from_pretrained(pretrained_model_name_or_path, **kwargs) +# TODO: figure out how this can be incorporated into the export patch system AutoConfig.from_pretrained = _from_pretrained_patched diff --git a/tensorrt_llm/_torch/auto_deploy/models/deepseek.py b/tensorrt_llm/_torch/auto_deploy/models/patches/deepseek.py similarity index 98% rename from tensorrt_llm/_torch/auto_deploy/models/deepseek.py rename to tensorrt_llm/_torch/auto_deploy/models/patches/deepseek.py index ae04bf6e592b..f30bc0c6fac5 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/deepseek.py +++ b/tensorrt_llm/_torch/auto_deploy/models/patches/deepseek.py @@ -181,4 +181,5 @@ def get_model_from_config_patched(config, **kwargs): return model +# TODO: figure out how this can be incorporated into the export patch system AutoModelForCausalLM.from_config = get_model_from_config_patched diff --git a/tensorrt_llm/_torch/auto_deploy/models/mixtral.py b/tensorrt_llm/_torch/auto_deploy/models/patches/mixtral.py similarity index 62% rename from tensorrt_llm/_torch/auto_deploy/models/mixtral.py rename to tensorrt_llm/_torch/auto_deploy/models/patches/mixtral.py index b0511a0ed946..b759fe6495d1 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/mixtral.py +++ b/tensorrt_llm/_torch/auto_deploy/models/patches/mixtral.py @@ -5,6 +5,8 @@ import torch.nn.functional as F from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock +from ...export.interface import BaseExportPatch, ExportPatchRegistry + def _forward_moe(self: MixtralSparseMoeBlock, hidden_states: torch.Tensor): # check if we can apply the patch @@ -46,5 +48,28 @@ def _forward_moe(self: MixtralSparseMoeBlock, hidden_states: torch.Tensor): return final_hidden_states, router_logits -MixtralSparseMoeBlock._original_forward = MixtralSparseMoeBlock.forward -MixtralSparseMoeBlock.forward = _forward_moe +@ExportPatchRegistry.register("hf_mixtral_moe") +class MixtralMoePatch(BaseExportPatch): + """Patch for Mixtral MoE to make it compatible with torch.export. + + This patch replaces the forward method of MixtralSparseMoeBlock with + a version that uses the torch_moe custom operator for better export compatibility. + """ + + def _apply_patch(self): + """Apply the Mixtral MoE patch.""" + # Store original forward method + self.original_values["MixtralSparseMoeBlock.forward"] = MixtralSparseMoeBlock.forward + + # Apply patch by replacing the forward method + MixtralSparseMoeBlock._original_forward = MixtralSparseMoeBlock.forward # type: ignore + MixtralSparseMoeBlock.forward = _forward_moe # type: ignore + + def _revert_patch(self): + """Revert the Mixtral MoE patch.""" + # Restore original forward method + MixtralSparseMoeBlock.forward = self.original_values["MixtralSparseMoeBlock.forward"] # type: ignore + + # Clean up the temporary attribute + if hasattr(MixtralSparseMoeBlock, "_original_forward"): + delattr(MixtralSparseMoeBlock, "_original_forward") diff --git a/tensorrt_llm/_torch/auto_deploy/models/phi.py b/tensorrt_llm/_torch/auto_deploy/models/patches/phi.py similarity index 99% rename from tensorrt_llm/_torch/auto_deploy/models/phi.py rename to tensorrt_llm/_torch/auto_deploy/models/patches/phi.py index dbb97db647c9..d7bf25ecee88 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/phi.py +++ b/tensorrt_llm/_torch/auto_deploy/models/patches/phi.py @@ -173,4 +173,5 @@ def get_model_from_config_patched(config, **kwargs): return model +# TODO: figure out how this can be incorporated into the export patch system AutoModelForCausalLM.from_config = get_model_from_config_patched diff --git a/tensorrt_llm/_torch/auto_deploy/models/qwen3.py b/tensorrt_llm/_torch/auto_deploy/models/patches/qwen3.py similarity index 60% rename from tensorrt_llm/_torch/auto_deploy/models/qwen3.py rename to tensorrt_llm/_torch/auto_deploy/models/patches/qwen3.py index 5befb20cf213..3870bc5bfd84 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/qwen3.py +++ b/tensorrt_llm/_torch/auto_deploy/models/patches/qwen3.py @@ -5,6 +5,8 @@ import torch.nn.functional as F from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock +from ...export.interface import BaseExportPatch, ExportPatchRegistry + def _forward_moe(self: Qwen3MoeSparseMoeBlock, hidden_states: torch.Tensor): # check if we can apply the patch @@ -43,5 +45,28 @@ def _forward_moe(self: Qwen3MoeSparseMoeBlock, hidden_states: torch.Tensor): return final_hidden_states, router_logits -Qwen3MoeSparseMoeBlock._original_forward = Qwen3MoeSparseMoeBlock.forward -Qwen3MoeSparseMoeBlock.forward = _forward_moe +@ExportPatchRegistry.register("hf_qwen3_moe") +class Qwen3MoePatch(BaseExportPatch): + """Patch for Qwen3 MoE to make it compatible with torch.export and reduce export time. + + This patch replaces the forward method of Qwen3MoeSparseMoeBlock with + a version that uses the torch_moe custom operator for better export compatibility. + """ + + def _apply_patch(self): + """Apply the Qwen3 MoE patch.""" + # Store original forward method + self.original_values["Qwen3MoeSparseMoeBlock.forward"] = Qwen3MoeSparseMoeBlock.forward + + # Apply patch by replacing the forward method + Qwen3MoeSparseMoeBlock._original_forward = Qwen3MoeSparseMoeBlock.forward # type: ignore + Qwen3MoeSparseMoeBlock.forward = _forward_moe # type: ignore + + def _revert_patch(self): + """Revert the Qwen3 MoE patch.""" + # Restore original forward method + Qwen3MoeSparseMoeBlock.forward = self.original_values["Qwen3MoeSparseMoeBlock.forward"] # type: ignore + + # Clean up the temporary attribute + if hasattr(Qwen3MoeSparseMoeBlock, "_original_forward"): + delattr(Qwen3MoeSparseMoeBlock, "_original_forward") diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index c1a0fb151d47..7f759d6796d6 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -25,7 +25,7 @@ ) from ..custom_ops.attention_interface import SequenceInfo from ..distributed import common as dist -from ..llm_args import LlmArgs +from ..llm_args import AutoDeployConfig, LlmArgs from ..transformations.transform import InferenceOptimizer from ..utils.logger import ad_logger from .interface import CachedSequenceInterface, GetInferenceModel @@ -82,14 +82,17 @@ def _device(self) -> DeviceLikeType: return self.cache_seq_interface.device @classmethod - def build_from_config(cls, ad_config: LlmArgs): - """Build the ADEngine using the AD LlmArgs that gets passed through from the LLM.""" + def build_from_config(cls, ad_config: AutoDeployConfig): + """Build the ADEngine using the AutoDeployConfig that gets passed through from the LLM.""" max_batch_size = ad_config.max_batch_size max_seq_len = ad_config.max_seq_len attn_page_size = ad_config.attn_page_size max_num_tokens = ad_config.max_num_tokens - ad_logger.info(f"{max_seq_len=}, {max_batch_size=}, {attn_page_size=}, {max_num_tokens=}") + max_beam_width = ad_config.max_beam_width + ad_logger.info( + f"{max_seq_len=}, {max_batch_size=}, {attn_page_size=}, {max_num_tokens=}, {max_beam_width=}" + ) # initialize seq info object seq_info = SequenceInfo( @@ -111,7 +114,7 @@ def build_from_config(cls, ad_config: LlmArgs): ) # construct engine - return cls(build_and_optimize, seq_info, device) + return cls(build_and_optimize, seq_info, device, max_beam_width) @torch.inference_mode() def __init__( @@ -119,6 +122,7 @@ def __init__( get_inference_model: GetInferenceModel, seq_info: SequenceInfo, device: DeviceLikeType, + max_beam_width: int = 1, ) -> None: """Initialize the engine with model and sequence information.""" # NOTE (lucaslie): create a fake Namespace to satisfy PyExecutor requirements... @@ -131,6 +135,7 @@ def __init__( self.iter_counter = 0 # NOTE (lucaslie): not a declared base member in the base class; required by PyExecutor... + self.max_beam_width = max_beam_width self.enable_attention_dp = False # construct cache sequence interface @@ -147,19 +152,25 @@ def __init__( @nvtx_range("ad_prepare_inputs") def _prepare_inputs( - self, scheduled_requests: ScheduledRequests, resource_manager: ResourceManager - ) -> bool: + self, + scheduled_requests: ScheduledRequests, + resource_manager: ResourceManager, + new_tokens: Optional[torch.Tensor] = None, + ) -> List[bool]: """Prepare inputs for AD Model from scheduled requests.""" # cache manager kv_cache_manager = resource_manager.get_resource_manager( ResourceManagerType.KV_CACHE_MANAGER ) - # requests in order of context, extend (generate with draft), generate + # requests in order of context, generate context_requests = scheduled_requests.context_requests - extend_requests = [r for r in scheduled_requests.generation_requests if r.draft_tokens] gen_requests = [r for r in scheduled_requests.generation_requests if not r.draft_tokens] + # new_tokens is a tensor on the device, we need to convert it to a list of lists. + # can we avoid this additional gpu->cpu transfer? + new_tokens_list = new_tokens.flatten().cpu().tolist() if new_tokens is not None else None + # info to be extracted input_ids: List[List[int]] = [] input_pos: List[int] = [] @@ -172,24 +183,27 @@ def _prepare_inputs( input_ids.append(request.get_tokens(0)) input_pos.append(request.context_current_position) - # only return last logit + request.py_batch_idx = request.seq_slot last_logit_only.append(True) - # look at extend+generate requests next - for request in chain(extend_requests, gen_requests): - # store input ids and pos of first token in sequence - input_ids.append([request.get_token(0, request.get_num_tokens(0) - 1)]) - input_pos.append(request.max_beam_num_tokens - 1) + # look at generate requests next + # TODO: we should also handle extend requests (for speculative decoding) here + for request in gen_requests: + # new_tokens are provided when the overlap scheduler is enabled. + if new_tokens_list is None or request.is_dummy or request.py_batch_idx is None: + input_ids.append([request.get_token(0, request.get_num_tokens(0) - 1)]) + input_pos.append(request.max_beam_num_tokens - 1) + else: + input_ids.append([new_tokens_list[request.py_batch_idx]]) + input_pos.append(request.max_beam_num_tokens) - # check for draft tokens - if request.draft_tokens: - input_ids[-1].extend([t for t in request.draft_tokens]) + request.py_batch_idx = request.seq_slot # return all logits last_logit_only.append(False) # extract cache information for all requests - for request in chain(context_requests, extend_requests, gen_requests): + for request in chain(context_requests, gen_requests): # get cache indices cache_indices = kv_cache_manager.get_cache_indices(request) page_assignments.append(cache_indices) @@ -199,7 +213,6 @@ def _prepare_inputs( si.nest_sequences(input_ids) si.update_pos(input_pos, reset=True) si.assign_cache_loc(page_assignments) - return last_logit_only def _compute_logits(self) -> List[torch.Tensor]: @@ -224,7 +237,8 @@ def forward( ): """Run forward from scheduled requests; main entrypoint that gets called by the executor.""" # convert requests and store in sequence info object - last_logit_only = self._prepare_inputs(scheduled_requests, resource_manager) + new_tokens = getattr(new_tokens_device, "new_tokens", None) + last_logit_only = self._prepare_inputs(scheduled_requests, resource_manager, new_tokens) # compute all logits logits = self._compute_logits() @@ -286,7 +300,9 @@ def create_autodeploy_executor(executor_config: ExecutorConfig, checkpoint_dir: resource_manager.resource_managers.move_to_end(ResourceManagerType.KV_CACHE_MANAGER, last=True) # scheduling - capacitor_scheduler = BindCapacityScheduler(ad_config.max_batch_size, kv_cache_manager.impl) + capacitor_scheduler = BindCapacityScheduler( + ad_config.max_batch_size, kv_cache_manager.impl, peft_cache_manager=None + ) mb_scheduler = BindMicroBatchScheduler( ad_config.max_batch_size, engine.cache_seq_interface.info.max_num_tokens ) @@ -301,7 +317,7 @@ def create_autodeploy_executor(executor_config: ExecutorConfig, checkpoint_dir: max_seq_len=ad_config.max_seq_len, max_draft_len=max_draft_len, max_num_sequences=max_num_sequences, - max_beam_width=executor_config.max_beam_width, + max_beam_width=ad_config.max_beam_width, enable_mixed_sampler=ad_config.enable_mixed_sampler, ) sampler = TorchSampler(sampler_args) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/__init__.py b/tensorrt_llm/_torch/auto_deploy/transform/__init__.py new file mode 100644 index 000000000000..796582270437 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/__init__.py @@ -0,0 +1,4 @@ +"""AutoDeploy's modular graph transform + inference optimizer pipeline.""" + +from . import library # ensure all transforms are registered +from .interface import * diff --git a/tensorrt_llm/_torch/auto_deploy/transform/interface.py b/tensorrt_llm/_torch/auto_deploy/transform/interface.py new file mode 100644 index 000000000000..294bd0c178d1 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/interface.py @@ -0,0 +1,361 @@ +"""The interface for all transforms. + +This module defines the base classes and interfaces for all transforms. +""" + +from abc import ABC, abstractmethod +from enum import Enum +from functools import total_ordering +from typing import Any, Callable, Dict, Mapping, Tuple, Type, Union, final + +from pydantic import BaseModel, Field +from torch.fx import GraphModule + +from ..models.factory import ModelFactory +from ..shim.interface import CachedSequenceInterface +from ..transformations._graph import canonicalize_graph, lift_to_meta +from ..utils.logger import ad_logger + + +class TransformError(Exception): + """An exception raised when a transform fails.""" + + pass + + +@total_ordering +class Stages(Enum): + """Enumerated (ordered!) stages of the transformation pipeline. + + This is used to classify and pre-order transforms. + """ + + FACTORY = "factory" # factory stage for building the model + EXPORT = "export" # export stage for exporting the model to a graph module + POST_EXPORT = "post_export" # low-level cleanups of the exported graph + PATTERN_MATCHER = "pattern_matcher" # high-level pattern matching to standardize graph + SHARDING = "sharding" # auto-sharding of the graph + WEIGHT_LOAD = "weight_load" # loading of the model weights + POST_LOAD_FUSION = "post_load_fusion" # post-loading fusion and perf optimizations of the graph + CACHE_INIT = "cache_init" # initialization of cached attention + (KV) cache initialization + COMPILE = "compile" # graph compilation stage using low-level compilers like torch.compile + + def __lt__(self, other): + """Enable sorting by definition order.""" + if self.__class__ is other.__class__: + return list(self.__class__).index(self) < list(other.__class__).index(other) + return NotImplemented + + +class TransformConfig(BaseModel): + """A simple configuration class that can be extended by a transform for configurability.""" + + model_config = { + # to provide an easy way to do config validation of child config classes with more fields + "extra": "allow", + } + + ### MANDATORY CONFIG ########################################################################### + stage: Stages = Field( + description="The stage of the transformation pipeline where this transform should run.", + ) + + ### OPTIONAL CONFIG ########################################################################### + enabled: bool = Field( + default=True, + description="Whether to enable this transform.", + ) + skip_on_error: bool = Field( + default=False, + description="Whether to skip the transform if an error occurs.", + ) + + run_graph_cleanup: bool = Field( + default=True, + description="Whether to run graph cleanup/canonicalization after this transform.", + ) + run_shape_prop: bool = Field( + default=False, + description="Whether to run shape propagation after this transform.", + ) + + requires_clean_graph: bool = Field( + default=True, + description="Whether this transform requires the graph to be clean before it is applied.", + ) + requires_shape_prop: bool = Field( + default=False, + description="Whether this transform requires shape propagation before it is applied.", + ) + + +AutodeployMeta = Dict[str, Any] +_UntypedInferenceOptimizerConfig = Dict[str, Any] +StrictInferenceOptimizerConfig = Dict[str, TransformConfig] +InferenceOptimizerConfig = Mapping[str, Union[TransformConfig, _UntypedInferenceOptimizerConfig]] + + +class TransformInfo(BaseModel): + """Information about the result of a transform.""" + + model_config = { + "frozen": True, # Make the model immutable after creation + } + + skipped: bool = Field( + description="Whether the transform was skipped.", + ) + num_matches: int = Field( + description="Number of matches found.", + ) + is_clean: bool = Field( + default=False, + description="Whether the graph is clean after the transform. This can be set by the " + "transform to indicate that the transform does not change the graph and it preserves the " + "is_clean flag of the last transform.", + ) + has_valid_shapes: bool = Field( + default=False, + description="Whether meta tensor shapes are valid after the transform. This can be set by " + "the transform to indicate that the transform does not affect the shapes in the meta " + "information of the graph. In other words, the transform does not change the shapes of the " + "tensors in the graph and it preserves the has_valid_shapes flag of the last transform.", + ) + + +TransformHistory = Dict[str, TransformInfo] + + +class BaseTransform(ABC): + """A base class for all transforms.""" + + config: TransformConfig # overwrite type hint if other config cls is used in subclass! + _autodeploy_meta_key: str = "_autodeploy" + _history_key: str = "transform_history" + _transform_key: str # Set by TransformRegistry.register() decorator + + @classmethod + def get_transform_key(cls) -> str: + """Get the short name of the transform. + + This is used to identify the transform in the transformation pipeline. + """ + if hasattr(cls, "_transform_key"): + return cls._transform_key + raise NotImplementedError( + f"Transform class {cls.__name__} must be registered with TransformRegistry.register() " + "or manually implement get_transform_key()" + ) + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + """Get the configuration class for the transform. + + This is used to validate the configuration of the transform. + """ + return TransformConfig + + @final + def __init__(self, config: TransformConfig): + """Initialize the transform. + + Args: + config: The configuration for the transform, either as base config object or the actual + config object. + + To customize the initialization, override the `_post_init` method. + """ + if not isinstance(config, self.get_config_class()): + config = self.get_config_class()(**config.model_dump()) + self.config = config + self._post_init() + + def _post_init(self): + """Post-initialization hook that can be overridden by subclasses.""" + pass + + @final + @classmethod + def from_kwargs(cls, **kwargs) -> "BaseTransform": + """Create a transform from kwargs. + + Args: + **kwargs: The configuration for the transform. + + Returns: + The transform instance. + """ + config = cls.get_config_class()(**kwargs) + return cls(config=config) + + @final + def __call__( + self, gm: GraphModule, cm: CachedSequenceInterface, factory: ModelFactory + ) -> GraphModule: + """Apply the transform to the graph. + + Args: + gm: The graph module to apply the transform to. + cm: The cached sequence interface defining the sequence interface. + factory: The model factory used to build the model. + + Returns: + GraphModule: The transformed graph module. + + NOTE: The transform can/should modify the graph module in place if possible. Returning the + graph is mostly to standardize the interface for transforms that cannot modify the graph + in place (e.g. the factory or export transform). + + This method is the main entry point for any transforms and is called by the + InferenceOptimizer pipeline. + """ + + # get the transform key + t_name = self.get_transform_key() + + # retrieve autodeploy metadata from the graphmodule + autodeploy_meta = self._get_autodeploy_meta(gm) + + # retrieve transform history and last transform info + history: TransformHistory = autodeploy_meta.get(self._history_key, {}) + h_keys = list(history.keys()) # preserves order of insertion/transform execution + info_last = history[h_keys[-1]] if h_keys else TransformInfo(skipped=False, num_matches=0) + + # show debug info for debug config + ad_logger.debug(f"{t_name} config: {self.config}") + + # run or skip the transform + if self.config.enabled: + # run graph pre-cleanup + self._run_pre_cleanup(gm, info_last) + + # run the transform in a error-handling wrapper + try: + gm, info = self._apply(gm, cm, factory) + except Exception as e: + error_msg = f"Transform {t_name} failed" + if self.config.skip_on_error: + ad_logger.warning(f"{error_msg}: {e}") + info = TransformInfo(skipped=True, num_matches=0) + else: + raise TransformError(error_msg) from e + + # run graph post-cleanup + info = self._run_post_cleanup(gm, info) + else: + # skip the transform and set info object using the last transform info + info_dict = info_last.model_dump() + info_dict["skipped"] = True + info_dict["num_matches"] = 0 + info = TransformInfo(**info_dict) + + # log the result of the transform + log_msgs = [ + f"stage={self.config.stage.value}", + f"transform={t_name}", + "skipped=True" if info.skipped else f"num_matches={info.num_matches}", + f"is_clean={info.is_clean}", + f"has_valid_shapes={info.has_valid_shapes}", + ] + ad_logger.info(", ".join(log_msgs)) + ad_logger.debug(f"Graph after {t_name}: {gm}") + + # update + store new meta data + history[t_name] = info + autodeploy_meta[self._history_key] = history + self._set_autodeploy_meta(gm, autodeploy_meta) + + # return the graph module + return gm + + @final + def _get_autodeploy_meta(self, gm: GraphModule) -> AutodeployMeta: + """Get the autodeploy metadata from the graphmodule.""" + return gm.meta.get(self._autodeploy_meta_key, {}) + + @final + def _set_autodeploy_meta(self, gm: GraphModule, autodeploy_meta: AutodeployMeta) -> None: + """Set the autodeploy metadata in the graphmodule.""" + gm.meta[self._autodeploy_meta_key] = autodeploy_meta + + @final + def _run_pre_cleanup(self, gm: GraphModule, info: TransformInfo) -> None: + """Run graph cleanup before the transform. + + This is used to ensure the transform is applied to a clean graph as needed by the transform. + """ + if not self.config.requires_clean_graph: + return + + # check if run cleanup depending on the config and info + if self.config.requires_shape_prop and not (info.is_clean and info.has_valid_shapes): + with lift_to_meta(gm): + canonicalize_graph(gm, shape_prop=True) + elif self.config.requires_clean_graph and not info.is_clean: + canonicalize_graph(gm) + + @final + def _run_post_cleanup(self, gm: GraphModule, info: TransformInfo) -> TransformInfo: + """Run graph cleanup after the transform. + + Cleanup is done as requested in the config and we will update the graph module and info + accordingly. + + Returns: + Updated TransformInfo with cleanup status. + """ + if not self.config.run_graph_cleanup: + return info + + # check if run cleanup depending on the config and info + if self.config.run_shape_prop and not (info.is_clean and info.has_valid_shapes): + with lift_to_meta(gm): + canonicalize_graph(gm, shape_prop=True) + elif self.config.run_graph_cleanup and not info.is_clean: + canonicalize_graph(gm) + + # create new info object with updated cleanup status + info_dict = info.model_dump() + info_dict["is_clean"] |= self.config.run_graph_cleanup + info_dict["has_valid_shapes"] |= self.config.run_shape_prop + return TransformInfo(**info_dict) + + @abstractmethod + def _apply( + self, gm: GraphModule, cm: CachedSequenceInterface, factory: ModelFactory + ) -> Tuple[GraphModule, TransformInfo]: + """Apply the transform to the graph. + + This is the core method that should be implemented by subclasses. + """ + + +class TransformRegistry: + """A registry for all transforms.""" + + _registry: Dict[str, Type[BaseTransform]] = {} + + @classmethod + def register(cls, name: str) -> Callable[[Type[BaseTransform]], Type[BaseTransform]]: + def inner(fn: Type[BaseTransform]) -> Type[BaseTransform]: + cls._registry[name] = fn + # Auto-store the transform key as a class attribute + fn._transform_key = name + return fn + + return inner + + @classmethod + def get(cls, name: str) -> Type[BaseTransform]: + """Get the transform class by name.""" + return cls._registry[name] + + @classmethod + def get_config_class(cls, name: str) -> Type[TransformConfig]: + """Get the configuration class for a transform by name.""" + return cls.get(name).get_config_class() + + @classmethod + def has(cls, name: str) -> bool: + """Check if a transform is registered.""" + return name in cls._registry diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/__init__.py b/tensorrt_llm/_torch/auto_deploy/transform/library/__init__.py new file mode 100644 index 000000000000..403e9ee401f2 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/__init__.py @@ -0,0 +1,16 @@ +"""AutoDeploy's library of transforms. + +This file ensures that all publicly listed files/transforms in the library folder are auto-imported +and the corresponding transforms are registered. +""" + +import importlib +import pkgutil + +__all__ = [] + +for _, module_name, is_pkg in pkgutil.iter_modules(__path__): + if module_name.startswith("_"): + continue + __all__.append(module_name) + importlib.import_module(f"{__name__}.{module_name}") diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/build_model.py b/tensorrt_llm/_torch/auto_deploy/transform/library/build_model.py new file mode 100644 index 000000000000..48a8accb20b0 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/build_model.py @@ -0,0 +1,41 @@ +"""A simple wrapper transform to build a model via the model factory.""" + +from typing import Tuple, Type + +from pydantic import Field +from torch.fx import GraphModule + +from ...models.factory import ModelFactory +from ...shim.interface import CachedSequenceInterface +from ..interface import BaseTransform, TransformConfig, TransformInfo, TransformRegistry + + +class BuildModelConfig(TransformConfig): + """Configuration for the build model transform.""" + + device: str = Field(default="meta", description="The device to build the model on.") + + +@TransformRegistry.register("build_model") +class BuildModel(BaseTransform): + """A simple wrapper transform to build a model via the model factory.""" + + config: BuildModelConfig + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return BuildModelConfig + + def _apply( + self, gm: GraphModule, cm: CachedSequenceInterface, factory: ModelFactory + ) -> Tuple[GraphModule, TransformInfo]: + # build the model + model = factory.build_model(self.config.device) + + # as wrapper to satisfy the interface we will register the model as a submodule + gm.add_module("factory_model", model) + + # by convention, we say this fake graph module is always clean + info = TransformInfo(skipped=False, num_matches=1, is_clean=True, has_valid_shapes=True) + + return gm, info diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/cleanup_input_constraints.py b/tensorrt_llm/_torch/auto_deploy/transform/library/cleanup_input_constraints.py new file mode 100644 index 000000000000..1e5963505e8c --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/cleanup_input_constraints.py @@ -0,0 +1,49 @@ +import math +from typing import List, Tuple + +import torch +from torch.fx import Graph, GraphModule +from torch.utils._sympy.value_ranges import ValueRanges + +from ...models.factory import ModelFactory +from ...shim.interface import CachedSequenceInterface +from ..interface import BaseTransform, TransformInfo, TransformRegistry + + +# TODO (lucaslie): consider reconfiguring this transform to run before we switch to flattened +# sequences which is done in update_in_out_nodes at the moment. +@TransformRegistry.register("cleanup_input_constraints") +class CleanupInputConstraints(BaseTransform): + """Cleanup input constraints from the graph. + + This transformations updates the input constraints of the graph. Specifically, we want to + account for flattened sequences and hence the max constraint should be updated to reflect the + flattened sequence length. + """ + + def _apply( + self, gm: GraphModule, cm: CachedSequenceInterface, factory: ModelFactory + ) -> Tuple[GraphModule, TransformInfo]: + graph: Graph = gm.graph + input_node = graph.find_nodes(op="placeholder")[0] + sym_shape: torch.Size = input_node.meta["val"].shape + + # get expressions in the symbolic shape + vrs: List[ValueRanges] = [] + for s in sym_shape: + if isinstance(s, int): + vrs.append(ValueRanges(0, s)) + elif isinstance(s, torch.SymInt): + vrs.append(gm.range_constraints[s.node.expr]) + else: + raise TypeError(f"Unexpected type {type(s)} in symbolic shape.") + + # update the max constraint for each vr + max_total = math.prod(vr.upper for vr in vrs) + for vr in vrs: + object.__setattr__(vr, "upper", max_total) + + # store info object about the transform + info = TransformInfo(skipped=False, num_matches=len(vrs)) + + return gm, info diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/cleanup_noop_add.py b/tensorrt_llm/_torch/auto_deploy/transform/library/cleanup_noop_add.py new file mode 100644 index 000000000000..4b2abf3106b5 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/cleanup_noop_add.py @@ -0,0 +1,52 @@ +from typing import Tuple + +import torch +from torch.fx import GraphModule + +from ...models.factory import ModelFactory +from ...shim.interface import CachedSequenceInterface +from ...utils.node_utils import is_op +from ..interface import BaseTransform, TransformInfo, TransformRegistry + + +@TransformRegistry.register("cleanup_noop_add") +class CleanupNoopAdd(BaseTransform): + """Eliminate add nodes from the graph that are no-ops. + + This would be any node that is just adding 0 to the input tensor. We can safely remove those. + + NOTE: this function has one failure mode when the op ``out = tensor + zero_tensor`` is used + in such a way that``out`` will be broadcast to the shape of zero_tensor. After removing this op + then, out won't have the right shape anymore. This should be a rare case and we can handle it + when it comes up or disable this transform. + """ + + def _apply( + self, gm: GraphModule, cm: CachedSequenceInterface, factory: ModelFactory + ) -> Tuple[GraphModule, TransformInfo]: + num_matches = 0 + for node in gm.graph.nodes: + # looking for add nodes + if not is_op(node, torch.ops.aten.add): + continue + # only handling this parameter combination for now + if len(node.all_input_nodes) != 2: + continue + + # check if any of the input nodes is just a constant tensor with value 0 + if is_op(node.all_input_nodes[0], torch.ops.aten.zeros): + zero_node, true_node = node.all_input_nodes + elif is_op(node.all_input_nodes[1], torch.ops.aten.zeros): + true_node, zero_node = node.all_input_nodes + else: + continue + + # do the replacement and clean-up + node.replace_all_uses_with(true_node) + gm.graph.erase_node(node) + num_matches += 1 + + # store info object about the transform + info = TransformInfo(skipped=False, num_matches=num_matches) + + return gm, info diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/cleanup_noop_slice.py b/tensorrt_llm/_torch/auto_deploy/transform/library/cleanup_noop_slice.py new file mode 100644 index 000000000000..4b58520931af --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/cleanup_noop_slice.py @@ -0,0 +1,49 @@ +from typing import Tuple + +import torch +from torch.fx import GraphModule + +from ...models.factory import ModelFactory +from ...shim.interface import CachedSequenceInterface +from ...utils.node_utils import is_op +from ..interface import BaseTransform, TransformInfo, TransformRegistry + + +@TransformRegistry.register("cleanup_noop_slice") +class CleanupNoopSlice(BaseTransform): + """Remove no-op slice nodes from the graph. + + Those will be nodes that are used to represent a slice operation like ``t[:, :5]``. The graph IR + will represent it as ``t[:][:5]``, i.e., two nodes and the first slice being a no-op. This + function gets rid of such instances. + """ + + def _apply( + self, gm: GraphModule, cm: CachedSequenceInterface, factory: ModelFactory + ) -> Tuple[GraphModule, TransformInfo]: + num_matches = 0 + for node in gm.graph.nodes: + # looking for slice nodes + if not is_op(node, torch.ops.aten.slice): + continue + # only handling this parameter combination for now + # 4 args will be (input, dim, start, end) + if len(node.args) != 4 or len(node.kwargs) != 0: + continue + # check if dim is just an integer + if not isinstance(node.args[1], int): + continue + # check if the slice op is indeed a no-op + if node.args[2] != 0 or node.args[3] != torch.iinfo(torch.long).max: + continue + # extract input tensor node and remove the slice node + in_node = node.args[0] + assert [in_node] == node.all_input_nodes, "Slice node has unexpected input nodes." + node.replace_all_uses_with(in_node) + gm.graph.erase_node(node) + num_matches += 1 + + # store info object about the transform + info = TransformInfo(skipped=False, num_matches=num_matches) + + return gm, info diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/export_to_gm.py b/tensorrt_llm/_torch/auto_deploy/transform/library/export_to_gm.py new file mode 100644 index 000000000000..bbe72650b4e2 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/export_to_gm.py @@ -0,0 +1,71 @@ +"""A simple wrapper transform to export a model to a graph module.""" + +from typing import List, Optional, Tuple, Type + +from pydantic import Field +from torch.fx import GraphModule + +from ...export import torch_export_to_gm +from ...models.factory import ModelFactory +from ...shim.interface import CachedSequenceInterface +from ..interface import BaseTransform, TransformConfig, TransformInfo, TransformRegistry + + +class ExportToGMConfig(TransformConfig): + """Configuration for the export to graph module transform.""" + + strict: bool = Field( + description="Whether to export in strict mode. NOTE: we generally export in non-strict mode" + "for now as it relaxes some assumptions around tracing. Strict mode uses torchdynamo" + "(symbolic bytecode analysis), which can be brittle since it relies on the exact bytecode" + "representation of the model see here as well: https://pytorch.org/docs/stable/export.html#non-strict-export", + default=False, + ) + clone_state_dict: bool = Field( + description="Whether to clone the state_dict of the model. This is useful to avoid" + "modifying the original state_dict of the model.", + default=False, + ) + patch_list: Optional[List[str]] = Field( + description="List of patch names to apply with export. " + "Default is to apply all registered patches.", + default=None, + ) + + +@TransformRegistry.register("export_to_gm") +class ExportToGM(BaseTransform): + """A simple wrapper transform to export a model to a graph module.""" + + config: ExportToGMConfig + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return ExportToGMConfig + + def _apply( + self, gm: GraphModule, cm: CachedSequenceInterface, factory: ModelFactory + ) -> Tuple[GraphModule, TransformInfo]: + # at this point we assume the gm is just a dummy graph module + assert len(gm.graph.nodes) == 0, "Expected empty graph module." + + # retrieve the actual model from the dummy graph module + model = gm.get_submodule("factory_model") + + # set the example sequence + cm.info.set_example_sequence() + + # export the model to a graph module + gm = torch_export_to_gm( + model, + args=cm.args, + dynamic_shapes=cm.dynamic_shapes, + clone=self.config.clone_state_dict, + strict=self.config.strict, + patch_list=self.config.patch_list, + ) + + # this is a clean graph by definition since it was just exported + info = TransformInfo(skipped=False, num_matches=1, is_clean=True, has_valid_shapes=True) + + return gm, info diff --git a/tensorrt_llm/_torch/auto_deploy/transform/optimizer.py b/tensorrt_llm/_torch/auto_deploy/transform/optimizer.py new file mode 100644 index 000000000000..2aac699327f4 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/optimizer.py @@ -0,0 +1,76 @@ +"""High-level entrypoint to transform a model into an efficient inference model.""" + +from typing import Optional + +import torch.nn as nn +from torch.fx import Graph, GraphModule + +from ..models.factory import ModelFactory +from ..shim.interface import CachedSequenceInterface +from .interface import ( + InferenceOptimizerConfig, + Stages, + StrictInferenceOptimizerConfig, + TransformConfig, + TransformRegistry, +) + + +class InferenceOptimizer: + def __init__(self, factory: ModelFactory, config: InferenceOptimizerConfig): + self.factory = factory + self.config = self._clean_config(config) + + def _clean_config(self, config: InferenceOptimizerConfig) -> StrictInferenceOptimizerConfig: + """Get a typed checked ("strict") config with sorted keys according to stages.""" + # convert to nested kwargs, no TransformConfig objects allowed + nested_kwargs = { + k: v.model_dump() if isinstance(v, TransformConfig) else v for k, v in config.items() + } + # sort by stage + keys_sorted = sorted(nested_kwargs.keys(), key=lambda k: Stages(nested_kwargs[k]["stage"])) + # create strict config with correct config classes and correct order + strict_config: StrictInferenceOptimizerConfig = { + k: TransformRegistry.get_config_class(k)(**nested_kwargs[k]) for k in keys_sorted + } + # return strict config + return strict_config + + @staticmethod + def _init_gm() -> GraphModule: + """Initialize a fake graph module. + + This is a dummy graph module that will be used to kick off the transforms. + """ + return GraphModule(nn.Module(), Graph()) + + def __call__( + self, cm: CachedSequenceInterface, gm: Optional[GraphModule] = None + ) -> GraphModule: + """Transform a model into an optimized inference model. + + Args: + cm: The cached sequence interface defining the sequence interface. + + Returns: + A GraphModule representing the optimized inference model. + """ + ############################################################################################ + # RUN THROUGH CONFIGURED TRANSFORMATIONS + ############################################################################################ + + # start with an empty fake graph module if not provided + if gm is None: + gm = self._init_gm() + + # iterate over all transforms sorted by stage in the config + for t_name, t_config in self.config.items(): + # instantiate transform + transform = TransformRegistry.get(t_name)(t_config) + # run transform + gm = transform(gm, cm, self.factory) + + ############################################################################################ + # RETURN OPTIMIZED GRAPH + ############################################################################################ + return gm diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/__init__.py b/tensorrt_llm/_torch/auto_deploy/transformations/__init__.py index e69de29bb2d1..d643d8bb0b60 100644 --- a/tensorrt_llm/_torch/auto_deploy/transformations/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/transformations/__init__.py @@ -0,0 +1 @@ +"""V1 Graph Transformations Module --> will be deprecated and replaced by auto_deploy.transform.""" diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/_graph.py b/tensorrt_llm/_torch/auto_deploy/transformations/_graph.py index 5b33a3816e84..5e92764079f5 100644 --- a/tensorrt_llm/_torch/auto_deploy/transformations/_graph.py +++ b/tensorrt_llm/_torch/auto_deploy/transformations/_graph.py @@ -59,7 +59,7 @@ def load_buffers_and_params( if clone: v_new = v.detach().clone() if isinstance(v, torch.nn.Parameter): - v_new = nn.Parameter(v_new) + v_new = nn.Parameter(v_new, requires_grad=False) else: v_new = state_dict[k] setattr(submod, name, v_new) @@ -192,7 +192,7 @@ def _canonicalize_single_gm( def canonicalize_graph( gm: GraphModule, shape_prop: bool = False, args_static: Optional[Tuple[Any, ...]] = None -) -> GraphModule: +) -> None: """Canonicalize the graph of the given GraphModule. Args: @@ -217,8 +217,6 @@ def canonicalize_graph( ad_logger.debug(f"After canonicalizing: {gm}") - return gm - def add_graph_input( gm: GraphModule, name: str, val: Optional[torch.Tensor] = None, dynamic_shape=None diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/export.py b/tensorrt_llm/_torch/auto_deploy/transformations/export.py deleted file mode 100644 index 495b3593ecc7..000000000000 --- a/tensorrt_llm/_torch/auto_deploy/transformations/export.py +++ /dev/null @@ -1,488 +0,0 @@ -import importlib.metadata -import math -from collections import defaultdict -from contextlib import contextmanager, nullcontext -from functools import partial -from typing import Any, Dict, List, Optional, Tuple - -import torch -import torch.export as te -import torch.nn as nn -import torch.nn.functional as F -from packaging import version -from torch import fx -from torch.utils._sympy.value_ranges import ValueRanges - -from ..utils.logger import ad_logger -from ..utils.node_utils import is_op -from ._graph import canonicalize_graph, lift_to_meta, load_buffers_and_params, tree_to - -try: - from modelopt.torch.quantization.utils import export_torch_mode as torch_export_context -except ImportError: - torch_export_context = nullcontext - - -def _clean_up_no_op_slice_nodes(gm: fx.GraphModule): - """Remove no-op slice nodes from the graph. - - Those will be nodes that are used to represent a slice operation like ``t[:, :5]``. The graph IR - will represent it as ``t[:][:5]``, i.e., two nodes and the first slice being a no-op. This - function gets rid of such instances. - """ - for node in gm.graph.nodes: - # looking for slice nodes - if not is_op(node, torch.ops.aten.slice): - continue - # only handling this parameter combination for now - # 4 args will be (input, dim, start, end) - if len(node.args) != 4 or len(node.kwargs) != 0: - continue - # check if dim is just an integer - if not isinstance(node.args[1], int): - continue - # check if the slice op is indeed a no-op - if node.args[2] != 0 or node.args[3] != torch.iinfo(torch.long).max: - continue - # extract input tensor node and remove the slice node - in_node = node.args[0] - assert [in_node] == node.all_input_nodes, "Slice node has unexpected input nodes." - node.replace_all_uses_with(in_node) - gm.graph.erase_node(node) - - canonicalize_graph(gm) - - -def _eliminate_no_op_add_nodes(gm: fx.GraphModule): - """Eliminate add nodes from the graph that are no-ops. - - This would be any node that is just adding 0 to the input tensor. We can safely remove those. - - NOTE: this function has one failure mode when the op ``out = tensor + zero_tensor`` is used - in such a way that``out`` will be broadcast to the shape of zero_tensor. After removing this op - then, out won't have the right shape anymore. This should e a rare case and we can handle it - when it comes up. - """ - for node in gm.graph.nodes: - # looking for add nodes - if not is_op(node, torch.ops.aten.add): - continue - # only handling this parameter combination for now - if len(node.all_input_nodes) != 2: - continue - - # check if any of the input nodes is just a constant tensor with value 0 - if is_op(node.all_input_nodes[0], torch.ops.aten.zeros): - zero_node, true_node = node.all_input_nodes - elif is_op(node.all_input_nodes[1], torch.ops.aten.zeros): - true_node, zero_node = node.all_input_nodes - else: - continue - - # do the replacement and clean-up - node.replace_all_uses_with(true_node) - gm.graph.erase_node(node) - - canonicalize_graph(gm) - - -def _clean_up_device_info(gm: fx.GraphModule): - """Correct device information in the graph.""" - devices = {t.device for _, t in gm.named_parameters()} - if len(devices) == 0: - return - elif len(devices) > 1: - raise AssertionError("All parameters should be on the same device.") - device = devices.pop() - meta_device = torch.device("meta") - - for node in gm.graph.nodes: - if any(a == meta_device for a in node.args): - new_args = list(node.args) - new_args = [a if a != meta_device else device for a in new_args] - node.args = tuple(new_args) - if any(a == meta_device for a in node.kwargs.values()): - new_kwargs = dict(node.kwargs) - new_kwargs = {k: v if v != meta_device else device for k, v in new_kwargs.items()} - node.kwargs = new_kwargs - - canonicalize_graph(gm) - - -def _load_hook_for_deduplication( - state_dict, prefix, *args, param_key_remaining: str, param_key_removed: str -): - """Check for removed param key and and put it into the key that is remaining.""" - ad_logger.debug(f"Loading hook for deduplication: {param_key_remaining} <- {param_key_removed}") - k_remaining = prefix + param_key_remaining - k_removed = prefix + param_key_removed - if k_removed in state_dict: - state_dict[k_remaining] = state_dict.pop(k_removed) - - -def _deduplicate_params_and_buffers(gm: fx.GraphModule): - """This will de-duplicate params and buffers that share the same tensor.""" - # get all get_attr nodes - get_attr_nodes = [n for n in gm.graph.nodes if n.op == "get_attr"] - - # sort by id of target - targets: Dict[int, List[fx.Node]] = defaultdict(list) - for n in get_attr_nodes: - submod, _, name = n.target.rpartition(".") - t_target = getattr(gm.get_submodule(submod), name) - targets[id(t_target)].append(n) - # now replace all instances of the same tensor with the same get_attr node (idx 0 in the list) - for nodes in targets.values(): - node_kept = nodes[0] - for n in nodes[1:]: - n.replace_all_uses_with(node_kept) - gm.graph.erase_node(n) - - # remove the param/buffer from the submodule - submod, _, name = n.target.rpartition(".") - delattr(gm.get_submodule(submod), name) - - # add load hooks to also load the weights correctly - gm._register_load_state_dict_pre_hook( - partial( - _load_hook_for_deduplication, - param_key_remaining=node_kept.target, - param_key_removed=n.target, - ) - ) - - ad_logger.debug(f"Deduplicated: {n.target} --> {node_kept.target}") - - canonicalize_graph(gm) - - -def _clean_up_checks(gm: fx.GraphModule): - """This transformations removes shape checks and assertions from the graph.""" - check_ops = { - torch.ops.aten._assert_scalar, - torch.ops.aten.sym_constrain_range, - torch.ops.aten.sym_constrain_range_for_size, - torch.ops.aten._assert_tensor_metadata, - # torch.ops.aten._functional_sym_constrain_range, - # torch.ops.aten._functional_sym_constrain_range_for_size - } - graph: fx.Graph = gm.graph - for node in reversed(graph.nodes): - if len(node.users) > 0 or not is_op(node, check_ops): - continue - graph.erase_node(node) - canonicalize_graph(gm) - - -def _clean_up_input_constraints(gm: fx.GraphModule): - """This transformations updates the input constraints of the graph. - - Specifically, we want to account for flattened sequences and hence the max constraint should - be updated to reflect the flattened sequence length. - """ - graph: fx.Graph = gm.graph - input_node = graph.find_nodes(op="placeholder")[0] - sym_shape: torch.Size = input_node.meta["val"].shape - - # get expressions in the symbolic shape - vrs: List[ValueRanges] = [] - for s in sym_shape: - if isinstance(s, int): - vrs.append(ValueRanges(0, s)) - elif isinstance(s, torch.SymInt): - vrs.append(gm.range_constraints[s.node.expr]) - else: - raise TypeError(f"Unexpected type {type(s)} in symbolic shape.") - - # update the max constraint for each vr - max_total = math.prod(vr.upper for vr in vrs) - for vr in vrs: - object.__setattr__(vr, "upper", max_total) - - canonicalize_graph(gm) - - -# TODO: remove once https://github.com/pytorch/pytorch/issues/140710 is resolved -def _torch_where_patch(condition: torch.Tensor, *args, **kwargs): - if len(args) == 0 and len(kwargs) == 0: - return torch.nonzero(condition, as_tuple=True) - return _torch_where_patch.where_original(condition, *args, **kwargs) - - -_torch_where_patch.where_original = torch.where - - -def _torch_linear_patch( - input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None -) -> torch.Tensor: - return torch.ops.auto_deploy.torch_linear_simple(input, weight, bias) - - -# TODO: remove once https://github.com/pytorch/pytorch/issues/142439 is resolved -def _torch_modulelist_getitem_patch(self: nn.ModuleList, idx): - if isinstance(idx, slice): - # return a simple list. - # NOTE: this obviously only works for any use case where we access the sliced module list - # like a regular list like a for-loop. For most other things, this hack will not work. - return list(self._modules.values())[idx] - else: - return _torch_modulelist_getitem_patch.getitem_original(self, idx) - - -_torch_modulelist_getitem_patch.getitem_original = nn.ModuleList.__getitem__ - - -def _torch_tensor_patch(data, **kwargs): - """Patch torch.tensor to handle 0.0 on meta device. - - ``torch.tensor(0.0, device="meta")`` does not work and hence we are patching it to use - ``torch.zeros((), device="meta")`` instead, which is equivalent. - """ - device = kwargs.get("device", None) - if data == 0.0 and device is not None and torch.device(device) == torch.device("meta"): - return torch.zeros((), **kwargs) - return _torch_tensor_patch.tensor_original(data, **kwargs) - - -_torch_tensor_patch.tensor_original = torch.tensor - - -def _transformers_version() -> str: - """Get the version of transformers.""" - return version.parse(importlib.metadata.version("transformers")).base_version - - -# TODO (@lucaslie): https://github.com/NVIDIA/TensorRT-LLM/issues/5728 -# not great that this patch is here but it's the least invasisve change until we make headway on the -# above issue. -@contextmanager -def _transformers_sdpa_mask_patch(): - """Patch transformers.masking_utils.sdpa_mask to be export-compatible.""" - # this patch is only needed+compatible for transformers >= 4.53.0 - if version.parse(_transformers_version()) < version.parse("4.53.0"): - yield # Just yield without doing anything (like nullcontext) - return - - # imports only after version check - from transformers import masking_utils - from transformers.integrations.executorch import sdpa_mask_without_vmap - - # recall original implementation - sdpa_mask_original = masking_utils.sdpa_mask - - # patch function and mask attention interface - masking_utils.sdpa_mask = sdpa_mask_without_vmap - if "sdpa" in masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._local_mapping: - sdpa_local_original = masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._local_mapping["sdpa"] - else: - sdpa_local_original = None - masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"] = sdpa_mask_without_vmap - - try: - yield - finally: - # revert patches - masking_utils.sdpa_mask = sdpa_mask_original - if sdpa_local_original is None: - del masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"] - else: - masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"] = sdpa_local_original - - -def add_missing_load_hooks(gm: fx.GraphModule, model: nn.Module) -> fx.GraphModule: - """Adds back the state dict load hooks stripped away during export.""" - hooks = { - k: mod._load_state_dict_pre_hooks - for k, mod in model.named_modules() - if mod._load_state_dict_pre_hooks - } - - for mod_name, mod in gm.named_modules(): - if mod_name in hooks: - for hook in hooks.pop(mod_name).values(): - mod._register_load_state_dict_pre_hook(hook.hook, with_module=hook.with_module) - assert not (bool(hooks)), f"""Mismatch in names of exported and source modules with hooks. - The following module names were not found in exported module {list(hooks.keys())}""" - - return gm - - -def add_load_hook_for_aliased_params(gm: fx.GraphModule, model: nn.Module): - """ - Add a load hook to handle aliased parameters in the model. - - When parameters are aliased (multiple parameter names point to the same tensor), - we need to ensure all aliases get the same value during loading. This hook: - 1. Identifies groups of aliased parameters - 2. For each group, finds a valid parameter value from the state dict - 3. Applies that value to all aliases in the group - - Args: - gm: The graph module to add the hook to - model: The source model containing the original parameter aliases - """ - # Find all parameter aliases in the source model - param_to_names = defaultdict(list) - for name, param in model.named_parameters(remove_duplicate=False): - param_to_names[id(param)].append(name) - - # Filter to only groups with multiple aliases - aliased_groups = [names for names in param_to_names.values() if len(names) > 1] - - if not aliased_groups: - return gm # No aliases to handle - - def find_valid_param_value( - state_dict: Dict[str, torch.Tensor], param_names: List[str] - ) -> Optional[torch.Tensor]: - """Find a valid parameter value from state dict for a group of aliased parameters. - - Args: - state_dict: The state dict being loaded - param_names: List of parameter names that are aliases of each other - - Returns: - A valid tensor value if found, None otherwise - """ - # First try to find a non-meta tensor value - value = None - for name in param_names: - if name in state_dict: - value = state_dict[name] - if value.device.type != "meta": - return value - - return value - - def aliasing_load_pre_hook(state_dict: Dict[str, torch.Tensor], prefix: str, *args, **kwargs): - """Load hook that ensures aliased parameters get the same value.""" - for group in aliased_groups: - # Find a valid value for this group of aliases - value = find_valid_param_value(state_dict, group) - assert value is not None, ( - f"No valid value found in state dict for aliased parameters: {group}" - ) - - # Apply the value to all aliases - for name in group: - state_dict[name] = value - - ad_logger.debug(f"Applied value from {group[0]} to aliased parameters: {group}") - - # Register the hook - gm._register_load_state_dict_pre_hook(aliasing_load_pre_hook) - - -@torch.inference_mode() -def torch_export(model: nn.Module, *export_args, **export_kwargs) -> te.ExportedProgram: - """Just like torch.export except we decorate it to be in inference_mode.""" - with torch_export_context(): - ep = te.export(model, *export_args, **export_kwargs) - - # return the result - return ep - - -def torch_export_to_gm( - model: nn.Module, - args: Tuple[Any, ...], - kwargs: Optional[Dict[str, Any]] = None, - clone: bool = False, # clone or don't clone the model state_dict - **export_kwargs, -) -> fx.GraphModule: - """torch_export with wrapping into GraphModule + useful additions to the resulting module.""" - # we need to better control how F.scaled_dot_product_attention is represented in the graph - # there is no guarantee how it is represented and we need to make sure it is easily identifiable - # in the graph. - sdpa_original = F.scaled_dot_product_attention - F.scaled_dot_product_attention = torch.ops.auto_deploy.torch_attention_sdpa - - # We overwrite the linear functional as well. This basically avoids exporting the view ops - # that are used to flatten/unflatten multiple batch dimensions of the input tensor. - linear_original = F.linear - # patch linear → always supply bias - F.linear = _torch_linear_patch - - # patch torch.where(condition) to torch.nonzero(condition, as_tuple=True) - torch.where = _torch_where_patch - - # patch nn.ModuleList.__getitem__ to handle slicing - nn.ModuleList.__getitem__ = _torch_modulelist_getitem_patch - - # overwrite autocast/sdpa contextmanagers to be no-ops - autocast_original = torch.autocast - sdpa_kernel_original = torch.nn.attention.sdpa_kernel - torch.autocast = lambda *args, **kwargs: nullcontext() - torch.nn.attention.sdpa_kernel = lambda *args, **kwargs: nullcontext() - - # patch torch.tensor to handle 0.0 on meta device - torch.tensor = _torch_tensor_patch - - # run export with sdpa masking patch and lifted to meta - with _transformers_sdpa_mask_patch(): - with lift_to_meta(model) as state_dict: - # clean up args, kwargs and move to correct device - args, kwargs = tree_to((args, kwargs or {}), device="meta") - - # NOTE: we always export in non-strict mode for now as it relaxes some - # assumptions around tracing. Strict mode uses torchdynamo (symbolic bytecode analysis), - # which can be brittle since it relies on the exact bytecode representation of the model - # see here as well: https://pytorch.org/docs/stable/export.html#non-strict-export - export_kwargs["strict"] = False - - # run export and extract graph module - egm: fx.GraphModule = torch_export(model, args, kwargs, **export_kwargs).module() - - # load state_dict into egm - # NOTE: export might have removed unused params/buffers (hence we allow unexpected keys) - load_buffers_and_params( - egm, state_dict, strict_missing=True, strict_unexpected=False, clone=clone - ) - - # revert sdpa back to original - F.scaled_dot_product_attention = sdpa_original - - # revert linear back to original - F.linear = linear_original - - # revert torch.where patch - torch.where = _torch_where_patch.where_original - - # revert nn.ModuleList.__getitem__ patch - nn.ModuleList.__getitem__ = _torch_modulelist_getitem_patch.getitem_original - - # revert autocast/sdpa back to original - torch.autocast = autocast_original - torch.nn.attention.sdpa_kernel = sdpa_kernel_original - - # revert torch.tensor patch - torch.tensor = _torch_tensor_patch.tensor_original - - # Export strips away all methods not traced during forward. The model could have - # load hooks that contain logic for correct state_dict loading. We need to add those - # hooks back to the exported graph module. - add_missing_load_hooks(egm, model) - - # Export will have LOTS of no-op slice nodes. Let's remove them to clean up the graph - # representation - _clean_up_no_op_slice_nodes(egm) - - # Export does not clean "no-op" element-wise add nodes. We can safely remove those. - _eliminate_no_op_add_nodes(egm) - - # clean up devices in the graph - _clean_up_device_info(egm) - - # Add load hook to correctly load parameters that are aliased in the source model. - add_load_hook_for_aliased_params(egm, model) - - # deduplicate params and buffers - _deduplicate_params_and_buffers(egm) - - # clean up shape checks and assertions - _clean_up_checks(egm) - - # clean up input constraints - _clean_up_input_constraints(egm) - - return egm diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/library/__init__.py b/tensorrt_llm/_torch/auto_deploy/transformations/library/__init__.py index 379f7d2b30c4..7662a3d58395 100644 --- a/tensorrt_llm/_torch/auto_deploy/transformations/library/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/transformations/library/__init__.py @@ -3,11 +3,12 @@ from .attention import * from .collectives import * from .eliminate_redundant_transposes import * -from .ep_sharding import * from .fused_moe import * from .fusion import * from .kvcache import * from .quantization import * +from .quantize_moe import * +from .rms_norm import * from .rope import * from .sharding import * diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/library/attention.py b/tensorrt_llm/_torch/auto_deploy/transformations/library/attention.py index 7e46bd652ce1..e6efb8e0e7fb 100644 --- a/tensorrt_llm/_torch/auto_deploy/transformations/library/attention.py +++ b/tensorrt_llm/_torch/auto_deploy/transformations/library/attention.py @@ -11,7 +11,7 @@ from .._graph import canonicalize_graph -def match_repeat_kv(gm: GraphModule) -> GraphModule: +def match_repeat_kv(gm: GraphModule) -> None: """ Match and replace the repeat_kv pattern in fx graphs. @@ -36,13 +36,11 @@ def match_repeat_kv(gm: GraphModule) -> GraphModule: # Clean up the graph if we made any replacements if num_kv_patterns: - gm = canonicalize_graph(gm) + canonicalize_graph(gm) ad_logger.info(f"Found {num_kv_patterns} repeat_kv patterns") - return gm - -def match_eager_attention(gm: GraphModule) -> GraphModule: +def match_eager_attention(gm: GraphModule) -> None: """ Match and replace the eager attention pattern in fx graphs. @@ -68,12 +66,11 @@ def match_eager_attention(gm: GraphModule) -> GraphModule: # Clean up the graph if we made any replacements if num_eager_patterns: - gm = canonicalize_graph(gm) + canonicalize_graph(gm) ad_logger.info(f"Found {num_eager_patterns} eager attention patterns") - return gm -def match_grouped_attention(gm: GraphModule) -> GraphModule: +def match_grouped_attention(gm: GraphModule) -> None: """ Match and replace the grouped attention pattern in fx graphs. @@ -101,12 +98,11 @@ def match_grouped_attention(gm: GraphModule) -> GraphModule: # Clean up the graph if we made any replacements if num_grouped_patterns: - gm = canonicalize_graph(gm) + canonicalize_graph(gm) ad_logger.info(f"Found {num_grouped_patterns} grouped attention patterns") - return gm -def match_causal_attn_mask(gm: GraphModule) -> GraphModule: +def match_causal_attn_mask(gm: GraphModule) -> None: """ Match attention operations with causal attention masks and optimize them. @@ -174,9 +170,8 @@ def match_causal_attn_mask(gm: GraphModule) -> GraphModule: # Clean up the graph if we made any replacements if num_causal_patterns: - gm = canonicalize_graph(gm) + canonicalize_graph(gm) ad_logger.info(f"Found {num_causal_patterns} causal mask attention patterns") - return gm def _match_repeat_kv_pattern(reshape_node: Node) -> Optional[Dict[str, Node]]: @@ -748,7 +743,7 @@ def _has_triu_ancestor(node: Node, offset: int = 1, depth: int = 0, max_depth: i return False -def match_attention_layout(gm: GraphModule, attention_op: Type[AttentionDescriptor]) -> GraphModule: +def match_attention_layout(gm: GraphModule, attention_op: Type[AttentionDescriptor]) -> None: """ Match and transform attention operations to match the layout expected by the attention backend. @@ -832,9 +827,7 @@ def match_attention_layout(gm: GraphModule, attention_op: Type[AttentionDescript # Clean up the graph if we made any replacements if num_bsnd_patterns: - gm = canonicalize_graph(gm) + canonicalize_graph(gm) ad_logger.debug(f"Transformed graph for bsnd layout: {gm}") ad_logger.info(f"Found and matched {num_bsnd_patterns} attention layouts") - - return gm diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/library/collectives.py b/tensorrt_llm/_torch/auto_deploy/transformations/library/collectives.py index bf6f804c4273..8cec047561f9 100644 --- a/tensorrt_llm/_torch/auto_deploy/transformations/library/collectives.py +++ b/tensorrt_llm/_torch/auto_deploy/transformations/library/collectives.py @@ -15,7 +15,7 @@ # * version above with fused GEMMs (i.e. with a split node) # * all_reduce(pointwise_op(linear(x))) # * ... -def fuse_collectives(gm: GraphModule) -> GraphModule: +def fuse_collectives(gm: GraphModule) -> None: num_gemm_collective_fusions = 0 ad_logger.debug("Before GEMM+Collective fusion: " + str(gm)) @@ -54,13 +54,12 @@ def fuse_collectives(gm: GraphModule) -> GraphModule: gm.graph.erase_node(parent_node) num_gemm_collective_fusions += 1 - gm = canonicalize_graph(gm) + canonicalize_graph(gm) ad_logger.info(f"Found {num_gemm_collective_fusions} GEMM+Collective fusions") ad_logger.debug("After GEMM+Collective fusion: " + str(gm)) - return gm -def fuse_allreduce_residual_rmsnorm(gm: GraphModule) -> GraphModule: +def fuse_allreduce_residual_rmsnorm(gm: GraphModule) -> None: """Essentially, this function fuses the following operators into one allreduce trtllm implementation. * target pattern: @@ -72,7 +71,7 @@ def fuse_allreduce_residual_rmsnorm(gm: GraphModule) -> GraphModule: """ if not is_trtllm_op_available(): - return gm + return num_ar_r_rms_fusions = 0 ad_logger.debug("Before allreduce+residual+rmsnorm fusion: " + str(gm)) @@ -158,14 +157,11 @@ def trace_and_fuse(allreduce_node, graph): nonlocal num_ar_r_rms_fusions num_ar_r_rms_fusions += 1 - return - # Traverse all nodes for node in gm.graph.nodes: if is_op(node, torch.ops.auto_deploy.torch_dist_all_reduce): trace_and_fuse(allreduce_node=node, graph=gm.graph) - gm = canonicalize_graph(gm) + canonicalize_graph(gm) ad_logger.info(f"Found {num_ar_r_rms_fusions} allreduce+residual+rmsnorm fusions") ad_logger.debug("After allreduce+residual+rmsnorm fusion: " + str(gm)) - return gm diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/library/eliminate_redundant_transposes.py b/tensorrt_llm/_torch/auto_deploy/transformations/library/eliminate_redundant_transposes.py index 5433afdbae01..a8c6668dde5a 100644 --- a/tensorrt_llm/_torch/auto_deploy/transformations/library/eliminate_redundant_transposes.py +++ b/tensorrt_llm/_torch/auto_deploy/transformations/library/eliminate_redundant_transposes.py @@ -40,7 +40,7 @@ def _are_transpose_args_same(node1: Node, node2: Node) -> bool: return dim1_node1 == dim1_node2 and dim2_node1 == dim2_node2 -def eliminate_redundant_transposes(gm: GraphModule) -> GraphModule: +def eliminate_redundant_transposes(gm: GraphModule) -> None: """Eliminate redundant transpose operations in the graph. This transformation identifies pairs of consecutive transpose operations with @@ -107,7 +107,6 @@ def eliminate_redundant_transposes(gm: GraphModule) -> GraphModule: # Clean up the graph if nodes_to_eliminate: gm.graph.eliminate_dead_code() - gm = canonicalize_graph(gm) + canonicalize_graph(gm) ad_logger.info(f"Found and eliminated {len(nodes_to_eliminate)} redundant transpose pairs") ad_logger.debug("After eliminating redundant transposes: " + str(gm)) - return gm diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/library/ep_sharding.py b/tensorrt_llm/_torch/auto_deploy/transformations/library/ep_sharding.py deleted file mode 100644 index acae157a6b7d..000000000000 --- a/tensorrt_llm/_torch/auto_deploy/transformations/library/ep_sharding.py +++ /dev/null @@ -1,130 +0,0 @@ -""" -Expert Parallel Sharding for Mixture-of-Experts (MoE) Graphs. - -This module implements graph transformations to enable expert sharding -for Mixture-of-Experts (MoE) models in a multi-GPU setting. The sharding -algorithm partitions the expert weights, as well as updates the routing -components (`selected_experts` and `final_scales`), so that each GPU only -processes a subset of experts. - -The sharding process consists of: - -1. Identify MoE nodes in the FX graph -2. Compute local sharding parameters (`selected_experts` and `final_scales`) to update the routing tensors. -3. Partition expert weight lists according to the current rank and world size, - and replace the MoE node’s arguments with these sharded versions. -4. Append an all_reduce node after each MoE node to aggregate outputs across devices, - then canonicalize the modified graph. - -""" - -import operator - -import torch -from torch.fx import GraphModule, Node - -from ...utils.logger import ad_logger -from ...utils.node_utils import is_op -from .._graph import canonicalize_graph - - -def ep_shard(gm: GraphModule, rank: int, world_size: int) -> GraphModule: - ad_logger.debug("Before sharding graph: " + str(gm)) - - if world_size < 2: - ad_logger.info("Skipping sharding for single device") - return gm - - assert isinstance(gm, GraphModule), "Expecting GraphModule" - num_moe_patterns = 0 - for node in list(gm.graph.nodes): - if not is_op(node, torch.ops.auto_deploy.torch_moe): - continue - _insert_sharded_moe(gm, node, rank, world_size) - num_moe_patterns += 1 - # canonicalize and return - gm = canonicalize_graph(gm) - - ad_logger.debug("After sharding: " + str(gm)) - ad_logger.info(f"Found {num_moe_patterns} MoE patterns") - return gm - - -def _insert_sharded_moe( - gm: GraphModule, - node: Node, - rank: int, - world_size: int, -): - """Update the torch_moe node with sharded weight lists, - sharded `selected_experts` and `final_scales(router_logics)`. - Add an all_reduce node after the moe node. - """ - num_experts = len(node.args[3]) - args = list(node.args) - - # -- Handle selected_experts and final_scales sharding -- - selected_experts = args[1] - final_scales = args[2] - - experts_per_rank = num_experts // world_size - - with gm.graph.inserting_before(node): - lower = experts_per_rank * rank - # selected_experts_local = selected_experts - low - selected_experts_local = gm.graph.create_node( - "call_function", operator.sub, args=(selected_experts, lower), kwargs={} - ) - - # For num_experts % world_size != 0 case, - # assign the last (num_experts % world_size) experts to the last rank - # if rank == world_size -1: - # rank_mask = (selected_experts // experts_per_rank) >= rank - # else: - # rank_mask = (selected_experts // experts_per_rank) == rank - div_node = gm.graph.create_node( - "call_function", operator.floordiv, args=(selected_experts, experts_per_rank), kwargs={} - ) - comp_op = torch.ge if rank == world_size - 1 else torch.eq - rank_mask = gm.graph.create_node("call_function", comp_op, args=(div_node, rank), kwargs={}) - - # final_scales_local = final_scales * rank_mask - final_scales_local = gm.graph.create_node( - "call_function", operator.mul, args=(final_scales, rank_mask), kwargs={} - ) - - # -- Shard expert weights -- - def get_partition(lst, world_size, rank): - num_experts = len(lst) - expert_size_per_partition = num_experts // world_size - expert_start = rank * expert_size_per_partition - # For num_experts % world_size != 0 case, - # assign the last (num_experts % world_size) experts to the last rank - expert_end = ( - num_experts if (rank == world_size - 1) else expert_start + expert_size_per_partition - ) - return lst[expert_start:expert_end] - - w1_list_sharded = get_partition(args[3], world_size, rank) - w2_list_sharded = get_partition(args[4], world_size, rank) - w3_list_sharded = get_partition(args[5], world_size, rank) - - # -- Update args -- - args[1] = selected_experts_local - args[2] = final_scales_local - args[3] = w1_list_sharded - args[4] = w2_list_sharded - args[5] = w3_list_sharded - - ad_logger.debug( - f"Updated node {node}: replaced original arguments {node.args} with sharded arguments {args}." - ) - node.args = tuple(args) - - # -- add an all_reduce node -- - with gm.graph.inserting_after(node): - dist_node = gm.graph.call_function( - torch.ops.auto_deploy.torch_dist_all_reduce, args=(node,) - ) - node.replace_all_uses_with(dist_node) - dist_node.replace_input_with(dist_node, node) diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/library/fused_moe.py b/tensorrt_llm/_torch/auto_deploy/transformations/library/fused_moe.py index 02e3e64e1704..e04997086223 100644 --- a/tensorrt_llm/_torch/auto_deploy/transformations/library/fused_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/transformations/library/fused_moe.py @@ -7,10 +7,11 @@ from ...utils.cuda_mem_tracker import cuda_memory_tracker from ...utils.logger import ad_logger from ...utils.node_utils import bfs, identify_regions_between_residuals, is_linear_op, is_op +from ...utils.quantization_utils import get_scales_and_type_from_node from .._graph import canonicalize_graph -def match_moe_pattern(gm: GraphModule) -> GraphModule: +def match_moe_pattern(gm: GraphModule) -> None: graph = gm.graph ad_logger.debug("Before MoE Pattern Matching: " + str(gm)) @@ -21,8 +22,8 @@ def match_moe_pattern(gm: GraphModule) -> GraphModule: for start_boundary, end_boundary in zip(boundary_nodes[:-1], boundary_nodes[1:]): # Step 1: Identify Expert Compute pattern - pattern_input_nodes, pattern_output_nodes, expert_weights = _match_expert_compute_pattern( - start_boundary, end_boundary + (pattern_input_nodes, pattern_output_nodes, expert_weights, expert_scales, weight_type) = ( + _match_expert_compute_pattern(start_boundary, end_boundary) ) if not expert_weights: continue @@ -56,29 +57,70 @@ def match_moe_pattern(gm: GraphModule) -> GraphModule: if final_hidden_state_node is None: continue - # Step 5: Insert the moe op into the graph. + # Step 5: Insert the MoE op into the graph. ad_logger.debug( - f"""Found MoE Pattern: between boundary {start_boundary} and {end_boundary}.\n - Capturing input hidden states node: {hidden_states}, - selected_experts node: {selected_experts}, routing_weights node: {normalized_routing_weights}, - expert weights : {expert_weights} """ + f"Found MoE Pattern: between boundary {start_boundary} and {end_boundary}.\n" + f"Input hidden states node: {hidden_states}, " + f"selected_experts node: {selected_experts}, " + f"routing_weights node: {normalized_routing_weights}, " + f"expert weights: {expert_weights}, weight type: {weight_type}" ) with graph.inserting_before(final_hidden_state_node): w1_list = expert_weights["w1"] w2_list = expert_weights["w2"] w3_list = expert_weights["w3"] - fused_moe_node = graph.call_function( - torch.ops.auto_deploy.torch_moe, - args=( - hidden_states, - selected_experts, - normalized_routing_weights, - w1_list, - w2_list, - w3_list, - ), - ) + if weight_type == "fp8": + fused_moe_node = graph.call_function( + torch.ops.auto_deploy.torch_quant_fp8_moe, + args=( + hidden_states, + selected_experts, + normalized_routing_weights, + w1_list, + w2_list, + w3_list, + expert_scales["w1_input_scale"], + expert_scales["w2_input_scale"], + expert_scales["w3_input_scale"], + expert_scales["w1_weight_scale"], + expert_scales["w2_weight_scale"], + expert_scales["w3_weight_scale"], + ), + ) + elif weight_type == "fp4": + fused_moe_node = graph.call_function( + torch.ops.auto_deploy.torch_quant_fp4_moe, + args=( + hidden_states, + selected_experts, + normalized_routing_weights, + w1_list, + w2_list, + w3_list, + expert_scales["w1_input_scale"], + expert_scales["w2_input_scale"], + expert_scales["w3_input_scale"], + expert_scales["w1_weight_scale"], + expert_scales["w2_weight_scale"], + expert_scales["w3_weight_scale"], + expert_scales["w1_alpha"], + expert_scales["w2_alpha"], + expert_scales["w3_alpha"], + ), + ) + else: + fused_moe_node = graph.call_function( + torch.ops.auto_deploy.torch_moe, + args=( + hidden_states, + selected_experts, + normalized_routing_weights, + w1_list, + w2_list, + w3_list, + ), + ) final_hidden_state_node.replace_all_uses_with(fused_moe_node) graph.erase_node(final_hidden_state_node) @@ -88,17 +130,15 @@ def match_moe_pattern(gm: GraphModule) -> GraphModule: num_moe_patterns += 1 - gm = canonicalize_graph(gm) + canonicalize_graph(gm) ad_logger.info(f"Found {num_moe_patterns} MoE Patterns") ad_logger.debug("After MoE Pattern Matching: " + str(gm)) - return gm - -def fuse_moe(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: +def fuse_moe(gm: torch.fx.GraphModule) -> None: """ - Scan the FX graph and replace all calls to torch.ops.moe.torch_moe with + Scan the FX graph and replace all calls to torch.ops.auto_deploy.torch_moe with torch.ops.auto_deploy.trtllm_moe_fused. """ ad_logger.debug("Before MoE fusion: " + str(gm)) @@ -106,11 +146,10 @@ def fuse_moe(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: with cuda_memory_tracker(): fused_key_counter = _insert_fused_moe_ops(gm) if fused_key_counter: - gm = canonicalize_graph(gm) + canonicalize_graph(gm) ad_logger.info(f"Found {fused_key_counter} MoE fusions") ad_logger.debug("After MoE fusion: " + str(gm)) - return gm def _insert_fused_moe_ops(gm: GraphModule) -> int: @@ -146,6 +185,7 @@ def _insert_fused_moe_ops(gm: GraphModule) -> int: with graph.inserting_before(node): new_node = graph.call_function( + # TODO(Fridah-nv): torch.ops.auto_deploy.trtllm_moe_fused for quantized models torch.ops.auto_deploy.trtllm_moe_fused, args=( hidden_states, @@ -227,6 +267,32 @@ def lca_two(a: Node, b: Node) -> Optional[Node]: return common +def _extract_linear_parameters(linear_node: Node) -> tuple[Node, torch.Tensor, Optional[dict], str]: + """ + Given a linear op node, extract the input tensor node, weight tensor, + any quantization scales (if the op is quantized), and return a weight type. + + For a torch.ops.auto_deploy.torch_linear_simple.default op: + - Returns (input_node, weight, None, "simple") + + For a torch.ops.auto_deploy.torch_quant_fp8_linear op: + - Returns (input_node, weight, {"input_scale": input_scale, "weight_scale": weight_scale}, "fp8") + For a torch.ops.auto_deploy.torch_quant_fp4_linear op: + - Returns (input_node, weight, {"input_scale": input_scale, "weight_scale": weight_scale, "alpha": alpha}, "fp4") + """ + input_node = linear_node.args[0] + if is_op(linear_node, torch.ops.auto_deploy.torch_linear_simple): + weight = linear_node.args[1] + return input_node, weight, None, "" + elif { + is_op(linear_node, torch.ops.auto_deploy.torch_quant_fp4_linear), + is_op(linear_node, torch.ops.auto_deploy.torch_quant_fp8_linear), + }: + weight = linear_node.args[1] + scales, quant_type = get_scales_and_type_from_node(linear_node) + return input_node, weight, scales, quant_type + + def _match_expert_compute_pattern(start_boundary: Node, end_boundary: Node): """ Match the expert compute pattern between the given boundaries. @@ -235,24 +301,39 @@ def _match_expert_compute_pattern(start_boundary: Node, end_boundary: Node): (F.silu(x @ w1.t()) * (x @ w3.t())) @ w2.t() - For each expert, the function returns: - - pattern_input_nodes: a list of input nodes (x) used for the expert compute. - - pattern_output_nodes: a list of final expert output nodes (the linear op with weight w2). - - expert_weights: a dict with keys "w1", "w2", and "w3" mapping to lists of - corresponding weight nodes from the w1, w2, and w3 branches. + For each expert, the function extracts the input node from the w1 branch and + collects the weight parameters from three linear ops (w1, w3, and w2 branches). + + This function supports both: + - torch.ops.auto_deploy.torch_linear_simple.default ops, and + - torch.ops.auto_deploy.torch_quant_fp8_linear ops (also extracts quantization scales). + - torch.ops.auto_deploy.torch_quant_fp4_linear ops (also extracts quantization scales). + + Returns: + A tuple: + (pattern_input_nodes, pattern_output_nodes, expert_weights, expert_scales, weight_type) + + - pattern_input_nodes: List of input nodes (x) used for the expert compute. + - pattern_output_nodes: List of final expert output nodes (the linear op with weight w2). + - expert_weights: Dict with keys "w1", "w2", "w3" mapping to lists of weight tensors. + - expert_scales: Dict with keys "w1_input_scale", "w1_weight_scale", etc., containing scale tensors + (empty if weight_type is "simple"). + - weight_type: "fp8" if FP8 ops were used, "simple" otherwise. """ pattern_input_nodes, pattern_output_nodes = [], [] expert_weights = defaultdict(list) + expert_scales = defaultdict(list) + weight_type = "simple" # default nodes = list(start_boundary.graph.nodes) region_nodes = nodes[nodes.index(start_boundary) + 1 : nodes.index(end_boundary)] for node in region_nodes: - if not is_linear_op(node): + # Accept both simple and quantized linear ops. + if not is_linear_op(node, include_quantization=True): continue final_linear = node - # Must have at least one argument, and that first argument must be a Node. if not final_linear.args or not isinstance(final_linear.args[0], Node): continue @@ -261,47 +342,68 @@ def _match_expert_compute_pattern(start_boundary: Node, end_boundary: Node): continue arg_a, arg_b = mul_node.args[:2] - # Pick the silu op from either arg_a or arg_b. silu_node = ( arg_a - if (isinstance(arg_a, Node) and is_op(arg_a, torch.ops.aten.silu)) + if is_op(arg_a, torch.ops.aten.silu) else arg_b - if (isinstance(arg_b, Node) and is_op(arg_b, torch.ops.aten.silu)) + if is_op(arg_b, torch.ops.aten.silu) else None ) if silu_node is None: continue - if not ( - silu_node.args - and isinstance(silu_node.args[0], Node) - and is_linear_op(silu_node.args[0]) - ): + if not (silu_node.args and is_linear_op(silu_node.args[0], include_quantization=True)): continue linear_w1_node = silu_node.args[0] # The other branch should be a linear op (w3 branch). linear_w3_node = arg_b if arg_a is silu_node else arg_a - if not (isinstance(linear_w3_node, Node) and is_linear_op(linear_w3_node)): + if not is_linear_op(linear_w3_node, include_quantization=True): continue if not (linear_w1_node.args and linear_w3_node.args): continue - input_node_w1 = linear_w1_node.args[0] - weight_w1 = linear_w1_node.args[1] if len(linear_w1_node.args) > 1 else None - weight_w3 = linear_w3_node.args[1] if len(linear_w3_node.args) > 1 else None - weight_w2 = final_linear.args[1] if len(final_linear.args) > 1 else None + # Extract parameters from each linear op. + input_node_w1, weight_w1, quant_params_w1, wt_type_w1 = _extract_linear_parameters( + linear_w1_node + ) + _, weight_w3, quant_params_w3, wt_type_w3 = _extract_linear_parameters(linear_w3_node) + _, weight_w2, quant_params_w2, wt_type_w2 = _extract_linear_parameters(final_linear) if None in (weight_w1, weight_w3, weight_w2): continue + # Ensure the weight type is consistent across branches. + if wt_type_w1 != wt_type_w3 or wt_type_w1 != wt_type_w2: + continue + weight_type = wt_type_w1 + pattern_input_nodes.append(input_node_w1) pattern_output_nodes.append(final_linear) expert_weights["w1"].append(weight_w1) expert_weights["w3"].append(weight_w3) expert_weights["w2"].append(weight_w2) - return pattern_input_nodes, pattern_output_nodes, expert_weights + # TODO: sanity check that all experts have same weight type + if weight_type == "fp8": + expert_scales["w1_input_scale"].append(quant_params_w1["input_scale"]) + expert_scales["w1_weight_scale"].append(quant_params_w1["weight_scale"]) + expert_scales["w3_input_scale"].append(quant_params_w3["input_scale"]) + expert_scales["w3_weight_scale"].append(quant_params_w3["weight_scale"]) + expert_scales["w2_input_scale"].append(quant_params_w2["input_scale"]) + expert_scales["w2_weight_scale"].append(quant_params_w2["weight_scale"]) + elif weight_type == "fp4": + expert_scales["w1_input_scale"].append(quant_params_w1["input_scale"]) + expert_scales["w1_weight_scale"].append(quant_params_w1["weight_scale"]) + expert_scales["w1_alpha"].append(quant_params_w1["alpha"]) + expert_scales["w3_input_scale"].append(quant_params_w3["input_scale"]) + expert_scales["w3_weight_scale"].append(quant_params_w3["weight_scale"]) + expert_scales["w3_alpha"].append(quant_params_w3["alpha"]) + expert_scales["w2_input_scale"].append(quant_params_w2["input_scale"]) + expert_scales["w2_weight_scale"].append(quant_params_w2["weight_scale"]) + expert_scales["w2_alpha"].append(quant_params_w2["alpha"]) + + return pattern_input_nodes, pattern_output_nodes, expert_weights, expert_scales, weight_type def _find_final_hidden_state_node( @@ -376,7 +478,7 @@ def _extract_index_branches_from_expert_outputs( if not mul or len(mul.args) < 2: continue idx_node = mul.args[1] - if not (isinstance(idx_node, Node) and is_op(idx_node, torch.ops.aten.index)): + if not is_op(idx_node, torch.ops.aten.index): continue routing_branches.append(idx_node.args[0]) experts = idx_node.args[1] diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/library/fusion.py b/tensorrt_llm/_torch/auto_deploy/transformations/library/fusion.py index 11cd1b6e54ad..e66ced8ae696 100644 --- a/tensorrt_llm/_torch/auto_deploy/transformations/library/fusion.py +++ b/tensorrt_llm/_torch/auto_deploy/transformations/library/fusion.py @@ -116,7 +116,7 @@ def split_output(tensor: torch.Tensor) -> Tuple[torch.Tensor, ...]: gm.delete_all_unused_submodules() -def fuse_gemms(gm: GraphModule) -> GraphModule: +def fuse_gemms(gm: GraphModule) -> None: ad_logger.info("GEMM fusion") ad_logger.debug("Before GEMM fusion: " + str(gm)) # sort linear nodes by parent node @@ -139,8 +139,7 @@ def fuse_gemms(gm: GraphModule) -> GraphModule: _insert_fused_gemm(gm, idx := idx + 1, parent_node, lin_children) # clean up and return - gm = canonicalize_graph(gm) + canonicalize_graph(gm) ad_logger.debug("After GEMM fusion: " + str(gm)) torch.cuda.empty_cache() - return gm diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/library/kvcache.py b/tensorrt_llm/_torch/auto_deploy/transformations/library/kvcache.py index 97a4ef3fdac0..62a9d355602f 100644 --- a/tensorrt_llm/_torch/auto_deploy/transformations/library/kvcache.py +++ b/tensorrt_llm/_torch/auto_deploy/transformations/library/kvcache.py @@ -1,7 +1,7 @@ """Graph transformation to automatically add kv cache into fused MHA op.""" import operator -from typing import Dict +from typing import Dict, Type import torch from torch.fx import Graph, GraphModule, Node @@ -14,7 +14,7 @@ from .._graph import add_graph_input, canonicalize_graph -def update_in_out_nodes(egm: GraphModule, cm: CachedSequenceInterface) -> GraphModule: +def update_in_out_nodes(egm: GraphModule, cm: CachedSequenceInterface) -> None: """Modify the graph module by adding new input nodes and canonicalizing the graph. The new input nodes correspond to the extra arguments needed for cached and flattened attention. @@ -22,9 +22,6 @@ def update_in_out_nodes(egm: GraphModule, cm: CachedSequenceInterface) -> GraphM Args: egm: The graph module to analyze and modify. cm: Cached sequence interface containing extra argument information. - - Returns: - The updated GraphModule with new input nodes and a canonicalized graph. """ # loop through nodes to get input, output, and get_attr nodes input_nodes, output_nodes = get_all_input_output_nodes(egm.graph) @@ -45,17 +42,15 @@ def update_in_out_nodes(egm: GraphModule, cm: CachedSequenceInterface) -> GraphM input_nodes.append(add_graph_input(egm, name)) ad_logger.info(f"Added {len(new_args)} new input nodes for cached attention metadata") - egm = canonicalize_graph(egm) - - return egm + canonicalize_graph(egm) def insert_cached_attention( egm: GraphModule, cm: CachedSequenceInterface, - attn_descriptor: AttentionDescriptor, + attn_descriptor: Type[AttentionDescriptor], cache_config: CacheConfig, -) -> GraphModule: +) -> None: """Replace uncached source attention node with corresponding cached attn node.""" # Get all attention nodes and their info objects source_op = attn_descriptor.get_source_attention_op() @@ -68,7 +63,7 @@ def insert_cached_attention( if not source_attn_nodes: # If there are no nodes for kv cache insertion found, return current graph - return egm + return # Sanity check if cm.info.is_paged: @@ -131,15 +126,13 @@ def insert_cached_attention( graph.erase_node(attn_node) num_cached_attn_replacements += 1 - egm = canonicalize_graph(egm) + canonicalize_graph(egm) ad_logger.info( f"Replaced {num_cached_attn_replacements} {source_op} ops " f"with {attn_descriptor.get_cached_attention_op()}" ) ad_logger.debug(f"After inserting {attn_descriptor=} with cache: {egm}") - return egm - def resize_kv_cache( egm: GraphModule, @@ -150,8 +143,13 @@ def resize_kv_cache( free_mem_ratio specifies the fraction of available memory to occupy. """ - free_mem, total_mem = torch.cuda.mem_get_info() - ad_logger.info(f"Free memory: {free_mem}, Total memory: {total_mem}") + + def _get_mem_info_in_mb(): + free_mem, total_mem = torch.cuda.mem_get_info() + return free_mem // 1024**2, total_mem // 1024**2 + + free_mem, total_mem = _get_mem_info_in_mb() + ad_logger.info(f"Free memory (MB): {free_mem}, Total memory (MB): {total_mem}") current_cache_size = cm.current_cache_size_bytes() current_num_pages = cm.info.num_pages ad_logger.info( @@ -165,14 +163,16 @@ def resize_kv_cache( try: # Let's run a forward pass to get the memory usage cm.info._set_max_num_tokens_sample() - free_mem_pre, _ = torch.cuda.mem_get_info() - ad_logger.info(f"Free memory before forward pass: {free_mem_pre}") + free_mem_pre, _ = _get_mem_info_in_mb() + ad_logger.info(f"Free memory before forward pass (MB): {free_mem_pre}") + egm(*cm.args) - free_mem_post, _ = torch.cuda.mem_get_info() - ad_logger.info(f"Free memory after forward pass: {free_mem_post}") + + free_mem_post, _ = _get_mem_info_in_mb() + ad_logger.info(f"Free memory after forward pass (MB): {free_mem_post}") memory_for_forward_pass = free_mem_pre - free_mem_post - ad_logger.info(f"Memory for forward pass: {memory_for_forward_pass}") + ad_logger.info(f"Memory for forward pass (MB): {memory_for_forward_pass}") new_cache_size = free_mem_post * free_mem_ratio + current_cache_size new_num_pages = int(new_cache_size // (current_cache_size // current_num_pages)) diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/library/quantization.py b/tensorrt_llm/_torch/auto_deploy/transformations/library/quantization.py index e63e58b7d8ad..0414ed2fe25d 100644 --- a/tensorrt_llm/_torch/auto_deploy/transformations/library/quantization.py +++ b/tensorrt_llm/_torch/auto_deploy/transformations/library/quantization.py @@ -11,7 +11,6 @@ get_quantization_params_from_linear_node, is_bmm_op, is_linear_op, - is_match, ) from ...utils.quantization_utils import ( QuantizationImpl, @@ -19,6 +18,7 @@ is_quantized_graph, is_quantized_op, remove_output_quantizers, + should_skip_quantization, ) from .._graph import canonicalize_graph @@ -169,23 +169,22 @@ def get_scale_name(scale_name): node.args = (*node.args, *scale_values) -def quantize(gm: GraphModule, quant_config: Dict[str, Any]): - """Quantize the GraphModule and replace linear and bmm with quantized versions.""" +def quantize(gm: GraphModule, quant_config: Dict[str, Any]) -> None: + """Quantize the GraphModule and replace linear with quantized linear.""" # extract info from quant_config is_quant_graph = is_quantized_graph(gm) quant_algo = quant_config.get("quant_algo") - skip = quant_config.get("exclude_modules", []) + excluded_patterns = quant_config.get("exclude_modules", []) # no quantization to do if not (is_quant_graph or quant_config): ad_logger.info("No quantization to do.") - return gm + return # tracking quantized operations in the graph quantized_nodes: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int)) for n in gm.graph.nodes: - # check if we should skip this node - if is_match(n, skip): + if should_skip_quantization(n, excluded_patterns): continue # Process linear operations @@ -215,10 +214,8 @@ def quantize(gm: GraphModule, quant_config: Dict[str, Any]): if is_quant_graph: remove_output_quantizers(gm) - gm = canonicalize_graph(gm) + canonicalize_graph(gm) for quant_algo in quantized_nodes: for op_type, count in quantized_nodes[quant_algo].items(): ad_logger.info(f"Found {count} {quant_algo} quantized {op_type} nodes.") ad_logger.debug("After quantization: " + str(gm)) - - return gm diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/library/quantize_moe.py b/tensorrt_llm/_torch/auto_deploy/transformations/library/quantize_moe.py new file mode 100644 index 000000000000..93890d1da8c3 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transformations/library/quantize_moe.py @@ -0,0 +1,167 @@ +from functools import partial +from typing import Any, Callable, Dict, List, Tuple + +import torch +import torch.nn as nn +from torch.fx import GraphModule, Node + +from ...utils.logger import ad_logger +from ...utils.node_utils import is_op +from ...utils.quantization_utils import QuantizationImpl, should_skip_quantization +from .._graph import canonicalize_graph + +quantized_moe_op_map = { + "FP8": torch.ops.auto_deploy.torch_quant_fp8_moe, + "NVFP4": torch.ops.auto_deploy.torch_quant_fp4_moe, +} + + +def _quantize_moe_node( + gm: GraphModule, + node: Node, + quant_impl: QuantizationImpl, + quantized_op: Callable[..., Node], +): + """ + Replace a torch.ops.auto_deploy.torch_moe node with its quantized version, + quantizing each expert weight list and registering scales + hooks. + Automatically handles different scale configurations per quantization type. + """ + w1_names, w2_names, w3_names = _extract_moe_weight_param_lists(node) + + scale_keys = quant_impl.scale_names() + + def quantize_param_list(weight_names: List[str]) -> Tuple[List[Node], List[List[Node]]]: + new_attrs = [] + scale_nodes_group = [] + for name in weight_names: + orig_weight = gm.get_parameter(name) + new_weight = quant_impl.quantize_weight(orig_weight) + + # Replace parameter in submodule + modname, _, attrname = name.rpartition(".") + submod = gm.get_submodule(modname) + setattr(submod, attrname, nn.Parameter(new_weight, requires_grad=False)) + + # Register new scale buffers + for scale_name, scale_val in quant_impl.default_scales(orig_weight.shape).items(): + submod.register_buffer(scale_name, scale_val) + + # Register load hook + gm._register_load_state_dict_pre_hook(partial(quant_impl.load_hook, weight_name=name)) + + # Create get_attr nodes for new param and each scale + with gm.graph.inserting_before(node): + new_weight_attr = gm.graph.get_attr(name) + new_attrs.append(new_weight_attr) + scales = [gm.graph.get_attr(modname + "." + s) for s in scale_keys] + scale_nodes_group.append(scales) + + return new_attrs, scale_nodes_group + + # Quantize all three expert weights + w1_attrs, w1_scales = quantize_param_list(w1_names) + w2_attrs, w2_scales = quantize_param_list(w2_names) + w3_attrs, w3_scales = quantize_param_list(w3_names) + + # Collect scale tensors per scale type across w1, w2, w3 + def collect_scales(index: int) -> Tuple[List[Node], List[Node], List[Node]]: + return ( + [s[index] for s in w1_scales], + [s[index] for s in w2_scales], + [s[index] for s in w3_scales], + ) + + # Prepare args + args = [ + node.args[0], # x + node.args[1], # selected_experts + node.args[2], # routing_weights + w1_attrs, + w2_attrs, + w3_attrs, + ] + + for idx in range(len(scale_keys)): + s1, s2, s3 = collect_scales(idx) + args.extend([s1, s2, s3]) + + # Replace the current node with the quantized version + with gm.graph.inserting_after(node): + new_node = gm.graph.call_function( + quantized_op, + args=tuple(args), + ) + ad_logger.debug(f"Updating {node.name} args to {new_node.args}") + node.replace_all_uses_with(new_node) + gm.graph.erase_node(node) + + +def quantize_moe(gm: GraphModule, quant_config: Dict[str, Any]) -> None: + """ + Traverse gm, find every torch.ops.auto_deploy.torch_moe, and replace it with the + quantized version using the quant_algo from quant_config. + """ + quant_algo = quant_config.get("quant_algo") + if not quant_algo: + ad_logger.info("No quantization to do.") + return gm + excluded_patterns = quant_config.get("exclude_modules", []) + + quant_impl = QuantizationImpl.create(quant_algo) + quantized_op = quantized_moe_op_map[quant_algo] + + count = 0 + + for node in list(gm.graph.nodes): + if is_op(node, torch.ops.auto_deploy.torch_moe): + # Check that all expert weights should be quantized + w1_names, w2_names, w3_names = _extract_moe_weight_param_lists(node) + if any( + should_skip_quantization(n, excluded_patterns) + for n in w1_names + w2_names + w3_names + ): + continue + _quantize_moe_node(gm, node, quant_impl, quantized_op) + count += 1 + + if count == 0: + return gm + + gm = canonicalize_graph(gm) + ad_logger.info(f"Found {count} {quant_algo} quantized {quantized_op} nodes.") + return + + +# TODO(Fridah-nv): robust handling similar to `extract_param_names_from_lin_node` or expand it +def _extract_moe_weight_param_lists(moe_node: Node) -> Tuple[List[str], List[str], List[str]]: + """ + Given a torch.ops.moe.torch_moe node in gm.graph, extract three lists of + the parameter names for w1_weight, w2_weight, and w3_weight. + + Returns: + (w1_names, w2_names, w3_names), each a list of strings like 'layer.expert_0.w1.weight' + """ + # args layout: (x, selected_experts, routing_weights, w1_list, w2_list, w3_list) + try: + w1_list, w2_list, w3_list = moe_node.args[3:6] + except ValueError: + raise RuntimeError( + f"Expected moe_node.args to have at least 6 entries, got {len(moe_node.args)}" + ) + + def _unwrap_list(arg) -> List[str]: + if not isinstance(arg, (list, tuple)): + raise TypeError(f"Expected a Python list/tuple of get_attr Nodes, got {type(arg)}") + names: List[str] = [] + for elt in arg: + if not isinstance(elt, Node) or elt.op != "get_attr": + raise RuntimeError(f"Expected each list element to be a get_attr Node, got {elt}") + names.append(elt.target) + return names + + w1_names = _unwrap_list(w1_list) + w2_names = _unwrap_list(w2_list) + w3_names = _unwrap_list(w3_list) + + return w1_names, w2_names, w3_names diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/library/rms_norm.py b/tensorrt_llm/_torch/auto_deploy/transformations/library/rms_norm.py new file mode 100644 index 000000000000..a94758b18193 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transformations/library/rms_norm.py @@ -0,0 +1,113 @@ +"""Graph transform to optimize RMSNorm execution using FlashInfer.""" + +from functools import partial + +import torch +from torch.fx import GraphModule + +from ...utils.logger import ad_logger + +# It is important to import ADPatternMatcherPass from pattern_matcher.py, not from torch._inductor.pattern_matcher +from ...utils.pattern_matcher import ADPatternMatcherPass, register_ad_pattern +from .._graph import canonicalize_graph + +_BACKEND_OPS = { + "flashinfer": torch.ops.auto_deploy.flashinfer_rms_norm, + "triton": torch.ops.auto_deploy.triton_rms_norm, + "torch": torch.ops.auto_deploy.torch_rmsnorm, +} + + +def _rms_norm_pattern(data: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """Implements the RMSNorm pattern for pattern matching. + + Args: + data: Input tensor to normalize. + weight: Scaling weights for the normalized output. + eps: Small constant for numerical stability. + + Returns: + Normalized and scaled tensor. + """ + input_dtype = data.dtype + data = data.to(torch.float32) + variance = data.pow(2).mean(-1, keepdim=True) + data = data * torch.rsqrt(variance + eps) + return weight * data.to(input_dtype) + + +def _rms_norm_replacement( + data: torch.Tensor, weight: torch.Tensor, eps: float, backend: str +) -> torch.Tensor: + """Backend-specific rms_norm implementation. + + Args: + data: Input tensor to normalize. + weight: Scaling weights for the normalized output. + eps: Small constant for numerical stability. + backend: Backend to use for RMSNorm computation ("flashinfer" or "triton"). + + Returns: + Normalized and scaled tensor using the specified backend implementation. + """ + + assert backend.lower() in _BACKEND_OPS, ( + f"Invalid {backend=}; must be one of {list(_BACKEND_OPS)}" + ) + return _BACKEND_OPS[backend.lower()](data, weight, eps) + + +def fuse_rmsnorm(gm: GraphModule, backend: str = "triton") -> None: + """Matches and replaces RMSNorm patterns in the graph with FlashInfer or Triton implementation. + + This function sets up pattern matching to identify RMSNorm operations in the graph + and replaces them with optimized implementations. It uses dummy tensors to register + the pattern matching rules. + + Args: + gm: Input graph module to transform. + backend: Backend to use for RMSNorm computation ("flashinfer" or "triton"). + + Returns: + Transformed graph module with optimized RMSNorm operations. + """ + if backend.lower() not in _BACKEND_OPS: + raise ValueError(f"Invalid backend, must be one of {list(_BACKEND_OPS)}, got {backend}") + ad_logger.info(f"Starting RMSNorm pattern matching with backend: {backend}") + + graph = gm.graph + patterns = ADPatternMatcherPass() + + # Create dummy tensors for pattern matching + bs = 2 + hidden_size = 512 + + def dummy_args(input_dtype: torch.dtype, weight_dtype: torch.dtype, eps: float = 1e-6): + return [ + torch.randn(bs, hidden_size, device="cuda", dtype=input_dtype), + torch.randn(hidden_size, device="cuda", dtype=weight_dtype), + eps, + ] + + # Define configurations for different data types + configs = [ + (torch.bfloat16, torch.bfloat16), + (torch.float16, torch.float16), + (torch.float32, torch.float32), + ] + + # Register patterns for each configuration + for input_dtype, weight_dtype in configs: + register_ad_pattern( + search_fn=_rms_norm_pattern, + replace_fn=partial(_rms_norm_replacement, backend=backend), + patterns=patterns, + dummy_args=dummy_args(input_dtype, weight_dtype), + op_ignore_types={}, + scalar_workaround={"eps": 1e-6}, + ) + + cnt = patterns.apply(graph) + ad_logger.info(f"RMSNorm pattern count: {cnt}") + canonicalize_graph(gm) + ad_logger.debug("RMSNorm pattern matching completed.") diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/library/rope.py b/tensorrt_llm/_torch/auto_deploy/transformations/library/rope.py index 651d0730e554..ae686690e8d7 100644 --- a/tensorrt_llm/_torch/auto_deploy/transformations/library/rope.py +++ b/tensorrt_llm/_torch/auto_deploy/transformations/library/rope.py @@ -119,7 +119,7 @@ def _explicit_not_interleaved(match: Match) -> bool: return not any(isinstance(n, Node) and _match_input_interleave_pattern(n) for n in (q, k)) -def match_rope_pattern(gm: GraphModule) -> GraphModule: +def match_rope_pattern(gm: GraphModule) -> int: graph = gm.graph patterns = ADPatternMatcherPass() @@ -174,12 +174,12 @@ def match_rope_pattern(gm: GraphModule) -> GraphModule: ) num_matches = patterns.apply(graph) - gm = canonicalize_graph(gm) + canonicalize_graph(gm) ad_logger.info(f"Found and matched {num_matches} RoPE patterns") - return gm, num_matches + return num_matches -def match_rope_layout(gm: GraphModule, expected_layout: str = "bsnd") -> GraphModule: +def match_rope_layout(gm: GraphModule, expected_layout: str = "bsnd") -> None: """ Match and transform input and output of rope ops to the layout specified to meet requirements of optimized ops. Supported layout is 'bsnd' (batch, seq, head, dim). @@ -189,7 +189,7 @@ def match_rope_layout(gm: GraphModule, expected_layout: str = "bsnd") -> GraphMo ad_logger.warning( f"Unsupported RoPE layout '{expected_layout}'; expected '{supported}'. Skipping RoPE layout matching." ) - return gm + return ad_logger.info(f"Match RoPE layout to {expected_layout}") @@ -291,12 +291,11 @@ def match_rope_layout(gm: GraphModule, expected_layout: str = "bsnd") -> GraphMo k_rope_new.args = (k_rope_old, 1, 2) if num_rope_layout_matches: - gm = canonicalize_graph(gm) + canonicalize_graph(gm) ad_logger.info(f"Found {num_rope_layout_matches} RoPE layout matches") - return gm -def optimize_rope(gm: GraphModule) -> GraphModule: +def optimize_rope(gm: GraphModule) -> None: """ Scan the FX graph and replace calls to the torch-reference RoPE ops with the optimized `rope::flashinfer` kernel. @@ -317,9 +316,8 @@ def optimize_rope(gm: GraphModule) -> GraphModule: continue num_rope_optimizations += 1 if num_rope_optimizations: - gm = canonicalize_graph(gm) + canonicalize_graph(gm) ad_logger.info(f"Found {num_rope_optimizations} RoPE optimizations") - return gm def _optimize_explicit( diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/library/sharding.py b/tensorrt_llm/_torch/auto_deploy/transformations/library/sharding.py index 3afa7f5064fe..d7ed5918a494 100644 --- a/tensorrt_llm/_torch/auto_deploy/transformations/library/sharding.py +++ b/tensorrt_llm/_torch/auto_deploy/transformations/library/sharding.py @@ -18,12 +18,15 @@ import math import operator +from abc import ABC, abstractmethod from collections import defaultdict +from enum import IntEnum from functools import partial -from typing import Callable, DefaultDict, Dict, List, Set +from typing import Callable, DefaultDict, Dict, List, Literal, Optional, Set import torch import torch.nn as nn +from pydantic import BaseModel, ConfigDict, Field from torch.fx import GraphModule, Node from ...utils.logger import ad_logger @@ -38,6 +41,249 @@ from .._graph import canonicalize_graph +class SplitDimension(IntEnum): + """Enum for tensor split dimensions in sharding.""" + + ROW = 0 # Split along rows (first dimension) + COLUMN = 1 # Split along columns (second dimension) + + +class ShardingTransformInfo(BaseModel, ABC): + """Abstract base class for transformation configurations.""" + + model_config = ConfigDict(frozen=True) # Makes the model immutable and hashable + + target_node: str + rank: int + world_size: int + + def validate(self, gm: GraphModule = None, node: Node = None) -> bool: + """ + Validate whether the transformation is valid. + Execute right before applying the transformation. + """ + return True + + @abstractmethod + def apply(self, gm: GraphModule, node: Node) -> None: + """Apply the transformation to the graph module. + + This method must be implemented by each transformation class. + """ + pass + + def check_and_apply(self, gm: GraphModule, node: Node) -> None: + """Check if the transformation is valid and apply it if it is.""" + if not self.validate(gm, node): + ad_logger.warning(f"Skipping invalid transformation {self}.") + return + self.apply(gm, node) + + +class TPShardingInfo(ShardingTransformInfo): + """Configuration for TP sharding transformations.""" + + split_dim: SplitDimension + dist_op: Optional[Literal["all_reduce", "all_gather"]] = None + min_local_shape: int = 1 + + def validate(self, gm: GraphModule = None, node: Node = None) -> bool: + """Validate the transformation configuration.""" + if self.dist_op is not None: + if self.split_dim == SplitDimension.ROW: + if self.dist_op == "all_reduce": + ad_logger.warning( + f"Row split is only supported for all_gather. Skipping {self}." + ) + return False + if self.split_dim == SplitDimension.COLUMN: + if self.dist_op == "all_gather": + ad_logger.warning( + f"Column split is only supported for all_reduce. Skipping {self}." + ) + return False + return True + + def apply(self, gm: GraphModule, node: Node) -> None: + """Apply TP sharding transformation to the graph module.""" + + _insert_sharded_matmul( + gm=gm, + node=node, + dim=self.split_dim.value, + rank=self.rank, + world_size=self.world_size, + add_dist=self.dist_op is not None, + min_local_shape=self.min_local_shape, + ) + + +class BMMShardingInfo(ShardingTransformInfo): + """Configuration for BMM sharding transformations.""" + + rank: int + world_size: int + start_idx: int + end_idx: int + + def validate(self, gm: GraphModule = None, node: Node = None) -> bool: + """Validate the transformation configuration.""" + if not is_op(node, torch.ops.aten.bmm): + ad_logger.warning(f"BMM sharding is only supported for BMM nodes. Skipping {self}.") + return False + + # Get the input tensors + lhs_tensor = node.args[0] + rhs_tensor = node.args[1] + + # Check batch sizes from meta information + lhs_batch_size = lhs_tensor.meta["val"].shape[0] + rhs_batch_size = rhs_tensor.meta["val"].shape[0] + + assert lhs_batch_size == rhs_batch_size, "Batch sizes of both tensors must match" + bmm_batch_size = lhs_batch_size + + # Check if the distribution is balanced + remainder = bmm_batch_size % self.world_size + + # NOTE: our torch.ops.auto_deploy.torch_dist_all_gather doesn't support uneven splits at the moment. + if remainder: + ad_logger.warning( + f"BMM batch size {bmm_batch_size} is not divisible by world size {self.world_size}. " + f"This will result in uneven distribution of work across devices. Skipping." + ) + return False + return True + + def apply(self, gm: GraphModule, node: Node) -> None: + """Apply BMM sharding transformation to the graph module.""" + + def handle_tensor( + bmm_node: Node, tensor_node: Node, arg_idx: int, start_idx: int, end_idx: int + ): + """Unified helper function to shard either a parameter tensor or a dynamic tensor. + + Args: + bmm_node: The BMM node that is being processed + tensor_node: The input tensor node to shard + arg_idx: The argument index of the tensor in the BMM node + start_idx: Start index for sharding + end_idx: End index for sharding + """ + + # Define slice function for the sharding + def slice_tensor(t: torch.Tensor) -> torch.Tensor: + return t[start_idx:end_idx] + + if tensor_node.op == "get_attr": + # Handle parameter tensor + weight_key = tensor_node.target + modname, _, param_name = weight_key.rpartition(".") + param = gm.get_parameter(weight_key) + + # Update the parameter with its shard + param_new = nn.Parameter(slice_tensor(param).detach().clone(), requires_grad=True) + gm.get_submodule(modname).register_parameter(param_name, param_new) + + # Register load state dict hook + gm._register_load_state_dict_pre_hook( + partial( + _load_hook, + f_split=slice_tensor, + param_key=weight_key, + param_shape=param_new.shape, + ) + ) + else: + # Handle dynamic tensor + with gm.graph.inserting_before(bmm_node): + tensor_slice = gm.graph.call_function( + torch.ops.aten.slice.Tensor, args=(tensor_node, 0, start_idx, end_idx, 1) + ) + # Update BMM node to use the sliced tensor + bmm_node.update_arg(arg_idx, tensor_slice) + + # Get the input tensors + lhs_tensor = node.args[0] + rhs_tensor = node.args[1] + # Handle both tensors + handle_tensor(node, lhs_tensor, 0, self.start_idx, self.end_idx) + handle_tensor(node, rhs_tensor, 1, self.start_idx, self.end_idx) + + # Add all_gather node after BMM to collect results + with gm.graph.inserting_after(node): + gather_node = gm.graph.call_function( + torch.ops.auto_deploy.torch_dist_all_gather, + args=(node, 0), # Gather along batch dimension (0) + ) + node.replace_all_uses_with(gather_node) + gather_node.replace_input_with(gather_node, node) + + +class EPShardingInfo(ShardingTransformInfo): + """Configuration for EP sharding transformations.""" + + rank: int + world_size: int + + def validate(self, gm: GraphModule = None, node: Node = None) -> bool: + """Validate the transformation configuration.""" + if not is_op( + node, + ( + torch.ops.auto_deploy.torch_moe, + torch.ops.auto_deploy.torch_quant_fp8_moe, + torch.ops.auto_deploy.torch_quant_fp4_moe, + ), + ): + ad_logger.warning(f"EP sharding is only supported for MOE nodes. Skipping {self}.") + return False + return True + + def apply(self, gm: GraphModule, node: Node) -> None: + """Apply EP sharding transformation to the graph module.""" + _insert_sharded_moe(gm, node, self.rank, self.world_size) + + +class ShardingConfig(BaseModel): + """Configuration for sharding the model.""" + + tp_transforms: List[TPShardingInfo] = Field(default_factory=list) + bmm_transforms: List[BMMShardingInfo] = Field(default_factory=list) + ep_transforms: List[EPShardingInfo] = Field(default_factory=list) + + +def sharding_transform_executor(gm: GraphModule, sharding_config: ShardingConfig) -> None: + """Apply transformations to the graph module. + + Args: + gm: Graph module to apply transformations to + sharding_config: Transformation configuration containing list of transformations to apply + """ + # create a node dict for faster lookup + node_dict = {n.name: n for n in gm.graph.nodes} + + def check_and_apply(transform: ShardingTransformInfo) -> None: + if transform.target_node is None or transform.target_node not in node_dict: + ad_logger.warning( + f"Skipping transformation {transform} because target node " + + f"{transform.target_node} not found in graph" + ) + return + transform.check_and_apply(gm, node_dict[transform.target_node]) + + for tp_transform in sharding_config.tp_transforms: + check_and_apply(tp_transform) + for bmm_transform in sharding_config.bmm_transforms: + check_and_apply(bmm_transform) + for ep_transform in sharding_config.ep_transforms: + check_and_apply(ep_transform) + + # canonicalize and return + gm = canonicalize_graph(gm) + ad_logger.debug("After applying sharding transformations: " + str(gm)) + + def _load_hook( state_dict, prefix, @@ -79,8 +325,8 @@ def _insert_sharded_matmul( world_size: int, add_dist: bool = False, min_local_shape: int = 1, -): - """Replaces the matmul node with a new matmul node that accepts sharded weights. +) -> None: + """Replace the matmul node with a new matmul node that accepts sharded weights. The state_dict is also updated to contain the sharded weights. """ @@ -200,22 +446,37 @@ def set_new_param(submod: nn.Module, param_key: str, remove: bool = False) -> to dist_node.replace_input_with(dist_node, node) -def _simple_shard( - gm: GraphModule, nodes_linear: Dict[Node, List[Node]], rank: int, world_size: int -): +def _append_simple_shard( + nodes_linear: Dict[Node, List[Node]], + rank: int, + world_size: int, + sharding_config: ShardingConfig, +) -> None: # for every linear node: # --> row_split (dim 0 of weight) + all_gather (dim -1 of output) + tp_shards: List[TPShardingInfo] = [] for node_group in nodes_linear.values(): for n in node_group: - _insert_sharded_matmul(gm, n, 0, rank, world_size, add_dist=True) + tp_shards.append( + TPShardingInfo( + target_node=n.name, + split_dim=SplitDimension.ROW, + rank=rank, + world_size=world_size, + dist_op="all_gather", + min_local_shape=1, + ) + ) + sharding_config.tp_transforms.extend(tp_shards) -def column_row_shard( +def detect_column_row_shard( gm: GraphModule, rank: int, world_size: int, + sharding_config: ShardingConfig, simple_shard_only: bool = False, -) -> GraphModule: +) -> None: """A transformation to apply sharding to the model following tensor parallelism. The transformation is based on the following steps: @@ -236,7 +497,7 @@ def column_row_shard( if world_size < 2: ad_logger.info("Skipping sharding for single device") - return gm + return assert isinstance(gm, GraphModule), "Expecting GraphModule" @@ -312,13 +573,13 @@ def column_row_shard( if simple_shard_only: ad_logger.debug(f"Forcing Simple Shard: Linear groups: {nodes_linear}") - _simple_shard(gm, nodes_linear, rank, world_size) + _append_simple_shard(nodes_linear, rank, world_size, sharding_config) continue # simple shard when we have != 2 groups of linear nodes if len(nodes_linear) != 2: ad_logger.debug(f"Linear groups: {nodes_linear}") - _simple_shard(gm, nodes_linear, rank, world_size) + _append_simple_shard(nodes_linear, rank, world_size, sharding_config) continue # let's look at the unnacounted nodes. They are okay as long as they fall before the @@ -348,7 +609,7 @@ def column_row_shard( # check if any unaccounted nodes are left. If so, do a simply shard if unaccounted_nodes or attention_related_nodes: ad_logger.debug(f"Unaccounted nodes: {unaccounted_nodes}") - _simple_shard(gm, nodes_linear, rank, world_size) + _append_simple_shard(nodes_linear, rank, world_size, sharding_config) continue # If we can account for all sharded nodes, we can do a two-way shard @@ -360,7 +621,7 @@ def column_row_shard( # Column-row shard boundary region detection is probably wrong - there should be # only one attention operation. Fall back to simple shard. ad_logger.debug(f"More than one attention node: {unaccounted_nodes}") - _simple_shard(gm, nodes_linear, rank, world_size) + _append_simple_shard(nodes_linear, rank, world_size, sharding_config) continue # Extract head dimension. We cannot shard below the head_dim size. # Assume that head_dim is the last (innermost) dimension of the tensor @@ -369,19 +630,27 @@ def column_row_shard( min_local_shape = 1 for i, group in enumerate(nodes_linear.values()): for n in group: - _insert_sharded_matmul( - gm, n, i, rank, world_size, add_dist=i > 0, min_local_shape=min_local_shape + if i > 0: + dist_op = "all_reduce" + else: + dist_op = None + sharding_config.tp_transforms.append( + TPShardingInfo( + target_node=n.name, + split_dim=i, + rank=rank, + world_size=world_size, + dist_op=dist_op, + min_local_shape=min_local_shape, + ) ) - # canonicalize and return - if num_shards: - gm = canonicalize_graph(gm) - ad_logger.debug("After sharding: " + str(gm)) ad_logger.info(f"Found {num_shards} TP shards") - return gm -def dp_bmm_shard(gm: GraphModule, rank: int, world_size: int) -> GraphModule: +def detect_dp_bmm_shard( + gm: GraphModule, rank: int, world_size: int, sharding_config: ShardingConfig +) -> None: """A transformation to apply sharding to batched matrix multiplications in the graph. We'll shard the BMM nodes by slicing the batch dimension of input tensors into world_size number of slices. @@ -394,57 +663,12 @@ def dp_bmm_shard(gm: GraphModule, rank: int, world_size: int) -> GraphModule: if world_size < 2: ad_logger.info("Skipping sharding for single device") - return gm + return assert isinstance(gm, GraphModule), "Expecting GraphModule" num_bmm_shards = 0 - def handle_tensor( - bmm_node: Node, tensor_node: Node, arg_idx: int, start_idx: int, end_idx: int - ): - """Unified helper function to shard either a parameter tensor or a dynamic tensor. - - Args: - bmm_node: The BMM node that is being processed - tensor_node: The input tensor node to shard - arg_idx: The argument index of the tensor in the BMM node - start_idx: Start index for sharding - end_idx: End index for sharding - """ - - # Define slice function for the sharding - def slice_tensor(t: torch.Tensor) -> torch.Tensor: - return t[start_idx:end_idx] - - if tensor_node.op == "get_attr": - # Handle parameter tensor - weight_key = tensor_node.target - modname, _, param_name = weight_key.rpartition(".") - param = gm.get_parameter(weight_key) - - # Update the parameter with its shard - param_new = nn.Parameter(slice_tensor(param).detach().clone(), requires_grad=True) - gm.get_submodule(modname).register_parameter(param_name, param_new) - - # Register load state dict hook - gm._register_load_state_dict_pre_hook( - partial( - _load_hook, - f_split=slice_tensor, - param_key=weight_key, - param_shape=param_new.shape, - ) - ) - else: - # Handle dynamic tensor - with gm.graph.inserting_before(bmm_node): - tensor_slice = gm.graph.call_function( - torch.ops.aten.slice.Tensor, args=(tensor_node, 0, start_idx, end_idx, 1) - ) - # Update BMM node to use the sliced tensor - bmm_node.update_arg(arg_idx, tensor_slice) - for node in gm.graph.nodes: if not is_op(node, {torch.ops.aten.bmm}): continue @@ -482,23 +706,19 @@ def slice_tensor(t: torch.Tensor) -> torch.Tensor: start_idx = remainder + rank * base_size end_idx = start_idx + base_size + sharding_config.bmm_transforms.append( + BMMShardingInfo( + target_node=node.name, + rank=rank, + world_size=world_size, + start_idx=start_idx, + end_idx=end_idx, + ) + ) ad_logger.debug( f"Sharding BMM for rank {rank}: batch_size={bmm_batch_size}, start_idx={start_idx}, end_idx={end_idx}" ) - # Handle both tensors - handle_tensor(node, lhs_tensor, 0, start_idx, end_idx) - handle_tensor(node, rhs_tensor, 1, start_idx, end_idx) - - # Add all_gather node after BMM to collect results - with gm.graph.inserting_after(node): - gather_node = gm.graph.call_function( - torch.ops.auto_deploy.torch_dist_all_gather, - args=(node, 0), # Gather along batch dimension (0) - ) - node.replace_all_uses_with(gather_node) - gather_node.replace_input_with(gather_node, node) - num_bmm_shards += 1 # Canonicalize and return @@ -506,4 +726,123 @@ def slice_tensor(t: torch.Tensor) -> torch.Tensor: gm = canonicalize_graph(gm) ad_logger.debug("After sharding BMM: " + str(gm)) ad_logger.info(f"Found {num_bmm_shards} BMM shards") - return gm + + +def detect_ep_shard( + gm: GraphModule, rank: int, world_size: int, sharding_config: ShardingConfig +) -> None: + ad_logger.debug("Before sharding graph: " + str(gm)) + + if world_size < 2: + ad_logger.info("Skipping sharding for single device") + return + + assert isinstance(gm, GraphModule), "Expecting GraphModule" + num_moe_patterns = 0 + for node in list(gm.graph.nodes): + if not is_op( + node, + ( + torch.ops.auto_deploy.torch_moe, + torch.ops.auto_deploy.torch_quant_fp8_moe, + torch.ops.auto_deploy.torch_quant_fp4_moe, + ), + ): + continue + sharding_config.ep_transforms.append( + EPShardingInfo( + target_node=node.name, + rank=rank, + world_size=world_size, + ) + ) + num_moe_patterns += 1 + + ad_logger.info(f"Found {num_moe_patterns} MoE patterns") + + +def _insert_sharded_moe( + gm: GraphModule, + node: Node, + rank: int, + world_size: int, +): + """Update the torch_moe node with sharded weight lists, + sharded `selected_experts` and `final_scales(router_logics)`. + Add an all_reduce node after the moe node. + """ + quant_impl = QuantizationImpl.create(node) + scale_names = quant_impl.scale_names() if quant_impl else [] + + num_experts = len(node.args[3]) + args = list(node.args) + + # -- Handle selected_experts and final_scales sharding -- + selected_experts = args[1] + final_scales = args[2] + + experts_per_rank = num_experts // world_size + + with gm.graph.inserting_before(node): + lower = experts_per_rank * rank + # selected_experts_local = selected_experts - low + selected_experts_local = gm.graph.create_node( + "call_function", operator.sub, args=(selected_experts, lower), kwargs={} + ) + + # For num_experts % world_size != 0 case, + # assign the last (num_experts % world_size) experts to the last rank + # if rank == world_size -1: + # rank_mask = (selected_experts // experts_per_rank) >= rank + # else: + # rank_mask = (selected_experts // experts_per_rank) == rank + div_node = gm.graph.create_node( + "call_function", operator.floordiv, args=(selected_experts, experts_per_rank), kwargs={} + ) + comp_op = torch.ge if rank == world_size - 1 else torch.eq + rank_mask = gm.graph.create_node("call_function", comp_op, args=(div_node, rank), kwargs={}) + + # final_scales_local = final_scales * rank_mask + final_scales_local = gm.graph.create_node( + "call_function", operator.mul, args=(final_scales, rank_mask), kwargs={} + ) + + # -- Shard expert weights -- + def get_partition(lst, world_size, rank): + num_experts = len(lst) + expert_size_per_partition = num_experts // world_size + expert_start = rank * expert_size_per_partition + # For num_experts % world_size != 0 case, + # assign the last (num_experts % world_size) experts to the last rank + expert_end = ( + num_experts if (rank == world_size - 1) else expert_start + expert_size_per_partition + ) + return lst[expert_start:expert_end] + + w1_list_sharded = get_partition(args[3], world_size, rank) + w2_list_sharded = get_partition(args[4], world_size, rank) + w3_list_sharded = get_partition(args[5], world_size, rank) + + # -- Update args -- + args[1] = selected_experts_local + args[2] = final_scales_local + args[3] = w1_list_sharded + args[4] = w2_list_sharded + args[5] = w3_list_sharded + + # Shard scales for quantized ops + for i in range(len(scale_names) * 3): # 3 layers (w1, w2, w3) × #scale_names per layer + args[6 + i] = get_partition(args[6 + i], world_size, rank) + + ad_logger.debug( + f"Updated node {node}: replaced original arguments {node.args} with sharded arguments {args}." + ) + node.args = tuple(args) + + # -- add an all_reduce node -- + with gm.graph.inserting_after(node): + dist_node = gm.graph.call_function( + torch.ops.auto_deploy.torch_dist_all_reduce, args=(node,) + ) + node.replace_all_uses_with(dist_node) + dist_node.replace_input_with(dist_node, node) diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/library/visualization.py b/tensorrt_llm/_torch/auto_deploy/transformations/library/visualization.py index d02cdecd4f29..aaf77ac8e8cd 100644 --- a/tensorrt_llm/_torch/auto_deploy/transformations/library/visualization.py +++ b/tensorrt_llm/_torch/auto_deploy/transformations/library/visualization.py @@ -5,12 +5,11 @@ import model_explorer import torch +import torch.export as te from model_explorer.graph_builder import GraphNode, KeyValue, MetadataItem from model_explorer.pytorch_exported_program_adater_impl import PytorchExportedProgramAdapterImpl from torch import fx -from ..export import torch_export - def print_tensor(self, tensor: torch.Tensor, size_limit: int = 16): shape = tensor.shape @@ -79,7 +78,7 @@ def add_outputs_metadata(self, fx_node: torch.fx.node.Node, node: GraphNode): # TODO(yudong): make viz as non-block call. def visualize_namespace(gm: fx.GraphModule, args: Tuple[torch.Tensor, ...], dynamic_shapes): - ep = torch_export(gm, args=args, dynamic_shapes=dynamic_shapes) + ep = te.export(gm, args=args, dynamic_shapes=dynamic_shapes) graph = ep.graph # Ensure the ops land up in the right module for better viz for n in graph.nodes: diff --git a/tensorrt_llm/_torch/auto_deploy/transformations/transform.py b/tensorrt_llm/_torch/auto_deploy/transformations/transform.py index 9d15af032543..a2f31644d5b8 100644 --- a/tensorrt_llm/_torch/auto_deploy/transformations/transform.py +++ b/tensorrt_llm/_torch/auto_deploy/transformations/transform.py @@ -3,24 +3,26 @@ import gc import torch -from torch.fx import GraphModule +import torch.nn as nn from ..compile import compile_and_capture from ..custom_ops.attention_interface import AttentionRegistry from ..distributed import common as dist_ad -from ..llm_args import LlmArgs +from ..llm_args import AutoDeployConfig from ..models.factory import ModelFactory from ..shim.interface import CachedSequenceInterface +from ..transform.optimizer import InferenceOptimizer as ModularInferenceOptimizer from ..utils.logger import ad_logger from ._graph import canonicalize_graph, lift_to_meta, move_to_device -from .export import torch_export_to_gm from .library import ( - column_row_shard, - dp_bmm_shard, + ShardingConfig, + detect_column_row_shard, + detect_dp_bmm_shard, + detect_ep_shard, eliminate_redundant_transposes, - ep_shard, fuse_allreduce_residual_rmsnorm, fuse_collectives, + fuse_rmsnorm, insert_cached_attention, match_attention_layout, match_causal_attn_mask, @@ -32,17 +34,19 @@ match_rope_pattern, optimize_rope, quantize, + quantize_moe, resize_kv_cache, + sharding_transform_executor, update_in_out_nodes, ) class InferenceOptimizer: - def __init__(self, factory: ModelFactory, ad_config: LlmArgs): + def __init__(self, factory: ModelFactory, ad_config: AutoDeployConfig): self.factory = factory self.ad_config = ad_config - def __call__(self, cm: CachedSequenceInterface) -> GraphModule: + def __call__(self, cm: CachedSequenceInterface) -> nn.Module: """Transform a model into an optimized inference model. Args: @@ -54,53 +58,46 @@ def __call__(self, cm: CachedSequenceInterface) -> GraphModule: quantization: The quantization method to use. Defaults to None. Returns: - A GraphModule representing the optimized inference model. + A nn.Module representing the optimized inference model. """ ############################################################################################ - # INITIALIZE MODEL + # RUN MODULAR INFERENCE OPTIMIZER FOR ALREADY-MIGRATED TRANSFORMS ############################################################################################ - model = self.factory.build_model(device="meta") + new_optimizer = ModularInferenceOptimizer(self.factory, self.ad_config.transforms) + egm = new_optimizer(cm) - ############################################################################################ - # EXPORT MODEL TO GRAPH MODULE - ############################################################################################ - - cm.info.set_example_sequence() - egm = torch_export_to_gm(model, args=cm.args, dynamic_shapes=cm.dynamic_shapes) - del model - ad_logger.debug("original graph: " + str(egm)) - local_rank, world_size = dist_ad.get_rank_world_size() + # TODO (lucaslie): continue moving legacy transforms to the new optimizer ############################################################################################ # RUN PATTERN MATCHER TRANSFORMATIONS TO STANDARDIZE GRAPH REPRESENTATION ############################################################################################ - # quantization - egm = quantize(egm, self.factory.get_quant_config()) + quantize(egm, self.factory.get_quant_config()) + quantize_moe(egm, self.factory.get_quant_config()) # Match MoE pattern - egm = match_moe_pattern(egm) + match_moe_pattern(egm) # Match repeat_kv pattern - egm = match_repeat_kv(egm) + match_repeat_kv(egm) # Match eager attention pattern - egm = match_eager_attention(egm) + match_eager_attention(egm) # Match grouped attention pattern - egm = match_grouped_attention(egm) + match_grouped_attention(egm) # Match and optimize causal attention masks - egm = match_causal_attn_mask(egm) + match_causal_attn_mask(egm) # Match attention layout expected by our backend - egm = match_attention_layout(egm, AttentionRegistry.get(self.ad_config.attn_backend)) + match_attention_layout(egm, AttentionRegistry.get(self.ad_config.attn_backend)) # Match rope - egm, _ = match_rope_pattern(egm) + match_rope_pattern(egm) # Match RoPE layout expected by our backend - egm = match_rope_layout( + match_rope_layout( egm, AttentionRegistry.get(self.ad_config.attn_backend).get_attention_layout() ) @@ -108,26 +105,35 @@ def __call__(self, cm: CachedSequenceInterface) -> GraphModule: # RUN TRANSFORMATIONS ON STANDARDIZED GRAPH REPRESENTATION ############################################################################################ + local_rank, world_size = dist_ad.get_rank_world_size() + # eliminate redundant transpose operations - egm = eliminate_redundant_transposes(egm) + eliminate_redundant_transposes(egm) # TODO (lucaslie): let's move this to perf optimization once TP sharding is improved # see https://github.com/NVIDIA/TensorRT-LLM/pull/3668#discussion_r2052714528 - egm = optimize_rope(egm) + optimize_rope(egm) + + # TODO: Infer sharding parameters (tp_size, row/column sharding) from the model config. + sharding_config = ShardingConfig() # run TP sharding across ranks - egm = column_row_shard(egm, local_rank, world_size, self.ad_config.simple_shard_only) + detect_column_row_shard( + egm, local_rank, world_size, sharding_config, self.ad_config.simple_shard_only + ) # run EP sharding across ranks - egm = ep_shard(egm, local_rank, world_size) + detect_ep_shard(egm, local_rank, world_size, sharding_config) # run BMM sharding across ranks - egm = dp_bmm_shard(egm, local_rank, world_size) + detect_dp_bmm_shard(egm, local_rank, world_size, sharding_config) + + sharding_transform_executor(egm, sharding_config) # let's run a shape propagation pass to update the graph with correct meta values for # subsequent optimization passes. Lift state_dict to meta as shape propagation involves device check with lift_to_meta(egm): - egm = canonicalize_graph(egm, shape_prop=True) + canonicalize_graph(egm, shape_prop=True) ############################################################################################ # MOVE MODEL AND LOAD WEIGHTS @@ -146,17 +152,21 @@ def __call__(self, cm: CachedSequenceInterface) -> GraphModule: # run MoE fusion # TODO: https://github.com/NVIDIA/TensorRT-LLM/issues/4674 this is causing OOMs - # egm = fuse_moe(egm) + # fuse_moe(egm) # run GEMM fusion # TODO: https://github.com/NVIDIA/TensorRT-LLM/issues/4674 this is causing OOMs - # egm = fuse_gemms(egm) + # fuse_gemms(egm) # check if we can fuse allreduce, residual and rmsnorm - egm = fuse_allreduce_residual_rmsnorm(egm) + fuse_allreduce_residual_rmsnorm(egm) # check if we can fuse collectives - egm = fuse_collectives(egm) + fuse_collectives(egm) + + # TODO (lucaslie): add backend selection as part of configurable inference optimizers + # check if we can fuse rmsnorm + fuse_rmsnorm(egm, "flashinfer") # visualize the final graph if self.ad_config.visualize: @@ -175,12 +185,12 @@ def __call__(self, cm: CachedSequenceInterface) -> GraphModule: # SWITCH TO CACHED+FLATTENED ATTENTION + INITIALIZE CACHES ############################################################################################ - egm = update_in_out_nodes(egm, cm) + update_in_out_nodes(egm, cm) # detect attention op and replace with cache-aware op for a_backend in [self.ad_config.attn_backend, self.ad_config.mla_backend]: attn_descriptor = AttentionRegistry.get(a_backend) - egm = insert_cached_attention(egm, cm, attn_descriptor, self.factory.get_cache_config()) + insert_cached_attention(egm, cm, attn_descriptor, self.factory.get_cache_config()) # initialize cache on correct device cm.initialize_caches() diff --git a/tensorrt_llm/_torch/auto_deploy/utils/_config.py b/tensorrt_llm/_torch/auto_deploy/utils/_config.py new file mode 100644 index 000000000000..1d618bf7ab58 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/utils/_config.py @@ -0,0 +1,122 @@ +"""Helper functions for config-related settings.""" + +import os +from pathlib import Path +from typing import Any, Dict, List, Union + +from omegaconf import DictConfig, OmegaConf +from pydantic import Field +from pydantic_settings import BaseSettings, PydanticBaseSettingsSource, YamlConfigSettingsSource +from pydantic_settings.sources.types import PathType + + +def deep_merge_dicts(*confs: Union[Dict, DictConfig]) -> Dict: + """Deep merge a list of dictionaries via OmegaConf.merge. + + Args: + *confs: A list of dictionaries or DictConfig objects to merge. + + Returns: + A merged dictionary. + """ + if len(confs) == 0: + return {} + merged_conf = OmegaConf.merge(*[OmegaConf.create(conf) for conf in confs]) + result = OmegaConf.to_container(merged_conf, resolve=True) + assert isinstance(result, Dict), f"Expected dict, got {type(result)}" + return result + + +class DynamicYamlWithDeepMergeSettingsSource(YamlConfigSettingsSource): + """YAML config settings source that dynamically loads files and merges them via deep update. + + We utilize the omegaconf library for deep merging. + """ + + def _read_files(self, files: PathType | None) -> dict[str, Any]: + if files is None: + return {} + if isinstance(files, (str, os.PathLike)): + files = [files] + + confs = [] + for file in files: + file_path = Path(file).expanduser() + if file_path.is_file(): + confs.append(OmegaConf.load(file_path)) + + return deep_merge_dicts(*confs) + + def __call__(self): + """Call additional config files based on current state.""" + yaml_data = self.yaml_data # this points to the default yaml data now + additional_files_data = self._read_files(self.current_state.get("yaml_configs", [])) + + return deep_merge_dicts(yaml_data, additional_files_data) + + +class DynamicYamlMixInForSettings: + """Mix-in class for settings providing dynamic yaml loading as lowest priority source. + + NOTE: This class must come FIRST in the MRO such that `yaml_configs` can be processed before + since otherwise we cannot load default values from the `yaml_configs` first. + + This mix-in enforces the following precedence order: + - init settings + - env settings + - dotenv settings + - file secret settings + - yaml configs + - default settings + + You can learn more about the different settings sources in + https://docs.pydantic.dev/latest/concepts/pydantic_settings/#field-value-priority. + + Note in particular how yaml settings have precedence only over default settings. You can hence + think of the yaml settings as a way to override default settings. + + Also consider the following consequences of precedence order in nested config settings: + - yaml configs for outer settings get converted to init settings for inner settings and hence + ALWAYS take precedence over yaml configs specified for inner settings. + - This implies inner settings from outer yaml configs also take precedence over outer inner + settings like env settings since they are now init settings from the view of the inner + settings. + - Explicitly initialized fields for inner settings take precedence over outer yaml configs for + inner settings since they are provided as init arguments. + - Check out ``tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_config.py`` for more + examples. + + + You can also provide multiple yaml config files to load. In this case, the files are deep merged + together in the order they are provided. Hence, the following order (decreasing precedence) for + multiple yaml config files is: + - default yaml provided as ``yaml_file`` argument in the ``model_config`` (``ConfigDict``) + - argument 0 of ``yaml_configs`` + - argument 1 of ``yaml_configs`` + - ... + - last argument of ``yaml_configs`` + """ + + yaml_configs: List[PathType] = Field( + default_factory=list, + description="Additional yaml config files to load.", + ) + + @classmethod + def settings_customise_sources( + cls, + settings_cls: type[BaseSettings], + init_settings: PydanticBaseSettingsSource, + env_settings: PydanticBaseSettingsSource, + dotenv_settings: PydanticBaseSettingsSource, + file_secret_settings: PydanticBaseSettingsSource, + ) -> tuple[PydanticBaseSettingsSource, ...]: + """Customise settings sources.""" + deferred_yaml_settings = DynamicYamlWithDeepMergeSettingsSource(settings_cls) + return ( + init_settings, + env_settings, + dotenv_settings, + file_secret_settings, + deferred_yaml_settings, # yaml files have lowest priority just before default values + ) diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py index 709ff91c80d2..48f06c70e60b 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py +++ b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py @@ -25,7 +25,8 @@ modelopt_quantize_op = None modelopt_dynamic_block_quantize_op = None -OperatorLike = Union[OpOverloadPacket, OpOverload, Callable] +OpOrOverload = Union[OpOverloadPacket, OpOverload] +OperatorLike = Union[OpOrOverload, Callable] @dataclass @@ -106,27 +107,17 @@ def get_quantization_params_from_linear_node(linear_op: torch.fx.node.Node): return input_params, weight_params, output_params -def is_match(node: Node, names_to_skip: List[str]): - if names_to_skip is None: - return False - for n in names_to_skip: - module_stack = node.meta.get("nn_module_stack", None) - if module_stack is None: - return False - module_stack = list(module_stack.keys()) - if n in module_stack[-1]: - return True - return False - - def extract_weight_node(mm_node: Node) -> int: - """Extracts the weight node from the given matmul node.""" + """Extracts the weight node from the given linear or BMM node. We assume torch.bmm(activation, weight)""" def find_get_attr_node(node: Node) -> Node: """Recursively traverse inputs of allowed nodes to find a node with 'get_attr' op.""" # If node is a get_attr node return node # List of nodes allowed in between a get_attr node and the matmul node - allowed_ops = {torch.ops.aten.to.dtype} + allowed_ops = { + torch.ops.aten.to.dtype, + torch.ops.aten.view.default, + } if node.op == "get_attr": return node @@ -161,8 +152,8 @@ def extract_param_names_from_lin_node(mm_node: Node) -> Tuple[str, Optional[str] Args: mm_node: Matmul node in the graph. """ - assert is_linear_op(mm_node, include_quantization=True), ( - f"Expecting linear node, Found: {mm_node}" + assert is_linear_op(mm_node, include_quantization=True) or is_bmm_op(mm_node), ( + f"Expecting linear or bmm node, Found: {mm_node}" ) weight_node = extract_weight_node(mm_node) @@ -215,6 +206,37 @@ def is_op(node: Node, ops: Union[OperatorLike, Iterable[OperatorLike]]) -> bool: return is_match +def filtered_nodes( + nodes: Iterable[Node], ops: Union[OperatorLike, Iterable[OperatorLike]] +) -> Iterable[Node]: + """Iterate over nodes that are filtered by the given operations. + + This utility function simplifies the common pattern of iterating through nodes + and filtering by operation type. + + Args: + nodes: Iterable of nodes to filter (e.g., gm.graph.nodes) + ops: Operation(s) to match against + + Yields: + Node: Nodes that match the given operations + + Example: + # Instead of: + for node in gm.graph.nodes: + if not is_op(node, torch.ops.aten.linear): + continue + # process node + + # Use: + for node in filtered_nodes(gm.graph.nodes, torch.ops.aten.linear): + # process node + """ + for node in nodes: + if is_op(node, ops): + yield node + + def is_linear_op(node: Node, include_quantization: bool = False) -> bool: """Check if the node is a linear op. diff --git a/tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher.py b/tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher.py index 011dfd33cb05..28e195b41ebb 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher.py +++ b/tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher.py @@ -30,7 +30,7 @@ ) from torch.fx import GraphModule -from tensorrt_llm._torch.auto_deploy.transformations.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm @contextlib.contextmanager diff --git a/tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py index 5b6acb6dafc6..f2075845187e 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py +++ b/tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py @@ -1,4 +1,5 @@ -from typing import Dict, List, Tuple, Union +from fnmatch import fnmatch +from typing import Dict, List, Optional, Tuple, Union import torch import torch.nn.functional as F @@ -12,7 +13,9 @@ ) from .logger import ad_logger from .node_utils import ( + extract_param_names_from_lin_node, get_quantization_params_from_linear_node, + is_bmm_op, is_linear_op, is_op, modelopt_dynamic_block_quantize_op, @@ -20,7 +23,7 @@ ) try: - from ...quantization.utils import float4_sf_dtype + from ....quantization.utils.fp4_utils import float4_sf_dtype except ImportError: float4_sf_dtype = None @@ -83,6 +86,7 @@ def create(quant_type_or_node: Union[str, Node], is_bmm: bool = False): quantization_impl_map = { "": None, "FP8": FP8QuantizationImpl, + "NVFP4": FP4QuantizationImpl, } return quantization_impl_map[quant_type_or_node] @@ -461,3 +465,48 @@ def post_load_hook(module, incompatible_keys, weight_name): attr_name, torch.nn.Parameter(param_cm, requires_grad=param.requires_grad), ) + + +def should_skip_quantization( + node_or_name: Union[Node, str], + excluded_patterns: list[str], +) -> bool: + """Check if a node or parameter name should be skipped based on excluded patterns.""" + if isinstance(node_or_name, str): + modname, _, _ = node_or_name.rpartition(".") + else: + if not (is_linear_op(node_or_name, include_quantization=False) or is_bmm_op(node_or_name)): + return True + param_name, _ = extract_param_names_from_lin_node(node_or_name) + modname, _, _ = param_name.rpartition(".") + + return any(fnmatch(modname, pattern) for pattern in excluded_patterns) + + +def extract_scales_from_node(node: Node, scale_names: list[str]) -> Dict[str, Optional[Node]]: + """ + Extracts scale tensors from node.args/kwargs using a fixed list of expected scale names. + """ + scales = {} + args = list(node.args) + + # Try kwargs first + for i, name in enumerate(scale_names): + scales[name] = node.kwargs.get(name, None) + + # Fallback to positional args (starting after input, weight, bias) + for i, name in enumerate(scale_names): + if scales[name] is None and len(args) > 3 + i: + scales[name] = args[3 + i] + + return scales + + +def get_scales_and_type_from_node(node: Node) -> Tuple[Dict[str, Node], str]: + """Returns a dict of scale args and quantization type string ('fp4', 'fp8', etc).""" + for qtype in [FP4QuantizationImpl, FP8QuantizationImpl]: + if is_op(node, qtype.target_op()): + return extract_scales_from_node( + node, qtype.scale_names() + ), qtype.__name__.lower().replace("quantizationimpl", "") + return None, "simple" diff --git a/tensorrt_llm/_torch/compilation/backend.py b/tensorrt_llm/_torch/compilation/backend.py index 1e06d553dc6b..f6e7ae64905d 100644 --- a/tensorrt_llm/_torch/compilation/backend.py +++ b/tensorrt_llm/_torch/compilation/backend.py @@ -12,9 +12,9 @@ import tensorrt_llm from tensorrt_llm import logger -from .patterns.ar_residual_norm import register_ar_residual_norm +from .multi_stream.auto_multi_stream import multi_stream_schedule +from .patterns.ar_residual_norm import register_ar_fusions from .patterns.residual_add_norm import register_add_norm -from .patterns.ub_allreduce import register_ub_patterns from .piecewise_optimizer import piecewise_optimizer from .recover_pass import recover_pass from .remove_copy_pass import remove_copy_for_mutates_args @@ -25,12 +25,20 @@ class Backend: _custom_pass_instances: List[PatternMatcherPass] = None _graph_pool_handle: tuple[int, int] = None + # Following classes are used to let weakref ref the stream and eventlist objects. + class Streams(list): + pass + + class Events(list): + pass + def __init__( self, enable_inductor=True, enable_userbuffers=False, enable_piecewise_cuda_graph: bool = False, cuda_graph_batch_sizes: Optional[List[int]] = None, + max_num_streams: int = 1, ) -> None: super().__init__() self.elapsed_time = 0 @@ -45,6 +53,10 @@ def __init__( else []) self.piecewise_cuda_graph = enable_piecewise_cuda_graph self.no_optimization = False + # We only need to create aux streams. + self.aux_streams = Backend.Streams( + [torch.cuda.Stream() for i in range(max_num_streams - 1)]) + self.events = Backend.Events() inductor_config.enable_auto_functionalized_v2 = False if Backend._graph_pool_handle is None: @@ -63,10 +75,9 @@ def get_custom_pass(cls, enable_userbuffers): # Currently torch compile cannot work properly with lamport fusion kernel # TO-DO: Fix this issue os.environ["DISABLE_LAMPORT_REDUCE_NORM_FUSION"] = "1" - register_ar_residual_norm(cls._custom_pass_instances[0]) - if enable_userbuffers and tensorrt_llm.bindings.internal.userbuffers.ub_supported( - ): - register_ub_patterns(cls._custom_pass_instances) + ub_enabled = enable_userbuffers and tensorrt_llm.bindings.internal.userbuffers.ub_supported( + ) + register_ar_fusions(cls._custom_pass_instances, ub_enabled) else: register_add_norm(cls._custom_pass_instances[0]) return cls._custom_pass_instances @@ -77,6 +88,12 @@ def bypass_optimization(self): def enable_optimization(self): self.no_optimization = False + def generate_events(self, num_events: int): + if num_events > len(self.events): + self.events += [ + torch.cuda.Event() for _ in range(num_events - len(self.events)) + ] + def optimize( self, gm: GraphModule, @@ -90,17 +107,30 @@ def optimize( graph.eliminate_dead_code() # After this pass, cannot run any dce!!! remove_copy_for_mutates_args(graph) + + # Do not apply multi-stream if enable piecewise cuda graph or inductor + # For piecewise cuda graph, we will apply the multi-stream optimization in piecewise_optimizer + # For inductor, we do not control the passes inside inductor. + if len( + self.aux_streams + ) > 0 and not self.piecewise_cuda_graph and not self.enable_inductor: + num_events = multi_stream_schedule(gm, len(self.aux_streams) + 1) + self.generate_events(num_events) + gm.recompile() if self.piecewise_cuda_graph: - return piecewise_optimizer( + gm, num_events = piecewise_optimizer( gm, example_inputs, self.enable_inductor, self.input_num_tokens, self.cuda_graph_batch_sizes, self._graph_pool_handle, + len(self.aux_streams) + 1, ) + self.generate_events(num_events) + return gm elif self.enable_inductor: return compile_fx(gm, example_inputs) else: diff --git a/tensorrt_llm/_torch/compilation/multi_stream/__init__.py b/tensorrt_llm/_torch/compilation/multi_stream/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py b/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py new file mode 100644 index 000000000000..c2d3cf012a05 --- /dev/null +++ b/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py @@ -0,0 +1,456 @@ +import time +from dataclasses import dataclass, field +from operator import getitem +from queue import PriorityQueue +from typing import Dict, List + +import torch +from torch.fx import Graph, GraphModule, Node + +from tensorrt_llm.logger import logger + +from ..utils import inplace_info + + +def is_symint_node(node: Node) -> bool: + if node is not None and 'val' in node.meta: + # This is a symint call that happens on host. No need to count time on stream. + if isinstance(node.meta['val'], torch.SymInt): + return True + return False + + +def estimate_time(node: Node) -> int: + if node is None: + return 0 + if is_symint_node(node): + # This is a symint call that happens on host. No need to count time on stream. + return 0 + + # Add cost model for ops that need special handling. + # We can start with rough estimation and refine it later. + + no_cost_ops = { + getitem, torch.ops.aten.view.default, torch.ops.aten.view.dtype, + torch.ops.aten.alias.default, torch.ops.aten.empty.memory_format, + torch.ops.aten.permute.default + } + + moe_ops = { + torch.ops.trtllm.fp4_block_scale_moe_runner.default, + torch.ops.trtllm.fused_moe.default, + } + + gemm_ops = { + torch.ops.aten.mm.default, + torch.ops.trtllm.nvfp4_gemm.default, + torch.ops.trtllm.fp8_batched_gemm_trtllmgen.default, + torch.ops.trtllm.w4a8_mxfp4_fp8_gemm.default, + torch.ops.trtllm.finegrained_mixed_dtype_gemm.default, + torch.ops.trtllm.bmm_out.default, + torch.ops.trtllm.cublas_scaled_mm.default, + torch.ops.trtllm.cublas_mm.default, + torch.ops.trtllm.dsv3_router_gemm_op.default, + torch.ops.trtllm.dsv3_fused_a_gemm_op.default, + torch.ops.trtllm.fp4_gemm.default, + torch.ops.trtllm.fp4_bmm.default, + torch.ops.trtllm.fp8_block_scaling_gemm.default, + torch.ops.trtllm.matmul_to_ub.default, + } + + # These ops are not counted in the time estimation. + if node.op == "call_function" and node.target in no_cost_ops: + return 0 + + # Add estimation below. With accurate estimation, the stream assignment + # can give the best performance. But it is hard to get accurate estimation. + # + # So currently, these estimations are not accurate. They just make sure the key path + # is correctly scheduled. Adjust the estimation or add new ones + # if the stream assignment is not desired. + + MOE_OP_COST = 20 + GEMM_OP_COST = 10 + DEFAULT_OP_COST = 1 + + # Adjust MOE weight to make the router -> MOE key path + if node.op == "call_function" and node.target in moe_ops: + return MOE_OP_COST + + # GEMM ops + if node.op == "call_function" and node.target in gemm_ops: + return GEMM_OP_COST + + # Refine the estimation of time for nodes. + return DEFAULT_OP_COST + + +@dataclass +class Stream: + # Stream id + id: int + + # Nodes running on the stream + nodes: List['MultiStreamNode'] = field(init=False, default_factory=list) + + # Current elapsed time of the stream + current_time: int = field(init=False, default=0) + + +class MultiStreamNode: + + def __init__(self, node: Node, in_edges: Dict[Node, 'MultiStreamNode']): + # The node in the original graph + self.node = node + + # The distance to the exit of DAG + self.distance = 0 + + # Weight for the node which represents the computation cost + self.weight = estimate_time(node) + + # The in edges of the node + self.in_edges = in_edges + + # The out edges of the node + self.out_edges = [] + + # end time of the node + self.end_time = 0 + + # Assigned stream for the node + self.stream = None + + # wait on events + self.wait_on = [] + + # trigger event + self.event = None + + +class MultiStreamDAG: + + def __init__(self, gm: GraphModule): + self.gm = gm + self.node_to_id = {} + self.node_in_degrees = {} + self.output_nodes = [] + self.placeholders = [] + self.nodes = {} + self.in_degrees = {} + self.work_list = [] + self.entry_node = None + self.exit_node = None + + self.create_dag_from_gm(gm) + assert self.entry_node is not None + assert self.exit_node is not None + + def create_dag_from_gm(self, gm: GraphModule) -> None: + """ + Create a DAG from the graph module. + """ + # Create node to id mapping + for node in gm.graph.nodes: + self.node_to_id[node] = len(self.node_to_id) + + # Fake entry node. + # All nodes without in edges will be connected to this node. + self.entry_node = MultiStreamNode(None, dict()) + + latest_inplace_stat = {} + inplace_map = inplace_info() + + def flatten_args(args): + """Recursively flatten nested arguments into a flat list.""" + args_new = [] + stack = list(args) + while stack: + arg = stack.pop() + if isinstance(arg, dict): + stack.extend(arg.values()) + elif isinstance(arg, (list, tuple)): + stack.extend(arg) + else: + args_new.append(arg) + return args_new + + # Pop all the placeholders from gm + # We know that the node is already in topological order + for node in gm.graph.nodes: + # We assume that all the placeholders are already synced with the base stream + if node.op == "placeholder": + self.placeholders.append(node) + continue + + args = flatten_args([a for a in node.args] + + [a for a in node.kwargs.values()]) + + in_edges = dict() + for arg in args: + if arg in latest_inplace_stat: + in_edges[arg] = latest_inplace_stat[arg] + elif isinstance(arg, torch.fx.Node) and arg.op != "placeholder": + in_edges[arg] = self.nodes[arg] + + # For node without in edge, connect it to the entry + if len(in_edges) == 0: + in_edges[None] = self.entry_node + + vertex = MultiStreamNode(node, in_edges) + if node.op == "output": + self.exit_node = vertex + vertex.distance = 0 + self.nodes[node] = vertex + self.in_degrees[vertex] = len(in_edges) + if node.op == "call_function": + func = node.target + if func in inplace_map: + for inplace_arg in inplace_map[func].values(): + # At this stage, all inplace op must be using kwargs for all params + assert inplace_arg in node.kwargs + latest_inplace_stat[node.kwargs[inplace_arg]] = vertex + + for edge in in_edges.values(): + edge.out_edges.append(vertex) + self.compute_distance() + + def compute_distance(self) -> None: + """ + Compute the distance to the exit node for each node. + """ + # Reverse topological sort to compute distance to exit node + work_list = [self.exit_node] + out_degrees = { + node: len(node.out_edges) + for node in self.nodes.values() + } + out_degrees[self.entry_node] = len(self.entry_node.out_edges) + + while len(work_list) > 0: + node = work_list.pop() + for in_edge in node.in_edges.values(): + out_degrees[in_edge] -= 1 + in_edge.distance = max(in_edge.distance, + node.weight + node.distance) + if out_degrees[in_edge] == 0: + work_list.append(in_edge) + + def assign_streams(self, max_num_streams: int) -> int: + """ + Assign streams to the nodes in the DAG. + Return the number of events created. + """ + worklist = PriorityQueue() + num_nodes = len(self.node_to_id) + + # When accessing node, the distance to the exit node is main priority. + # The node with largest distance means currently this is the bottleneck of the whole graph. + def calc_priority(node_id: int, distance: int) -> int: + # We keep the node order by default. + # It also gives deterministic order for priority queue. + return (-distance) * num_nodes + node_id + + streams = [Stream(i) for i in range(max_num_streams)] + + def pick_stream(start_time, node) -> Stream: + if node.weight == 0: + # This is a symint node or a getitem node. + # It always assigns to the stream that produce the node. + for n in node.in_edges.values(): + if is_symint_node(n.node): + continue + return n.stream + return streams[0] + + closest_stream = None + least_time = float('inf') + for st in streams: + if st.current_time <= start_time: + return st + else: + if st.current_time < least_time: + least_time = st.current_time + closest_stream = st + return closest_stream + + # We just start from the out_edges of the entry node. Entry node is just a fake node + # For entry, we assign to the primary stream. + self.entry_node.stream = streams[0] + streams[0].nodes.append(self.entry_node) + for out_edge in self.entry_node.out_edges: + worklist.put((calc_priority(self.node_to_id[out_edge.node], + out_edge.distance), out_edge)) + + sync_event_id = 0 + + while not worklist.empty(): + _, node = worklist.get() + assert node.stream is None + + # Get when current node can start. + # Start time is the max of the end time of all the in edges. + start_time = max( + [in_edge.end_time for in_edge in node.in_edges.values()]) + node.stream = pick_stream(start_time, node) + node.end_time = max(start_time, + node.stream.current_time) + node.weight + node.stream.current_time = node.end_time + node.stream.nodes.append(node) + + for in_edge_tensor, in_edge in node.in_edges.items(): + if in_edge.stream != node.stream and not is_symint_node( + in_edge.node): + if in_edge.event is None: + in_edge.event = sync_event_id + sync_event_id += 1 + node.wait_on.append((in_edge, in_edge_tensor)) + + # Now, for any in edge running on different stream, we need to create a sync event. + for out_edge in node.out_edges: + self.in_degrees[out_edge] -= 1 + if self.in_degrees[out_edge] == 0: + worklist.put((calc_priority(self.node_to_id[out_edge.node], + out_edge.distance), out_edge)) + self.streams = streams + return sync_event_id + + def create_new_graph(self) -> Graph: + """ + Create new graph with the nodes assigned to the streams. + """ + # Now each node should have been assigned a stream. We will now create a new graph and insert all nodes + # As torch need to create node for switching stream, need to group nodes as much as possible. + remap = {} + new_graph = Graph() + + for st in self.streams: + logger.debug(f"{len(st.nodes)} nodes running on stream {st.id}") + + # First, push all placeholders to the new graph. + for placeholder in self.placeholders: + remap[placeholder] = new_graph.node_copy(placeholder, + lambda n: remap[n]) + + # Then, we will push all the nodes into the new graph. + # Build in_degrees again as we need to check whether a stream is ready to run. + self.in_degrees = { + node: len(node.in_edges) + for node in self.nodes.values() + } + self.in_degrees[self.entry_node] = 0 + + stream_pos = [0] * len(self.streams) + + def has_more_nodes() -> bool: + for st in self.streams: + if len(st.nodes) > stream_pos[st.id]: + return True + return False + + last_stream = 0 + + # The nodes in stream are already in topological order. + while has_more_nodes(): + for st in self.streams: + if len(st.nodes) == stream_pos[st.id]: + continue + node = st.nodes[stream_pos[st.id]] + if self.in_degrees[node] != 0: + # This stream is not ready to run now. + continue + + # Any time the stream is changed, set the stream. + if node.stream.id != last_stream: + # Change stream + new_graph.create_node("call_function", + torch.ops.trtllm.set_stream, + args=(node.stream.id, )) + last_stream = node.stream.id + + for _ in range(stream_pos[st.id], len(st.nodes)): + node = st.nodes[stream_pos[st.id]] + if self.in_degrees[node] != 0: + break + for out_edge in node.out_edges: + self.in_degrees[out_edge] -= 1 + stream_pos[st.id] += 1 + # It could be the fake entry node. + if node.node is not None: + # Wait on all the events that the node is waiting on. + for wait in node.wait_on: + new_graph.create_node("call_function", + torch.ops.trtllm.wait_event, + args=(wait[0].event, )) + remap[node.node] = new_graph.node_copy( + node.node, lambda n: remap[n]) + for wait in node.wait_on: + # wait[1] is the actual tensor that the op is waiting on. + # Need to record stream for that tensor. + if wait[1] is None: + continue + new_graph.create_node( + "call_function", + torch.ops.trtllm.record_stream, + args=(remap[wait[1]], st.id)) + if node.event is not None: + new_graph.create_node("call_function", + torch.ops.trtllm.record_event, + args=(node.event, )) + + # After each handling, start again to make sure primary stream is pushed first. + break + return new_graph + + def optimize(self, max_num_streams: int) -> int: + """ + Run multistream optimize for MultiStreamDAG. The graph module that used to create the DAG will be updated. + Return the number of events created. + """ + num_events = self.assign_streams(max_num_streams) + new_graph = self.create_new_graph() + self.gm.graph = new_graph + return num_events + + +def multi_stream_schedule(gm: GraphModule, max_num_streams: int) -> int: + """ + Schedule the graph module for multi stream execution. + gm is the graph module to be scheduled. The gm will be updated by this function. + max_num_streams is the maximum number of streams to be used. The scheduler may not use all the streams. + Return the number of events created. + """ + dag = MultiStreamDAG(gm) + return dag.optimize(max_num_streams) + + +# Following code is for debug purpose. Use print_dag_to_dot to print a MultiStreamDAG to dot file. + + +def dump_dag_as_dot(dag: MultiStreamDAG, max_num_nodes: int = 500) -> None: + COLORS = [ + "red", "chocolate", "cyan", "gold", "coral", "green", "blue", "orange", + "purple", "brown" + ] + filename = f"dag_{int(time.time())}.dot" + with open(filename, 'w') as f: + f.write("digraph G {\n") + f.write( + f"id_entry [label=\"node=entry, distance={dag.entry_node.distance}\"]\n" + ) + cnt = 0 + for node in dag.nodes.values(): + color = "white" if node.stream is None else COLORS[node.stream.id] + f.write( + f"id_{dag.node_to_id[node.node]} [label=\"node={node.node}, " + f"distance={node.distance}, weight={node.weight}\", " + f"color={color}, shape=oval]\n") + for in_edge in node.in_edges.values(): + id = str(dag.node_to_id[ + in_edge.node]) if in_edge.node is not None else "entry" + f.write(f"id_{id} -> id_{dag.node_to_id[node.node]}\n") + if cnt > max_num_nodes: + break + cnt += 1 + f.write("}\n") + f.flush() diff --git a/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py b/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py index 411eed4bdc93..afbaa0949df3 100644 --- a/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py +++ b/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py @@ -1,4 +1,5 @@ from operator import getitem +from typing import List, Optional import torch from torch._inductor.pattern_matcher import (MULTIPLE, CallFunction, Ignored, @@ -9,7 +10,7 @@ import tensorrt_llm -from ...distributed import AllReduceFusionOp +from ...distributed import AllReduceFusionOp, AllReduceStrategy aten = torch.ops.aten from tensorrt_llm.mapping import Mapping @@ -95,3 +96,637 @@ def extra_check(match: Match) -> bool: search_fn_pattern=ar_residual_norm_pattern, extra_check=extra_check, ) + + +def check_f16_bf16_input(match, input_node) -> bool: + input = match.ctx.pattern_to_node[input_node] + if not isinstance(input, torch.fx.graph.Node): + return False + dtype = input.meta["tensor_meta"].dtype + if dtype != torch.float16 and dtype != torch.bfloat16: + return False + return True + + +def check_non_ub_strategy(match, strategy_node) -> bool: + strategy = match.ctx.pattern_to_node[strategy_node] + if not isinstance(strategy, int): + return False + if strategy == int(AllReduceStrategy.UB): + return False + return True + + +def register_ar_residual_norm_out_fp8_quant(custom_pass: PatternMatcherPass): + # TODO: add pp + tp support + mapping = Mapping( + world_size=tensorrt_llm.mpi_world_size(), + tp_size=tensorrt_llm.mpi_world_size(), + rank=tensorrt_llm.mpi_rank(), + ) + + input_node = KeywordArg("input") + strategy_node = KeywordArg("strategy") + allreduce_default = CallFunction(torch.ops.trtllm.allreduce.default, + input_node, + KeywordArg("residual"), + KeywordArg("gamma"), + None, + None, + KeywordArg("workspace"), + mapping.tp_group, + strategy_node, + int(AllReduceFusionOp.RESIDUAL_RMS_NORM), + KeywordArg("eps"), + KeywordArg("trigger_completion_at_end"), + _users=2) + getitem_0 = CallFunction(getitem, allreduce_default, 0, _users=2) + getitem_1 = CallFunction(getitem, allreduce_default, 1) + static_quantize_e4m3_per_tensor_default = CallFunction( + torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor.default, + getitem_0, + KeywordArg("scale"), + _users=2) + getitem_2 = CallFunction(getitem, + static_quantize_e4m3_per_tensor_default, + 0, + _users=2) + getitem_3 = CallFunction(getitem, static_quantize_e4m3_per_tensor_default, + 1) + pattern = MultiOutputPattern([getitem_0, getitem_1, getitem_2, getitem_3 + ]) # norm_out, residual_out, quant_out, scale + + def empty_pattern( + input: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + workspace: torch.LongTensor, + strategy: int, + eps: float, + scale: torch.Tensor, + trigger_completion_at_end: bool, + ): + return + + def target_pattern( + input: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + workspace: torch.LongTensor, + strategy: int, + eps: float, + scale: torch.Tensor, + trigger_completion_at_end: bool, + ): + allreduce = torch.ops.trtllm.allreduce( + input, residual, gamma, scale, None, workspace, mapping.tp_group, + int(strategy), + int(AllReduceFusionOp.RESIDUAL_RMS_NORM_OUT_QUANT_FP8), float(eps), + trigger_completion_at_end) + return allreduce[0], allreduce[2], allreduce[1], scale + + def extra_check(match: Match) -> bool: + return check_f16_bf16_input( + match, input_node) and check_non_ub_strategy(match, strategy_node) + + register_replacement( + empty_pattern, + target_pattern, + [], + fwd_only, + custom_pass, + search_fn_pattern=pattern, + extra_check=extra_check, + ) + + +def register_ar_residual_norm_fp8_quant(custom_pass: PatternMatcherPass): + # TODO: add pp + tp support + mapping = Mapping( + world_size=tensorrt_llm.mpi_world_size(), + tp_size=tensorrt_llm.mpi_world_size(), + rank=tensorrt_llm.mpi_rank(), + ) + + input_node = KeywordArg("input") + strategy_node = KeywordArg("strategy") + allreduce_default = CallFunction(torch.ops.trtllm.allreduce.default, + input_node, + KeywordArg("residual"), + KeywordArg("gamma"), + None, + None, + KeywordArg("workspace"), + mapping.tp_group, + strategy_node, + int(AllReduceFusionOp.RESIDUAL_RMS_NORM), + KeywordArg("eps"), + KeywordArg("trigger_completion_at_end"), + _users=2) + getitem_0 = CallFunction(getitem, allreduce_default, 0) + getitem_1 = CallFunction(getitem, allreduce_default, 1) + static_quantize_e4m3_per_tensor_default = CallFunction( + torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor.default, + getitem_0, + KeywordArg("scale"), + _users=2) + getitem_2 = CallFunction(getitem, + static_quantize_e4m3_per_tensor_default, + 0, + _users=2) + getitem_3 = CallFunction(getitem, static_quantize_e4m3_per_tensor_default, + 1) + pattern = MultiOutputPattern([getitem_1, getitem_2, + getitem_3]) # residual_out, quant_out, scale + + def empty_pattern( + input: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + workspace: torch.LongTensor, + strategy: int, + eps: float, + scale: torch.Tensor, + trigger_completion_at_end: bool, + ): + return + + def target_pattern( + input: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + workspace: torch.LongTensor, + strategy: int, + eps: float, + scale: torch.Tensor, + trigger_completion_at_end: bool, + ): + allreduce = torch.ops.trtllm.allreduce( + input, residual, gamma, scale, None, workspace, mapping.tp_group, + int(strategy), int(AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8), + float(eps), trigger_completion_at_end) + return allreduce[1], allreduce[0], scale + + def extra_check(match: Match) -> bool: + return check_f16_bf16_input( + match, input_node) and check_non_ub_strategy(match, strategy_node) + + register_replacement( + empty_pattern, + target_pattern, + [], + fwd_only, + custom_pass, + search_fn_pattern=pattern, + extra_check=extra_check, + ) + + +def register_ar_residual_norm_out_fp4_quant(custom_pass: PatternMatcherPass): + # TODO: add pp + tp support + mapping = Mapping( + world_size=tensorrt_llm.mpi_world_size(), + tp_size=tensorrt_llm.mpi_world_size(), + rank=tensorrt_llm.mpi_rank(), + ) + + input_node = KeywordArg("input") + strategy_node = KeywordArg("strategy") + allreduce_default = CallFunction(torch.ops.trtllm.allreduce.default, + input_node, + KeywordArg("residual"), + KeywordArg("gamma"), + None, + None, + KeywordArg("workspace"), + mapping.tp_group, + strategy_node, + int(AllReduceFusionOp.RESIDUAL_RMS_NORM), + KeywordArg("eps"), + KeywordArg("trigger_completion_at_end"), + _users=2) + getitem_0 = CallFunction(getitem, allreduce_default, 0, _users=2) + getitem_1 = CallFunction(getitem, allreduce_default, 1) + fp4_quant_default = CallFunction(torch.ops.trtllm.fp4_quantize.default, + getitem_0, + KeywordArg("scale"), + 16, + _users=2) + getitem_2 = CallFunction(getitem, fp4_quant_default, 0, _users=2) + getitem_3 = CallFunction(getitem, fp4_quant_default, 1) + pattern = MultiOutputPattern([getitem_0, getitem_1, getitem_2, getitem_3]) + + def empty_pattern( + input: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + workspace: torch.LongTensor, + strategy: int, + eps: float, + scale: torch.Tensor, + trigger_completion_at_end: bool, + ): + return + + def target_pattern( + input: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + workspace: torch.LongTensor, + strategy: int, + eps: float, + scale: torch.Tensor, + trigger_completion_at_end: bool, + ): + allreduce = torch.ops.trtllm.allreduce( + input, residual, gamma, scale, None, workspace, mapping.tp_group, + int(strategy), + int(AllReduceFusionOp.RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4), + float(eps), trigger_completion_at_end) + return allreduce[0], allreduce[3], allreduce[1], allreduce[2] + + def extra_check(match: Match) -> bool: + return check_f16_bf16_input( + match, input_node) and check_non_ub_strategy(match, strategy_node) + + register_replacement( + empty_pattern, + target_pattern, + [], + fwd_only, + custom_pass, + search_fn_pattern=pattern, + extra_check=extra_check, + ) + + +def register_ar_residual_norm_fp4_quant(custom_pass: PatternMatcherPass): + # TODO: add pp + tp support + mapping = Mapping( + world_size=tensorrt_llm.mpi_world_size(), + tp_size=tensorrt_llm.mpi_world_size(), + rank=tensorrt_llm.mpi_rank(), + ) + + input_node = KeywordArg("input") + strategy_node = KeywordArg("strategy") + allreduce_default = CallFunction(torch.ops.trtllm.allreduce.default, + input_node, + KeywordArg("residual"), + KeywordArg("gamma"), + None, + None, + KeywordArg("workspace"), + mapping.tp_group, + strategy_node, + int(AllReduceFusionOp.RESIDUAL_RMS_NORM), + KeywordArg("eps"), + KeywordArg("trigger_completion_at_end"), + _users=2) + getitem_0 = CallFunction(getitem, allreduce_default, 0) + getitem_1 = CallFunction(getitem, allreduce_default, 1) + fp4_quant_default = CallFunction(torch.ops.trtllm.fp4_quantize.default, + getitem_0, + KeywordArg("scale"), + 16, + _users=2) + getitem_2 = CallFunction(getitem, fp4_quant_default, 0, _users=2) + getitem_3 = CallFunction(getitem, fp4_quant_default, 1) + pattern = MultiOutputPattern([getitem_1, getitem_2, getitem_3]) + + def empty_pattern( + input: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + workspace: torch.LongTensor, + strategy: int, + eps: float, + scale: torch.Tensor, + trigger_completion_at_end: bool, + ): + return + + def target_pattern( + input: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + workspace: torch.LongTensor, + strategy: int, + eps: float, + scale: torch.Tensor, + trigger_completion_at_end: bool, + ): + allreduce = torch.ops.trtllm.allreduce( + input, residual, gamma, scale, None, workspace, mapping.tp_group, + int(strategy), int(AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4), + float(eps), trigger_completion_at_end) + return allreduce[2], allreduce[0], allreduce[1] + + def extra_check(match: Match) -> bool: + return check_f16_bf16_input( + match, input_node) and check_non_ub_strategy(match, strategy_node) + + register_replacement( + empty_pattern, + target_pattern, + [], + fwd_only, + custom_pass, + search_fn_pattern=pattern, + extra_check=extra_check, + ) + + +def register_ub_patterns(custom_passes: List[PatternMatcherPass]): + mapping = Mapping( + world_size=tensorrt_llm.mpi_world_size(), + tp_size=tensorrt_llm.mpi_world_size(), + rank=tensorrt_llm.mpi_rank(), + ) + + def register_convert_supported_ar_to_ub(custom_pass: PatternMatcherPass): + strategy = int(AllReduceStrategy.AUTO) + input_node = KeywordArg('input') + fusion = KeywordArg('fusion_op') + trtllm_allreduce_default = CallFunction( + torch.ops.trtllm.allreduce.default, input_node, + KeywordArg('residual_in'), KeywordArg('gamma'), KeywordArg('scale'), + None, Ignored(), mapping.tp_group, strategy, fusion, + KeywordArg('eps'), Ignored()) + + def empty_convert_supported_ar_to_ub( + input: torch.Tensor, + residual_in: torch.Tensor, + gamma: torch.Tensor, + scale: Optional[torch.Tensor], + fusion_op: int, + eps: float, + ): + return + + def target_convert_supported_ar_to_ub( + input: torch.Tensor, + residual_in: torch.Tensor, + gamma: torch.Tensor, + scale: Optional[torch.Tensor], + fusion_op: int, + eps: float, + ): + input = torch.ops.trtllm.copy_to_userbuffers(input) + all_reduce_output = torch.ops.trtllm.allreduce( + input, residual_in, gamma, scale, None, None, mapping.tp_group, + int(AllReduceStrategy.UB), fusion_op, eps, False) + finalize_output = torch.ops.trtllm.userbuffers_allreduce_finalize( + all_reduce_output[-1], False) + all_reduce_output[-1] = finalize_output + return all_reduce_output + + def extra_check_convert_supported_ar_to_ub(match: Match) -> bool: + if not check_f16_bf16_input(match, input_node): + return False + + fusion_value = match.ctx.pattern_to_node[fusion] + if not isinstance(fusion_value, int): + return False + if fusion_value != int( + AllReduceFusionOp.RESIDUAL_RMS_NORM + ) and fusion_value != int( + AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8 + ) and fusion_value != int( + AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4): + return False + + return True + + register_replacement( + empty_convert_supported_ar_to_ub, + target_convert_supported_ar_to_ub, + [], + fwd_only, + custom_pass, + search_fn_pattern=trtllm_allreduce_default, + extra_check=extra_check_convert_supported_ar_to_ub, + ) + + def register_ub_prologue_patterns(custom_pass: PatternMatcherPass): + + def register_scaled_mm_prologue(custom_pass: PatternMatcherPass): + trtllm_cublas_scaled_mm_default = CallFunction( + torch.ops.trtllm.cublas_scaled_mm.default, KeywordArg('mm0_a'), + KeywordArg('mm0_b'), KeywordArg('mm0_a_scale'), + KeywordArg('mm0_b_scale'), KeywordArg('mm0_bias'), + KeywordArg('mm_dtype')) + ub_copy = CallFunction(torch.ops.trtllm.copy_to_userbuffers, + trtllm_cublas_scaled_mm_default) + + def empty_scaled_mm_prologue_pattern( + mm0_a: torch.Tensor, + mm0_b: torch.Tensor, + mm0_a_scale: torch.Tensor, + mm0_b_scale: torch.Tensor, + mm0_bias: Optional[torch.Tensor], + mm_dtype: torch.dtype, + ): + return + + def target_scaled_mm_prologue_pattern( + mm0_a: torch.Tensor, + mm0_b: torch.Tensor, + mm0_a_scale: torch.Tensor, + mm0_b_scale: torch.Tensor, + mm0_bias: Optional[torch.Tensor], + mm_dtype: torch.dtype, + ): + scaled_mm_output = torch.ops.trtllm.cublas_scaled_mm( + mm0_a, mm0_b, mm0_a_scale, mm0_b_scale, mm0_bias, mm_dtype, + True) + return scaled_mm_output + + # No extra check needed as the output dtype of scaled_mm has been verified when + # ub_copy is inserted. + register_replacement( + empty_scaled_mm_prologue_pattern, + target_scaled_mm_prologue_pattern, + [], + fwd_only, + custom_pass, + search_fn_pattern=ub_copy, + ) + + def register_nvfp4_gemm_prologue(custom_pass: PatternMatcherPass): + trtllm_nvfp4_gemm_default = CallFunction( + torch.ops.trtllm.nvfp4_gemm.default, KeywordArg('act_fp4'), + KeywordArg('weight'), KeywordArg('act_sf'), + KeywordArg('weight_scale'), KeywordArg('alpha'), + KeywordArg('output_dtype')) + ub_copy = CallFunction(torch.ops.trtllm.copy_to_userbuffers, + trtllm_nvfp4_gemm_default) + + def empty_nvfp4_gemm_prologue_pattern( + act_fp4: torch.Tensor, + weight: torch.Tensor, + act_sf: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + ): + return + + def target_nvfp4_gemm_prologue_pattern( + act_fp4: torch.Tensor, + weight: torch.Tensor, + act_sf: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + ): + nvfp4_gemm_output = torch.ops.trtllm.nvfp4_gemm( + act_fp4, weight, act_sf, weight_scale, alpha, output_dtype, + True) + return nvfp4_gemm_output + + # No extra check needed as the output dtype of nvfp4_gemm has been verified when + # ub_copy is inserted. + register_replacement( + empty_nvfp4_gemm_prologue_pattern, + target_nvfp4_gemm_prologue_pattern, + [], + fwd_only, + custom_pass, + search_fn_pattern=ub_copy, + ) + + def register_mm_prologue(custom_pass: PatternMatcherPass): + aten_mm_default = CallFunction(aten.mm.default, KeywordArg('mm0_a'), + KeywordArg('mm0_b')) + ub_copy = CallFunction(torch.ops.trtllm.copy_to_userbuffers, + aten_mm_default) + + def empty_mm_prologue_pattern( + mm0_a: torch.Tensor, + mm0_b: torch.Tensor, + ): + return + + def target_mm_prologue_pattern( + mm0_a: torch.Tensor, + mm0_b: torch.Tensor, + ): + mm_output = torch.ops.trtllm.matmul_to_ub(mm0_a, mm0_b) + return mm_output + + # No extra check needed as the output dtype of mm has been verified when + # ub_copy is inserted. + register_replacement( + empty_mm_prologue_pattern, + target_mm_prologue_pattern, + [], + fwd_only, + custom_pass, + search_fn_pattern=ub_copy, + ) + + def register_add_prologue(custom_pass: PatternMatcherPass): + aten_add_default = CallFunction(aten.add.Tensor, + KeywordArg('add_a'), + KeywordArg('add_b')) + ub_copy = CallFunction(torch.ops.trtllm.copy_to_userbuffers, + aten_add_default) + + def empty_add_prologue_pattern( + add_a: torch.Tensor, + add_b: torch.Tensor, + ): + return + + def target_add_prologue_pattern( + add_a: torch.Tensor, + add_b: torch.Tensor, + ): + add_output = torch.ops.trtllm.add_to_ub(add_a, add_b) + return add_output + + # No extra check needed as the output dtype of add has been verified when + # ub_copy is inserted. + register_replacement( + empty_add_prologue_pattern, + target_add_prologue_pattern, + [], + fwd_only, + custom_pass, + search_fn_pattern=ub_copy, + ) + + register_scaled_mm_prologue(custom_pass) + register_nvfp4_gemm_prologue(custom_pass) + register_mm_prologue(custom_pass) + register_add_prologue(custom_pass) + + def register_ub_finalize_patterns(custom_pass: PatternMatcherPass): + trtllm_userbuffers_allreduce_finalize_default = CallFunction( + torch.ops.trtllm.userbuffers_allreduce_finalize.default, + KeywordArg("sharded_residual"), False) + trtllm_allreduce_default = CallFunction( + torch.ops.trtllm.allreduce.default, KeywordArg("input"), + trtllm_userbuffers_allreduce_finalize_default, KeywordArg("gamma"), + KeywordArg("scale"), Ignored(), Ignored(), mapping.tp_group, + int(AllReduceStrategy.UB), KeywordArg("fusion_op"), + KeywordArg("eps"), Ignored()) + + def empty_finalize_pattern( + input: torch.Tensor, + sharded_residual: torch.Tensor, + gamma: torch.Tensor, + scale: Optional[torch.Tensor], + fusion_op: int, + eps: float, + ): + return + + def target_finalize_pattern( + input: torch.Tensor, + sharded_residual: torch.Tensor, + gamma: torch.Tensor, + scale: Optional[torch.Tensor], + fusion_op: int, + eps: float, + ): + all_reduce_output = torch.ops.trtllm.allreduce( + input, sharded_residual, + gamma, scale, None, None, mapping.tp_group, + int(AllReduceStrategy.UB), fusion_op, eps, False) + return all_reduce_output + + register_replacement( + empty_finalize_pattern, + target_finalize_pattern, + [], + fwd_only, + custom_pass, + search_fn_pattern=trtllm_allreduce_default, + ) + + custom_passes.append(PatternMatcherPass()) + register_convert_supported_ar_to_ub(custom_passes[-1]) + + custom_passes.append(PatternMatcherPass()) + register_ub_prologue_patterns(custom_passes[-1]) + + custom_passes.append(PatternMatcherPass()) + register_ub_finalize_patterns(custom_passes[-1]) + + +def register_ar_fusions(custom_passes: List[PatternMatcherPass], + enable_ub: bool): + register_ar_residual_norm(custom_passes[-1]) + + custom_passes.append(PatternMatcherPass()) + register_ar_residual_norm_fp8_quant(custom_passes[-1]) + register_ar_residual_norm_fp4_quant(custom_passes[-1]) + # AR-Residual-Norm-Out-Quant-X is not supported by Userbuffers kernel. + if not enable_ub: + register_ar_residual_norm_out_fp8_quant(custom_passes[-1]) + register_ar_residual_norm_out_fp4_quant(custom_passes[-1]) + + if enable_ub: + register_ub_patterns(custom_passes) diff --git a/tensorrt_llm/_torch/compilation/patterns/ub_allreduce.py b/tensorrt_llm/_torch/compilation/patterns/ub_allreduce.py deleted file mode 100644 index 54a04c17ee48..000000000000 --- a/tensorrt_llm/_torch/compilation/patterns/ub_allreduce.py +++ /dev/null @@ -1,526 +0,0 @@ -from operator import getitem -from typing import List, Optional - -import torch -from torch._inductor.pattern_matcher import (CallFunction, Ignored, KeywordArg, - Match, MultiOutputPattern, - PatternMatcherPass, fwd_only, - register_replacement) - -import tensorrt_llm - -from ...distributed import AllReduceFusionOp, AllReduceStrategy - -aten = torch.ops.aten -from tensorrt_llm.mapping import Mapping - - -def register_ub_patterns(custom_passes: List[PatternMatcherPass]): - mapping = Mapping( - world_size=tensorrt_llm.mpi_world_size(), - tp_size=tensorrt_llm.mpi_world_size(), - rank=tensorrt_llm.mpi_rank(), - ) - - def register_ub_allreduce_quantize_fusion(custom_pass: PatternMatcherPass): - strategy = int(AllReduceStrategy.AUTO) - fusion = int(AllReduceFusionOp.RESIDUAL_RMS_NORM) - - def register_fp8_quant_pattern(custom_pass: PatternMatcherPass): - input_node = KeywordArg('input') - trtllm_allreduce_default = CallFunction( - torch.ops.trtllm.allreduce.default, - input_node, - KeywordArg('residual_in'), - KeywordArg('gamma'), - None, - None, - Ignored(), - mapping.tp_group, - strategy, - fusion, - KeywordArg('eps'), - Ignored(), - _users=2) - allreduce_output = CallFunction(getitem, trtllm_allreduce_default, - 0) - residual_out = CallFunction(getitem, trtllm_allreduce_default, 1) - tensorrt_llm_static_quantize_e4m3_per_tensor_default = CallFunction( - torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor.default, - allreduce_output, - KeywordArg('scale'), - _users=2) - quant_output = CallFunction( - getitem, tensorrt_llm_static_quantize_e4m3_per_tensor_default, - 0) - scale_out = CallFunction( - getitem, tensorrt_llm_static_quantize_e4m3_per_tensor_default, - 1) - fp8_quant_pattern = MultiOutputPattern( - [quant_output, scale_out, residual_out]) - - def empty_fp8_quant_pattern( - input: torch.Tensor, - residual_in: torch.Tensor, - gamma: torch.Tensor, - eps: float, - scale: torch.Tensor, - ): - return - - def target_fp8_quant_pattern( - input: torch.Tensor, - residual_in: torch.Tensor, - gamma: torch.Tensor, - eps: float, - scale: torch.Tensor, - ): - input = torch.ops.trtllm.copy_to_userbuffers(input) - all_reduce_output = torch.ops.trtllm.allreduce( - input, residual_in, gamma, scale, None, None, - mapping.tp_group, int(AllReduceStrategy.UB), - int(AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8), eps, - True) - finalize_output = torch.ops.trtllm.userbuffers_allreduce_finalize( - all_reduce_output[1], False) - return all_reduce_output[0], scale, finalize_output - - def extra_check_fp8_quant_pattern(match: Match) -> bool: - input = match.ctx.pattern_to_node[input_node] - if not isinstance(input, torch.fx.graph.Node): - return False - dtype = input.meta["tensor_meta"].dtype - # UB only supports FP16/BF16 input - if dtype != torch.float16 and dtype != torch.bfloat16: - return False - return True - - register_replacement( - empty_fp8_quant_pattern, - target_fp8_quant_pattern, - [], - fwd_only, - custom_pass, - search_fn_pattern=fp8_quant_pattern, - extra_check=extra_check_fp8_quant_pattern, - ) - - def register_fp4_quant_pattern(custom_pass: PatternMatcherPass): - input_node = KeywordArg('input') - trtllm_allreduce_default = CallFunction( - torch.ops.trtllm.allreduce.default, - input_node, - KeywordArg('residual_in'), - KeywordArg('gamma'), - None, - Ignored(), - Ignored(), - mapping.tp_group, - strategy, - fusion, - KeywordArg('eps'), - Ignored(), - _users=2) - allreduce_output = CallFunction(getitem, trtllm_allreduce_default, - 0) - residual_out = CallFunction(getitem, trtllm_allreduce_default, 1) - tensorrt_llm_fp4_quantize_default = CallFunction( - torch.ops.trtllm.fp4_quantize.default, - allreduce_output, - KeywordArg('scale'), - 16, - _users=2) - quant_output = CallFunction(getitem, - tensorrt_llm_fp4_quantize_default, 0) - scale_out = CallFunction(getitem, tensorrt_llm_fp4_quantize_default, - 1) - fp4_quant_pattern = MultiOutputPattern( - [quant_output, scale_out, residual_out]) - - def empty_fp4_quant_pattern( - input: torch.Tensor, - residual_in: torch.Tensor, - gamma: torch.Tensor, - eps: float, - scale: torch.Tensor, - ): - return - - def target_fp4_quant_pattern( - input: torch.Tensor, - residual_in: torch.Tensor, - gamma: torch.Tensor, - eps: float, - scale: torch.Tensor, - ): - input = torch.ops.trtllm.copy_to_userbuffers(input) - all_reduce_output = torch.ops.trtllm.allreduce( - input, residual_in, gamma, scale, None, None, - mapping.tp_group, int(AllReduceStrategy.UB), - int(AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4), eps, - True) - finalize_output = torch.ops.trtllm.userbuffers_allreduce_finalize( - all_reduce_output[-1], False) - return all_reduce_output[0], all_reduce_output[ - 1], finalize_output - - def extra_check_fp4_quant_pattern(match: Match) -> bool: - input = match.ctx.pattern_to_node[input_node] - if not isinstance(input, torch.fx.graph.Node): - return False - dtype = input.meta["tensor_meta"].dtype - # UB only supports FP16/BF16 input - if dtype != torch.float16 and dtype != torch.bfloat16: - return False - return True - - register_replacement( - empty_fp4_quant_pattern, - target_fp4_quant_pattern, - [], - fwd_only, - custom_pass, - search_fn_pattern=fp4_quant_pattern, - extra_check=extra_check_fp4_quant_pattern, - ) - - register_fp8_quant_pattern(custom_pass) - register_fp4_quant_pattern(custom_pass) - - def register_convert_supported_ar_to_ub(custom_pass: PatternMatcherPass): - strategy = int(AllReduceStrategy.AUTO) - # TODO: Also handle scale once the allreduce interface does not contain - # dynamic number of tensors. - input_node = KeywordArg('input') - fusion = KeywordArg('fusion_op') - trtllm_allreduce_default = CallFunction( - torch.ops.trtllm.allreduce.default, input_node, - KeywordArg('residual_in'), KeywordArg('gamma'), KeywordArg('scale'), - None, Ignored(), mapping.tp_group, strategy, fusion, - KeywordArg('eps'), Ignored()) - convert_pattern = MultiOutputPattern([trtllm_allreduce_default]) - - def empty_convert_supported_ar_to_ub( - input: torch.Tensor, - residual_in: torch.Tensor, - gamma: torch.Tensor, - scale: torch.Tensor, - fusion_op: int, - eps: float, - ): - return - - def target_convert_supported_ar_to_ub( - input: torch.Tensor, - residual_in: torch.Tensor, - gamma: torch.Tensor, - scale: torch.Tensor, - fusion_op: int, - eps: float, - ): - input = torch.ops.trtllm.copy_to_userbuffers(input) - all_reduce_output = torch.ops.trtllm.allreduce( - input, residual_in, gamma, scale, None, None, mapping.tp_group, - int(AllReduceStrategy.UB), fusion_op, eps, True) - finalize_output = torch.ops.trtllm.userbuffers_allreduce_finalize( - all_reduce_output[-1], False) - all_reduce_output[-1] = finalize_output - return all_reduce_output - - def extra_check_convert_supported_ar_to_ub(match: Match) -> bool: - input = match.ctx.pattern_to_node[input_node] - if not isinstance(input, torch.fx.graph.Node): - return False - dtype = input.meta["tensor_meta"].dtype - # UB only supports FP16/BF16 input - if dtype != torch.float16 and dtype != torch.bfloat16: - return False - - fusion_value = match.ctx.pattern_to_node[fusion] - if not isinstance(fusion_value, int): - return False - if fusion_value != int( - AllReduceFusionOp.RESIDUAL_RMS_NORM - ) and fusion_value != int( - AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8 - ) and fusion_value != int( - AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4): - return False - - return True - - register_replacement( - empty_convert_supported_ar_to_ub, - target_convert_supported_ar_to_ub, - [], - fwd_only, - custom_pass, - search_fn_pattern=convert_pattern, - extra_check=extra_check_convert_supported_ar_to_ub, - ) - - def register_ub_prologue_patterns(custom_pass: PatternMatcherPass): - - def register_scaled_mm_prologue(custom_pass: PatternMatcherPass): - trtllm_cublas_scaled_mm_default = CallFunction( - torch.ops.trtllm.cublas_scaled_mm.default, KeywordArg('mm0_a'), - KeywordArg('mm0_b'), KeywordArg('mm0_a_scale'), - KeywordArg('mm0_b_scale'), KeywordArg('mm0_bias'), - KeywordArg('mm_dtype')) - ub_copy = CallFunction(torch.ops.trtllm.copy_to_userbuffers, - trtllm_cublas_scaled_mm_default) - scaled_mm_prologue_pattern = MultiOutputPattern([ub_copy]) - - def empty_scaled_mm_prologue_pattern( - mm0_a: torch.Tensor, - mm0_b: torch.Tensor, - mm0_a_scale: torch.Tensor, - mm0_b_scale: torch.Tensor, - mm0_bias: Optional[torch.Tensor], - mm_dtype: torch.dtype, - ): - return - - def target_scaled_mm_prologue_pattern( - mm0_a: torch.Tensor, - mm0_b: torch.Tensor, - mm0_a_scale: torch.Tensor, - mm0_b_scale: torch.Tensor, - mm0_bias: Optional[torch.Tensor], - mm_dtype: torch.dtype, - ): - scaled_mm_output = torch.ops.trtllm.cublas_scaled_mm( - mm0_a, mm0_b, mm0_a_scale, mm0_b_scale, mm0_bias, mm_dtype, - True) - return scaled_mm_output - - # No extra check needed as the output dtype of scaled_mm has been verified when - # ub_copy is inserted. - register_replacement( - empty_scaled_mm_prologue_pattern, - target_scaled_mm_prologue_pattern, - [], - fwd_only, - custom_pass, - search_fn_pattern=scaled_mm_prologue_pattern, - ) - - def register_nvfp4_prologue(custom_pass: PatternMatcherPass): - trtllm_nvfp4_gemm_default = CallFunction( - torch.ops.trtllm.nvfp4_gemm.default, KeywordArg('act_fp4'), - KeywordArg('weight'), KeywordArg('act_sf'), - KeywordArg('weight_scale'), KeywordArg('alpha'), - KeywordArg('output_dtype')) - ub_copy = CallFunction(torch.ops.trtllm.copy_to_userbuffers, - trtllm_nvfp4_gemm_default) - nvfp4_gemm_prologue_pattern = MultiOutputPattern([ub_copy]) - - def empty_nvfp4_gemm_prologue_pattern( - act_fp4: torch.Tensor, - weight: torch.Tensor, - act_sf: torch.Tensor, - weight_scale: torch.Tensor, - alpha: torch.Tensor, - output_dtype: torch.dtype, - ): - return - - def target_nvfp4_gemm_prologue_pattern( - act_fp4: torch.Tensor, - weight: torch.Tensor, - act_sf: torch.Tensor, - weight_scale: torch.Tensor, - alpha: torch.Tensor, - output_dtype: torch.dtype, - ): - nvfp4_gemm_output = torch.ops.trtllm.nvfp4_gemm( - act_fp4, weight, act_sf, weight_scale, alpha, output_dtype, - True) - return nvfp4_gemm_output - - # No extra check needed as the output dtype of nvfp4_gemm has been verified when - # ub_copy is inserted. - register_replacement( - empty_nvfp4_gemm_prologue_pattern, - target_nvfp4_gemm_prologue_pattern, - [], - fwd_only, - custom_pass, - search_fn_pattern=nvfp4_gemm_prologue_pattern, - ) - - def register_mm_prologue(custom_pass: PatternMatcherPass): - aten_mm_default = CallFunction(torch.ops.aten.mm.default, - KeywordArg('mm0_a'), - KeywordArg('mm0_b')) - ub_copy = CallFunction(torch.ops.trtllm.copy_to_userbuffers, - aten_mm_default) - mm_prologue_pattern = MultiOutputPattern([ub_copy]) - - def empty_mm_prologue_pattern( - mm0_a: torch.Tensor, - mm0_b: torch.Tensor, - ): - return - - def target_mm_prologue_pattern( - mm0_a: torch.Tensor, - mm0_b: torch.Tensor, - ): - mm_output = torch.ops.trtllm.matmul_to_ub(mm0_a, mm0_b) - return mm_output - - # No extra check needed as the output dtype of mm has been verified when - # ub_copy is inserted. - register_replacement( - empty_mm_prologue_pattern, - target_mm_prologue_pattern, - [], - fwd_only, - custom_pass, - search_fn_pattern=mm_prologue_pattern, - ) - - def register_add_prologue(custom_pass: PatternMatcherPass): - aten_add_default = CallFunction(torch.ops.aten.add.Tensor, - KeywordArg('add_a'), - KeywordArg('add_b')) - ub_copy = CallFunction(torch.ops.trtllm.copy_to_userbuffers, - aten_add_default) - add_prologue_pattern = MultiOutputPattern([ub_copy]) - - def empty_add_prologue_pattern( - add_a: torch.Tensor, - add_b: torch.Tensor, - ): - return - - def target_add_prologue_pattern( - add_a: torch.Tensor, - add_b: torch.Tensor, - ): - add_output = torch.ops.trtllm.add_to_ub(add_a, add_b) - return add_output - - # No extra check needed as the output dtype of add has been verified when - # ub_copy is inserted. - register_replacement( - empty_add_prologue_pattern, - target_add_prologue_pattern, - [], - fwd_only, - custom_pass, - search_fn_pattern=add_prologue_pattern, - ) - - register_scaled_mm_prologue(custom_pass) - register_nvfp4_prologue(custom_pass) - register_mm_prologue(custom_pass) - register_add_prologue(custom_pass) - - def register_ub_finalize_patterns(custom_pass: PatternMatcherPass): - # TODO: Unify the finalize patterns once the allreduce interface does not contain - # dynamic number of tensors. - def allreduce_quant_finalize_pattern(custom_pass: PatternMatcherPass): - trtllm_userbuffers_allreduce_finalize_default = CallFunction( - torch.ops.trtllm.userbuffers_allreduce_finalize.default, - KeywordArg("sharded_residual"), False) - trtllm_allreduce_default = CallFunction( - torch.ops.trtllm.allreduce.default, KeywordArg("input"), - trtllm_userbuffers_allreduce_finalize_default, - KeywordArg("gamma"), KeywordArg("scale"), Ignored(), Ignored(), - mapping.tp_group, int(AllReduceStrategy.UB), - KeywordArg("fusion_op"), KeywordArg("eps"), Ignored()) - ub_ar_finalize_pattern = MultiOutputPattern( - [trtllm_allreduce_default]) - - def empty_quant_finalize_pattern( - input: torch.Tensor, - sharded_residual: torch.Tensor, - gamma: torch.Tensor, - scale: torch.Tensor, - fusion_op: int, - eps: float, - ): - return - - def target_quant_finalize_pattern( - input: torch.Tensor, - sharded_residual: torch.Tensor, - gamma: torch.Tensor, - scale: torch.Tensor, - fusion_op: int, - eps: float, - ): - all_reduce_output = torch.ops.trtllm.allreduce( - input, sharded_residual, gamma, - scale, None, None, mapping.tp_group, - int(AllReduceStrategy.UB), fusion_op, eps, True) - return all_reduce_output - - register_replacement( - empty_quant_finalize_pattern, - target_quant_finalize_pattern, - [], - fwd_only, - custom_pass, - search_fn_pattern=ub_ar_finalize_pattern, - ) - - def allreduce_half_finalize_pattern(custom_pass: PatternMatcherPass): - trtllm_userbuffers_allreduce_finalize_default = CallFunction( - torch.ops.trtllm.userbuffers_allreduce_finalize.default, - KeywordArg("sharded_residual"), False) - trtllm_allreduce_default = CallFunction( - torch.ops.trtllm.allreduce.default, KeywordArg("input"), - trtllm_userbuffers_allreduce_finalize_default, - KeywordArg("gamma"), Ignored(), Ignored(), Ignored(), - mapping.tp_group, int(AllReduceStrategy.UB), - int(AllReduceFusionOp.RESIDUAL_RMS_NORM), KeywordArg("eps"), - Ignored()) - ub_ar_finalize_pattern = MultiOutputPattern( - [trtllm_allreduce_default]) - - def empty_half_finalize_pattern( - input: torch.Tensor, - sharded_residual: torch.Tensor, - gamma: torch.Tensor, - eps: float, - ): - return - - def target_half_finalize_pattern( - input: torch.Tensor, - sharded_residual: torch.Tensor, - gamma: torch.Tensor, - eps: float, - ): - all_reduce_output = torch.ops.trtllm.allreduce( - input, sharded_residual, gamma, None, None, None, - mapping.tp_group, int(AllReduceStrategy.UB), - int(AllReduceFusionOp.RESIDUAL_RMS_NORM), eps, True) - return all_reduce_output - - register_replacement( - empty_half_finalize_pattern, - target_half_finalize_pattern, - [], - fwd_only, - custom_pass, - search_fn_pattern=ub_ar_finalize_pattern, - ) - - allreduce_quant_finalize_pattern(custom_pass) - allreduce_half_finalize_pattern(custom_pass) - - custom_passes.append(PatternMatcherPass()) - register_ub_allreduce_quantize_fusion(custom_passes[-1]) - - custom_passes.append(PatternMatcherPass()) - register_convert_supported_ar_to_ub(custom_passes[-1]) - - custom_passes.append(PatternMatcherPass()) - register_ub_prologue_patterns(custom_passes[-1]) - - custom_passes.append(PatternMatcherPass()) - register_ub_finalize_patterns(custom_passes[-1]) diff --git a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py index 75a9aeff8e5c..8e60b6bd36b5 100644 --- a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py +++ b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py @@ -12,7 +12,9 @@ from tensorrt_llm.llmapi.utils import enable_llm_debug from tensorrt_llm.logger import logger -from ..utils import get_piecewise_cuda_graph_flag, make_weak_ref +from ..utils import (get_model_extra_attrs, get_piecewise_cuda_graph_flag, + make_weak_ref) +from .multi_stream.auto_multi_stream import multi_stream_schedule from .utils import (get_enable_piecewise_cuda_graph_capture_flag, is_call_function) @@ -29,6 +31,7 @@ def __init__( graph_pool_handle: tuple[int, int], garbage_collect_values: bool = True, graph=None, + max_num_streams: int = 1, ): super().__init__(module, garbage_collect_values, graph) @@ -39,6 +42,8 @@ def __init__( self.exclude_modules = [f"submod_{i}" for i in exclude_modules_id] self.graph_pool_handle = graph_pool_handle self.enable_inductor = enable_inductor + self.num_events = 0 + self.max_num_streams = max_num_streams def run(self, *args): fake_args = [ @@ -72,6 +77,11 @@ def call_module(self, target, args, kwargs): found_dynamic_shape = True break + if self.max_num_streams > 1 and not self.enable_inductor: + num_events = multi_stream_schedule(submod, self.max_num_streams) + self.num_events = max(self.num_events, num_events) + submod.recompile() + self.module.__dict__[target] = PiecewiseRunner( submod, target, @@ -179,8 +189,12 @@ def __call__(self, *args): with patch("gc.collect", lambda: None): # TODO: consider to use `make_graphed_callables()` when # it's ready rather than capture it ourselves + # Graph Capture would override the stream. We need to setup the stream correctly. + extra_attrs = get_model_extra_attrs() with torch.cuda.graph(graph, pool=self.graph_pool_handle): + extra_attrs["global_stream"] = torch.cuda.current_stream() output = entry.callable(*args) + extra_attrs["global_stream"] = torch.cuda.current_stream() entry.cuda_graph = graph # Mark weak ref here. The intermediate activation tensor should be freed properly. @@ -218,7 +232,8 @@ def piecewise_optimizer( input_num_tokens: Union[int | torch.SymInt], cuda_graph_batch_sizes: Sequence[int], graph_pool_handle: tuple[int, int], -) -> GraphModule: + max_num_streams: int = 1, +) -> tuple[GraphModule, int]: graph_pool_handle = torch.cuda.graph_pool_handle() graph = gm.graph @@ -253,13 +268,16 @@ def piecewise_optimizer( lambda node: node_to_graph_id[node], keep_original_order=True) - PiecewiseInterpreter( + interpreter = PiecewiseInterpreter( gm, enable_inductor, input_num_tokens, cuda_graph_batch_sizes, exclude_modules_id, graph_pool_handle, - ).run(*example_inputs) + max_num_streams=max_num_streams, + ) + + interpreter.run(*example_inputs) - return gm + return gm, interpreter.num_events diff --git a/tensorrt_llm/_torch/compilation/remove_copy_pass.py b/tensorrt_llm/_torch/compilation/remove_copy_pass.py index fe968f020be0..8e5eb7a81148 100644 --- a/tensorrt_llm/_torch/compilation/remove_copy_pass.py +++ b/tensorrt_llm/_torch/compilation/remove_copy_pass.py @@ -5,7 +5,7 @@ auto_functionalized_v2) from torch.fx import Graph, Node -from .utils import is_call_function +from .utils import inplace_info, is_call_function aten = torch.ops.aten @@ -46,19 +46,12 @@ def remove_functionalize_inner(node: Node, mutates_args: dict, is_v2=False): inplace_func = node.args[0] - if inplace_func == torch.ops.trtllm.flashinfer_fused_add_rmsnorm.default: - remove_functionalize_inner( - node, - { - 1: "input", - 2: "residual" - }, - is_v2=node.target == auto_functionalized_v2, - ) - if inplace_func == torch.ops.trtllm.attention_inplace.default: - remove_functionalize_inner(node, {1: "output", 2: "output_sf"}) - if inplace_func == torch.ops.trtllm.mla_custom_op_inplace.default: - remove_functionalize_inner(node, {1: "output"}) + inplace_map = inplace_info() + if inplace_func not in inplace_map: + # We do not know the inplace op + continue + + remove_functionalize_inner(node, inplace_map[inplace_func]) for node in nodes_to_remove: graph.erase_node(node) diff --git a/tensorrt_llm/_torch/compilation/utils.py b/tensorrt_llm/_torch/compilation/utils.py index 6e900b9e3fd4..f00d689458af 100644 --- a/tensorrt_llm/_torch/compilation/utils.py +++ b/tensorrt_llm/_torch/compilation/utils.py @@ -41,3 +41,20 @@ def set_enable_piecewise_cuda_graph_capture_flag(enable: bool): def get_enable_piecewise_cuda_graph_capture_flag() -> bool: global _enable_piecewise_cuda_graph_capture return _enable_piecewise_cuda_graph_capture + + +def inplace_info(): + inplace_map = { + torch.ops.trtllm.flashinfer_fused_add_rmsnorm.default: { + 1: "input", + 2: "residual" + }, + torch.ops.trtllm.attention_inplace.default: { + 1: "output", + 2: "output_sf" + }, + torch.ops.trtllm.mla_custom_op_inplace.default: { + 1: "output" + } + } + return inplace_map diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 35eb19acf5f5..31fa33d3084d 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -501,51 +501,6 @@ def _(input, sizes, group): shape[0] = sizes[local_rank] return input.new_empty(shape) - @torch.library.register_fake("trtllm::fp4_block_scale_moe_runner") - def _( - routing_logits, - routing_bias, - hidden_states, - hidden_states_scale, - gemm1_weights, - gemm1_weights_scale, - gemm2_weights, - gemm2_weights_scale, - output1_scale_scalar, - output1_scale_gate_scalar, - output2_scale_scalar, - num_experts, - top_k, - n_group, - topk_group, - intermediate_size, - local_expert_offset, - local_num_experts, - routed_scaling_factor, - tile_tokens_dim, - routing_method_type, - do_finalize, - ) -> List[torch.Tensor]: - num_tokens = hidden_states.shape[0] - hidden_size = hidden_states.shape[1] * 2 - if do_finalize: - return [ - hidden_states.new_empty((num_tokens, hidden_size), - dtype=torch.bfloat16) - ] - - expanded_row_count = num_tokens * top_k - max_padding_required = (tile_tokens_dim - 1) * num_experts - max_num_padded_tokens = fp4_utils.pad_up( - expanded_row_count + max_padding_required, tile_tokens_dim) - wt_dtype = routing_bias.dtype if routing_bias is not None else torch.bfloat16 - return [ - hidden_states.new_empty((max_num_padded_tokens, hidden_size), - dtype=torch.bfloat16), - hidden_states.new_empty((num_tokens, top_k), dtype=wt_dtype), - hidden_states.new_empty((num_tokens, top_k), dtype=torch.int32) - ] - @torch.library.register_fake("trtllm::nvfp4_block_scale_interleave") def _(sf: torch.Tensor): rows = sf.shape[-2] @@ -559,3 +514,12 @@ def _(sf: torch.Tensor): @torch.library.register_fake("trtllm::nvfp4_block_scale_interleave_reverse") def _(sf: torch.Tensor): return torch.empty_like(sf, dtype=torch.uint8) + + @torch.library.register_fake("trtllm::moe_finalize_allreduce") + def _(input, residual, norm_weight, expanded_idx_to_permuted_idx, + shared_expert_output, expert_scale_factor, workspace, rank, nranks, + eps) -> List[torch.Tensor]: + return [ + torch.empty_like(residual), + torch.empty_like(residual), + ] diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index ffeb90c2fd3e..60ef215fe386 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -675,24 +675,114 @@ def _( dtype=output_dtype) -class W4A16GemmRunner(TunableRunner): +class WeightOnlyQuantGemmRunner(TunableRunner): + runner_dict = dict() + tuning_config = TuningConfig(dynamic_tensor_specs=( + DynamicTensorSpec(0, 0, get_last_power_of_2_num_tokens_buckets, + last_positive_power_of_2), )) + + def __init__( + self, + activation_dtype: torch.dtype, + weight_dtype: torch.dtype, + output_dtype: torch.dtype, + to_userbuffers: bool, + ): + self.output_dtype = output_dtype + self.to_userbuffers = to_userbuffers + instance_key = (activation_dtype, weight_dtype) + if instance_key not in WeightOnlyQuantGemmRunner.runner_dict: + WeightOnlyQuantGemmRunner.runner_dict[ + instance_key] = torch.classes.trtllm.WeightOnlyQuantGemmRunner( + activation_dtype, weight_dtype) + self.weight_only_quant_gemm_runner = WeightOnlyQuantGemmRunner.runner_dict[ + instance_key] + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + ) -> List[int]: + return list(range(self.weight_only_quant_gemm_runner.get_num_configs())) + + def forward( + self, + inputs: List[torch.Tensor], + tactic: int = -1, + ) -> torch.Tensor: + activation, weight, weight_scale = inputs + return self.weight_only_quant_gemm_runner.run_gemm( + activation, + weight, + weight_scale, + tactic, + self.to_userbuffers, + self.output_dtype, + ) + + +@torch.library.custom_op("trtllm::weight_only_quant_gemm", mutates_args=()) +def weight_only_quant_gemm( + activation: torch.Tensor, + weight: torch.Tensor, + weight_dtype: torch.dtype, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + to_userbuffers: bool = False, +) -> torch.Tensor: + + tuner = AutoTuner.get() + + # allocate workspace for profiling + weight_only_quant_gemm_runner = WeightOnlyQuantGemmRunner( + activation.dtype, weight_dtype, output_dtype, to_userbuffers) + + _, best_tactic = tuner.choose_one( + "trtllm::weight_only_quant_gemm::gemm", + [weight_only_quant_gemm_runner], + WeightOnlyQuantGemmRunner.tuning_config, + [activation, weight, weight_scale], + ) + + return weight_only_quant_gemm_runner( + inputs=[activation, weight, weight_scale], tactic=best_tactic) + + +@weight_only_quant_gemm.register_fake +def _( + activation: torch.Tensor, + weight: torch.Tensor, + weight_type: torch.dtype, + weight_scale: torch.Tensor, + output_dtype: torch.dtype = None, + to_userbuffers: bool = False, +) -> torch.Tensor: + dtype = output_dtype if output_dtype is not None else activation.dtype + return activation.new_empty((activation.size(0), weight.size(1)), + dtype=dtype) + + +class FinegrainedMixedDtypeGemm(TunableRunner): _runner_dict = dict() MAX_SUPPORTED_SM_VERSION = 90 - def __init__(self, activation_dtype: torch.dtype, quant_mode: int): - instance_key = (activation_dtype, quant_mode) - if instance_key not in W4A16GemmRunner._runner_dict: - W4A16GemmRunner._runner_dict[ - instance_key] = torch.classes.trtllm.W4A16GemmRunner( - activation_dtype, quant_mode) - self._w4a16_gemm_runner = W4A16GemmRunner._runner_dict[instance_key] + def __init__(self, activation_dtype: torch.dtype, output_dtype: torch.dtype, + quant_mode: int): + instance_key = (activation_dtype, output_dtype, quant_mode) + if instance_key not in FinegrainedMixedDtypeGemm._runner_dict: + FinegrainedMixedDtypeGemm._runner_dict[ + instance_key] = torch.classes.trtllm.finegrainedMixedDtypeGemmRunner( + activation_dtype, output_dtype, quant_mode) + self._finegrained_mixed_dtype_gemm_runner = FinegrainedMixedDtypeGemm._runner_dict[ + instance_key] def get_valid_tactics( self, inputs: List[torch.Tensor], profile: OptimizationProfile, ) -> List[int]: - return list(range(self._w4a16_gemm_runner.get_num_configs())) + return list( + range(self._finegrained_mixed_dtype_gemm_runner.get_num_configs())) def forward(self, inputs: List[torch.Tensor], @@ -707,25 +797,25 @@ def forward(self, activation, weights_packed, scales = inputs - return self._w4a16_gemm_runner.run_gemm( - activation, - weights_packed, - scales, - kwargs["group_size"], - tactic, - kwargs["bias"], - kwargs["zeros"], - ) + alpha = 1.0 if kwargs.get("alpha") is None else kwargs["alpha"] + + return self._finegrained_mixed_dtype_gemm_runner.run_gemm( + activation, weights_packed, scales, kwargs["group_size"], tactic, + kwargs["bias"], kwargs["zeros"], alpha) -@torch.library.custom_op("trtllm::w4a16_gemm", mutates_args=()) -def w4a16_gemm(input: torch.Tensor, - weight: torch.Tensor, - scales: torch.Tensor, - group_size: int, - has_zero_point: bool, - bias: Optional[torch.Tensor] = None, - zeros: Optional[torch.Tensor] = None) -> torch.Tensor: +@torch.library.custom_op("trtllm::finegrained_mixed_dtype_gemm", + mutates_args=()) +def finegrained_mixed_dtype_gemm( + input: torch.Tensor, + weight: torch.Tensor, + scales: torch.Tensor, + group_size: int, + has_zero_point: bool, + output_dtype: torch.dtype, + alpha: Optional[float] = None, + bias: Optional[torch.Tensor] = None, + zeros: Optional[torch.Tensor] = None) -> torch.Tensor: assert not has_zero_point or zeros is not None, "Expected 'zeros' tensor when has_zero_point is True" @@ -741,16 +831,44 @@ def w4a16_gemm(input: torch.Tensor, if quant_mode == 0: assert zeros is None, "When quant_mode is 0 (FINEGRAINED_SCALE_ONLY), zeros must be None" - w4a16_gemm_runner = W4A16GemmRunner(input.dtype, quant_mode) + finegrained_mixed_dtype_gemm_runner = FinegrainedMixedDtypeGemm( + input.dtype, output_dtype, quant_mode) - kwargs = {"group_size": group_size, "zeros": zeros, "bias": bias} - _, best_tactic = tuner.choose_one("trtllm::w4a16_gemm::gemm", - [w4a16_gemm_runner], tuning_config, - [input, weight, scales], **kwargs) + kwargs = { + "group_size": group_size, + "zeros": zeros, + "bias": bias, + "alpha": alpha + } + + _, best_tactic = tuner.choose_one( + "trtllm::finegrained_mixed_dtype_gemm::gemm", + [finegrained_mixed_dtype_gemm_runner], tuning_config, + [input, weight, scales], **kwargs) - return w4a16_gemm_runner(inputs=[input, weight, scales], - tactic=best_tactic, - **kwargs) + return finegrained_mixed_dtype_gemm_runner(inputs=[input, weight, scales], + tactic=best_tactic, + **kwargs) + + +@finegrained_mixed_dtype_gemm.register_fake +def _( + input: torch.Tensor, + weight: torch.Tensor, + scales: torch.Tensor, + group_size: int, + has_zero_point: bool, + output_dtype: torch.dtype, + alpha: Optional[float] = None, + bias: Optional[torch.Tensor] = None, + zeros: Optional[torch.Tensor] = None, +) -> torch.Tensor: + # For a typical GEMM: input [M, K] @ weight [K, N] -> output [M, N] + # Weight is typically packed, so we need to infer the output dimension + M = input.size(0) + # Assuming weight is packed and the output dimension can be inferred from weight.size(1) + N = weight.size(1) if weight.dim() > 1 else weight.size(0) + return input.new_empty((M, N), dtype=output_dtype) @torch.library.custom_op("trtllm::attention", mutates_args=()) @@ -958,3 +1076,45 @@ def _( output_sf = torch.empty(()) # Create a placeholder, which is not used. return output_act, output_sf + + +def get_event(event_idx: int): + from ..utils import get_model_extra_attrs + extra_attrs = get_model_extra_attrs() + assert "events" in extra_attrs, "Missing Event Book" + return extra_attrs["events"]()[event_idx] + + +def get_stream(stream_id: int): + from ..utils import get_model_extra_attrs + extra_attrs = get_model_extra_attrs() + if stream_id == 0: + return extra_attrs["global_stream"] + assert "aux_streams" in extra_attrs, "Missing Aux Streams" + return extra_attrs["aux_streams"]()[stream_id - 1] + + +@torch.library.custom_op("trtllm::set_stream", mutates_args=()) +def set_stream(stream_id: int) -> None: + stream = get_stream(stream_id) + assert stream is not None + torch.cuda.set_stream(stream) + + +@torch.library.custom_op("trtllm::record_event", mutates_args=()) +def record_event(event_idx: int) -> None: + event = get_event(event_idx) + event.record() + + +@torch.library.custom_op("trtllm::wait_event", mutates_args=()) +def wait_event(event_idx: int) -> None: + event = get_event(event_idx) + event.wait() + + +@torch.library.custom_op("trtllm::record_stream", mutates_args=()) +def record_stream(tensor: torch.Tensor, stream_id: int) -> None: + stream = get_stream(stream_id) + assert stream is not None + tensor.record_stream(stream) diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index a8d3b7e7ce0f..622fa12c5150 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -4,13 +4,28 @@ import torch -from tensorrt_llm._torch.utils import (get_last_power_of_2_num_tokens_buckets, - last_positive_power_of_2) +from tensorrt_llm._torch.utils import (fp4_utils, + get_last_power_of_2_num_tokens_buckets, + last_positive_power_of_2, + next_positive_power_of_2) from ..autotuner import (AutoTuner, ConstraintSpec, DynamicTensorSpec, OptimizationProfile, TunableRunner, TuningConfig) +def calculate_tile_tokens_dim(num_tokens: int, num_experts: int, + top_k: int) -> int: + # Guess tokens per expert assuming perfect expert distribution first. + num_tokens_per_expert = num_tokens * top_k // num_experts + + # And pad the number to the next power of 2. + tile_tokens_dim = next_positive_power_of_2(num_tokens_per_expert) + # Cap to 8-64 tokens per CTA tile as it's the range supported by the kernel. + tile_tokens_dim = min(max(tile_tokens_dim, 8), 64) + + return tile_tokens_dim + + @dataclass(frozen=True) class FP4BlockScaleMoEInputs: @@ -220,11 +235,14 @@ def fp4_block_scale_moe_runner(routing_logits: torch.Tensor, intermediate_size: int, local_expert_offset: int, local_num_experts: int, routed_scaling_factor: Optional[float], - tile_tokens_dim: int, routing_method_type: int, + routing_method_type: int, do_finalize: bool) -> List[torch.Tensor]: tuner = AutoTuner.get() + num_tokens = hidden_states.shape[0] + tile_tokens_dim = calculate_tile_tokens_dim(num_tokens, num_experts, top_k) + kernel_runner = FP4BlockScaleMoERunner( num_experts, top_k, n_group, topk_group, intermediate_size, local_expert_offset, local_num_experts, routed_scaling_factor, @@ -254,6 +272,53 @@ def fp4_block_scale_moe_runner(routing_logits: torch.Tensor, return kernel_runner(inputs, tactic=best_tactic) +@fp4_block_scale_moe_runner.register_fake +def _( + routing_logits, + routing_bias, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + gemm2_weights, + gemm2_weights_scale, + output1_scale_scalar, + output1_scale_gate_scalar, + output2_scale_scalar, + num_experts, + top_k, + n_group, + topk_group, + intermediate_size, + local_expert_offset, + local_num_experts, + routed_scaling_factor, + routing_method_type, + do_finalize, +) -> List[torch.Tensor]: + num_tokens = hidden_states.shape[0] + hidden_size = hidden_states.shape[1] * 2 + if do_finalize: + return [ + hidden_states.new_empty((num_tokens, hidden_size), + dtype=torch.bfloat16) + ] + + tile_tokens_dim = calculate_tile_tokens_dim(num_tokens, num_experts, top_k) + + expanded_row_count = num_tokens * top_k + max_padding_required = (tile_tokens_dim - 1) * num_experts + max_num_padded_tokens = fp4_utils.pad_up( + expanded_row_count + max_padding_required, tile_tokens_dim) + wt_dtype = routing_bias.dtype if routing_bias is not None else torch.bfloat16 + return [ + hidden_states.new_empty((max_num_padded_tokens, hidden_size), + dtype=torch.bfloat16), + hidden_states.new_empty((num_tokens, top_k), dtype=wt_dtype), + hidden_states.new_empty((num_tokens, top_k), dtype=torch.int32) + ] + + @dataclass(frozen=True) class FP8BlockScaleMoEInputs: @@ -420,23 +485,31 @@ def get_tuning_config(cls) -> TuningConfig: @torch.library.custom_op("trtllm::fp8_block_scale_moe_runner", mutates_args=()) -def fp8_block_scale_moe_runner(routing_logits: torch.Tensor, - routing_bias: torch.Tensor, - hidden_states: torch.Tensor, - hidden_states_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm1_weights_scale: torch.Tensor, - gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, - num_experts: int, top_k: int, n_group: int, - topk_group: int, intermediate_size: int, - local_expert_offset: int, local_num_experts: int, - routed_scaling_factor: float, - tile_tokens_dim: int, - routing_method_type: int) -> torch.Tensor: +def fp8_block_scale_moe_runner( + routing_logits: torch.Tensor, + routing_bias: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + num_experts: int, + top_k: int, + n_group: int, + topk_group: int, + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: float, + routing_method_type: int, +) -> torch.Tensor: tuner = AutoTuner.get() + num_tokens = hidden_states.shape[0] + tile_tokens_dim = calculate_tile_tokens_dim(num_tokens, num_experts, top_k) + kernel_runner = FP8BlockScaleMoERunner(num_experts, top_k, n_group, topk_group, intermediate_size, local_expert_offset, @@ -463,3 +536,30 @@ def fp8_block_scale_moe_runner(routing_logits: torch.Tensor, ) return kernel_runner(inputs, tactic=best_tactic) + + +@fp8_block_scale_moe_runner.register_fake +def _( + routing_logits: torch.Tensor, + routing_bias: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + num_experts: int, + top_k: int, + n_group: int, + topk_group: int, + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: float, + routing_method_type: int, +) -> torch.Tensor: + num_tokens = hidden_states.shape[0] + hidden_size = hidden_states.shape[1] * 2 + + return hidden_states.new_empty((num_tokens, hidden_size), + dtype=torch.bfloat16) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 83fbf5f91efe..ba713a7d566b 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -88,8 +88,8 @@ def get_allreduce_mnnvl_workspace( # This is a buffer to maintain the state of this allreduce Op # Should have the same lifetime with self._buffer - # [Buffer_ptr, Clear_ptr, Buffer_size, atomic access counter] - buffer_flags = torch.tensor([0, 2, max_num_elements, 0], + # [Buffer_ptr, Clear_ptr, Buffer_size, num_tokens_to_clear,atomic access counter] + buffer_flags = torch.tensor([0, 2, max_num_elements, 0, 0], dtype=torch.uint32, device=torch.device("cuda", mapping.local_rank)) @@ -305,7 +305,7 @@ def __init__(self, mapping: Mapping, dtype: torch.dtype): @staticmethod def get_supported_dtypes(): - return (torch.bfloat16, torch.float32) + return (torch.float16, torch.bfloat16, torch.float32) def forward( self, @@ -458,6 +458,7 @@ def forward( == False): return input + allreduce_strategy = self.strategy if all_reduce_params is None: all_reduce_params = AllReduceParams() @@ -469,6 +470,9 @@ def forward( return mnnvl_output # Fall back to regular AllReduce if MNNVL is not available or not applicable + # Make sure the strategy is AUTO since allreduceOp does not have the branch for MNNVL + if allreduce_strategy == AllReduceStrategy.MNNVL: + allreduce_strategy = AllReduceStrategy.AUTO output = torch.ops.trtllm.allreduce( input=input, residual=all_reduce_params.residual, @@ -477,7 +481,7 @@ def forward( bias=all_reduce_params.bias, workspace=self.workspace, group=self.mapping.tp_group, - strategy=self.strategy, + strategy=allreduce_strategy, op=all_reduce_params.fusion_op, eps=all_reduce_params.eps, trigger_completion_at_end=all_reduce_params. diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 830cd5bda6a7..3d0175a3c234 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -202,6 +202,9 @@ def from_pretrained(cls, json_quant_configs = quant_config_dict['quantization'] quant_config.quant_algo = json_quant_configs.get('quant_algo', None) + # fp8_pb_wo from modelopt is the same as FP8_BLOCK_SCALES + if quant_config.quant_algo == "fp8_pb_wo": + quant_config.quant_algo = 'FP8_BLOCK_SCALES' quant_config.kv_cache_quant_algo = json_quant_configs.get( 'kv_cache_quant_algo', None) quant_config.group_size = json_quant_configs.get('group_size', None) @@ -294,6 +297,49 @@ def get_bindings_model_config(self, num_heads = self.pretrained_config.num_attention_heads // ( self.mapping.tp_size * self.mapping.cp_size) + + # Handle both uniform and per-layer KV heads + num_kv_heads_per_layer = getattr(self.pretrained_config, + 'num_kv_heads_per_layer', None) + if num_kv_heads_per_layer is not None: + # For models with per-layer KV heads, like nemotron-nas + kv_heads_per_layer_raw = num_kv_heads_per_layer + use_per_layer_kv_heads = True + else: + # Check if num_key_value_heads is a list (per-layer) or scalar (uniform) + num_kv_heads_raw = getattr(self.pretrained_config, + 'num_key_value_heads', None) + + if num_kv_heads_raw is not None and isinstance( + num_kv_heads_raw, list): + # num_key_value_heads is a list - treat as per-layer KV heads + kv_heads_per_layer_raw = num_kv_heads_raw + use_per_layer_kv_heads = True + else: + # num_key_value_heads is scalar or None - treat as uniform KV heads + if num_kv_heads_raw is None: + # For uniform models, check: num_key_value_heads (standard) -> num_query_groups (NeMo) -> num_attention_heads + num_kv_heads_raw = getattr( + self.pretrained_config, 'num_query_groups', + self.pretrained_config.num_attention_heads) + + num_kv_heads = num_kv_heads_raw // (self.mapping.tp_size * + self.mapping.cp_size) + use_per_layer_kv_heads = False + + if use_per_layer_kv_heads: + # TRT-LLM LoRA requires uniform KV heads across layers + if self.lora_config is not None and len( + set(kv_heads_per_layer_raw)) > 1: + raise ValueError( + f"TRT-LLM LoRA requires uniform KV heads across layers, " + f"got: {kv_heads_per_layer_raw}") + # Apply TP/CP scaling to each layer + num_kv_heads_per_layer = [ + kv_heads // (self.mapping.tp_size * self.mapping.cp_size) + for kv_heads in kv_heads_per_layer_raw + ] + hidden_size = self.pretrained_config.hidden_size // self.mapping.tp_size model_config_cpp = ModelConfigCpp( @@ -305,6 +351,8 @@ def get_bindings_model_config(self, hidden_size=hidden_size, data_type=torch_dtype_to_binding( self.pretrained_config.torch_dtype)) + + # For kv cache size calculation: set tokens_per_block if tokens_per_block is None: logger.warning( f"tokens_per_block is not set, using default value {model_config_cpp.tokens_per_block}" @@ -312,6 +360,11 @@ def get_bindings_model_config(self, else: model_config_cpp.tokens_per_block = tokens_per_block + if use_per_layer_kv_heads: + model_config_cpp.num_kv_heads_per_layer = num_kv_heads_per_layer + else: + model_config_cpp.set_num_kv_heads(num_kv_heads) + mlp_hidden_size = None if self.pretrained_config.intermediate_size is not None: mlp_hidden_size = self.pretrained_config.intermediate_size // self.mapping.tp_size @@ -333,9 +386,16 @@ def get_bindings_model_config(self, f"Failed to infer mlp hidden size for model: {self.pretrained_config.model_type}" ) - if "head_size" in self.pretrained_config: - head_size = self.pretrained_config.head_size + # For kv cache size calculation: set size_per_head + head_dim_names = ["head_size", "head_dim"] + for head_dim_name in head_dim_names: + if head_dim_name in self.pretrained_config: + head_size = getattr(self.pretrained_config, head_dim_name) + break else: + logger.warning( + f"head_size/head_dim is not set, using default value {hidden_size // num_heads}" + ) head_size = hidden_size // num_heads model_config_cpp.mlp_hidden_size = mlp_hidden_size @@ -353,8 +413,10 @@ def _infer_nemotron_ffn_mult(self): # Nemotron-NAS has variable ffn_mult for each layer, we need to find the maximum # so that we don't set a too small mlp_hidden_size. This solution leads to a memory # consumption that is higher than required. - biggest_ffn_mult = max( - [x.ffn.ffn_mult for x in self.pretrained_config.block_configs]) + biggest_ffn_mult = max([ + (x.ffn.ffn_mult if x.ffn.ffn_mult is not None else 0) + for x in self.pretrained_config.block_configs + ]) from tensorrt_llm._torch.models.modeling_nemotron_nas import \ _ffn_mult_to_intermediate_size diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index cc75dffb5c0c..e4da7aff5a9a 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -10,11 +10,13 @@ from .modeling_hyperclovax import HCXVisionForCausalLM from .modeling_llama import LlamaForCausalLM from .modeling_llava_next import LlavaNextModel -from .modeling_mistral import MistralForCausalLM +from .modeling_mistral import Mistral3VLM, MistralForCausalLM from .modeling_mixtral import MixtralForCausalLM from .modeling_nemotron import NemotronForCausalLM from .modeling_nemotron_h import NemotronHForCausalLM from .modeling_nemotron_nas import NemotronNASForCausalLM +from .modeling_phi3 import Phi3ForCausalLM +from .modeling_phi4mm import Phi4MMForCausalLM from .modeling_qwen import (Qwen2ForCausalLM, Qwen2ForProcessRewardModel, Qwen2ForRewardModel) from .modeling_qwen2vl import Qwen2_5_VLModel, Qwen2VLModel @@ -37,11 +39,14 @@ "HCXVisionForCausalLM", "LlamaForCausalLM", "LlavaNextModel", + "Mistral3VLM", "MistralForCausalLM", "MixtralForCausalLM", "NemotronForCausalLM", "NemotronHForCausalLM", "NemotronNASForCausalLM", + "Phi3ForCausalLM", + "Phi4MMForCausalLM", "Qwen2ForCausalLM", "Qwen2ForProcessRewardModel", "Qwen2ForRewardModel", diff --git a/tensorrt_llm/_torch/models/checkpoints/__init__.py b/tensorrt_llm/_torch/models/checkpoints/__init__.py new file mode 100644 index 000000000000..58789f36458b --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/__init__.py @@ -0,0 +1,18 @@ +from .base_checkpoint_loader import BaseCheckpointLoader +from .hf.checkpoint_loader import HfCheckpointLoader +from .hf.config_loader import HfConfigLoader +from .hf.gemma3_weight_mapper import Gemma3HfWeightMapper +from .hf.llama4_weight_mapper import Llama4HfWeightMapper +from .hf.mixtral_weight_mapper import MixtralHfWeightMapper +from .hf.nemotron_h_weight_mapper import NemotronHHfWeightMapper +from .hf.qwen2_moe_weight_mapper import Qwen2MoeHfWeightMapper +from .hf.qwen3_moe_weight_mapper import Qwen3MoeHfWeightMapper +from .hf.weight_loader import HfWeightLoader +from .hf.weight_mapper import HfWeightMapper + +__all__ = [ + "HfConfigLoader", "HfWeightLoader", "HfWeightMapper", + "BaseCheckpointLoader", "HfCheckpointLoader", "NemotronHHfWeightMapper", + "Gemma3HfWeightMapper", "MixtralHfWeightMapper", "Llama4HfWeightMapper", + "Qwen2MoeHfWeightMapper", "Qwen3MoeHfWeightMapper" +] diff --git a/tensorrt_llm/_torch/models/checkpoints/auto_mapper.py b/tensorrt_llm/_torch/models/checkpoints/auto_mapper.py new file mode 100644 index 000000000000..8bb5bb92a8af --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/auto_mapper.py @@ -0,0 +1,17 @@ +from typing import Optional + +from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPER_MAPPING + + +class AutoCheckpointMapper(): + + @staticmethod + def get(format: str, name: Optional[str] = None) -> "BaseWeightMapper": + if name is not None: + try: + return MODEL_CLASS_MAPPER_MAPPING[f'{name}_{format}']() + except KeyError: # no mapper for this model architecture, resort to default + # TODO smor- a potential bug here, if the class isn't added to __init__, it will return the default mapper + return MODEL_CLASS_MAPPER_MAPPING[format]() + else: + return MODEL_CLASS_MAPPER_MAPPING[format]() diff --git a/tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py new file mode 100644 index 000000000000..c1bfec0144a9 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py @@ -0,0 +1,87 @@ +from abc import ABC, abstractmethod +from typing import Any + +from torch import nn + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.checkpoints.auto_mapper import \ + AutoCheckpointMapper +from tensorrt_llm._torch.models.checkpoints.base_config_loader import \ + BaseConfigLoader +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import \ + BaseWeightLoader +from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ + BaseWeightMapper +from tensorrt_llm._torch.models.modeling_utils import \ + CHECKPOINT_LOADER_FORMAT_DEFAULT_MAPPING + + +class BaseCheckpointLoader(ABC): + + @abstractmethod + def get_default_weight_loader(self) -> BaseWeightLoader: + raise NotImplementedError + + @abstractmethod + def get_default_config_loader(self) -> BaseConfigLoader: + raise NotImplementedError + + @abstractmethod + def cleanup(self) -> None: + raise NotImplementedError + + @property + @abstractmethod + def weight_loader(self) -> BaseWeightLoader: + ... + + @property + @abstractmethod + def weight_mapper(self) -> BaseWeightMapper: + ... + + @property + @abstractmethod + def config_loader(self) -> BaseConfigLoader: + ... + + @property + @abstractmethod + def checkpoint_format(self) -> str: + ... + + def load_config(self, checkpoint_dir: str, **kwargs) -> ModelConfig: + return self.config_loader.load(checkpoint_dir, **kwargs) + + def load_weights(self, checkpoint_dir: str, **kwargs) -> dict[str, Any]: + return self.weight_loader.load_weights(checkpoint_dir, **kwargs) + + @classmethod + def get(cls, checkpoint_format: str, **kwargs) -> "BaseCheckpointLoader": + try: + return CHECKPOINT_LOADER_FORMAT_DEFAULT_MAPPING[checkpoint_format]( + **kwargs) + except KeyError: + raise ValueError( + f"Checkpoint loader for format {checkpoint_format} not found, " + f"available formats are: {CHECKPOINT_LOADER_FORMAT_DEFAULT_MAPPING.keys()}" + ) + + def get_initilized_weight_mapper(self, model: nn.Module, + config: ModelConfig) -> BaseWeightMapper: + weight_mapper = None + if self.weight_mapper is not None: + self.weight_mapper.init_model_and_config(model, config) + return self.weight_mapper + else: + # The name of the registered mapper should be the model architecture + if config.pretrained_config and config.pretrained_config.architectures: + model_arch = config.pretrained_config.architectures[0] + else: + raise ValueError( + "Cannot determine model architecture from config") + weight_mapper = AutoCheckpointMapper.get(self.checkpoint_format, + model_arch) + weight_mapper.init_model_and_config(model, config) + self.weight_mapper = weight_mapper + return weight_mapper diff --git a/tensorrt_llm/_torch/models/checkpoints/base_config_loader.py b/tensorrt_llm/_torch/models/checkpoints/base_config_loader.py new file mode 100644 index 000000000000..4af1dadaddc6 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/base_config_loader.py @@ -0,0 +1,13 @@ +from abc import ABC, abstractmethod + +from tensorrt_llm._torch.model_config import ModelConfig + + +class BaseConfigLoader(ABC): + + @abstractmethod + def load(self, checkpoint_dir: str, **kwargs) -> ModelConfig: + pass + + def cleanup(self) -> None: + pass diff --git a/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py new file mode 100644 index 000000000000..c6c88d16bdcd --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/base_weight_loader.py @@ -0,0 +1,20 @@ +from abc import ABC, abstractmethod +from typing import Any + + +class BaseWeightLoader(ABC): + + @abstractmethod + def load_weights(self, checkpoint_dir: str) -> dict[str, Any]: + """ + Loads weights from a checkpoint directory. + + Args: + checkpoint_dir: A path to the checkpoint directory. + + Returns: + A dictionary where keys are tensor names and values are the tensors. + """ + + def cleanup(self) -> None: + pass diff --git a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py new file mode 100644 index 000000000000..708493eb5321 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py @@ -0,0 +1,165 @@ +from abc import ABC, abstractmethod +from typing import Callable, List, Union + +from torch import nn + +from tensorrt_llm._torch.model_config import ModelConfig, TConfig +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM + + +class BaseWeightMapper(ABC): + + def __init__(self): + self._callbacks: list[Callable] = [] + self._mapping: dict = {} + self._skip_modules = [] + self._model: Union[nn.Module, DecoderModelForCausalLM] | None = None + self._config: TConfig | None = None + + def init_model_and_config(self, model: Union[nn.Module, + DecoderModelForCausalLM], + config: TConfig): + self._model = model + self._config = config + + if not hasattr(model, 'model_config') or not isinstance( + model.model_config, ModelConfig): + raise ValueError("model must have a model_config attribute") + if not hasattr(model, 'config'): + raise ValueError("model must have a config attribute") + + self._tp_size = 1 if model.model_config.mapping.enable_attention_dp else model.model_config.mapping.tp_size + self._num_kv_heads = model.config.num_key_value_heads if hasattr( + model.config, 'num_key_value_heads' + ) and model.config.num_key_value_heads is not None else model.config.num_attention_heads + + self.map_weights() + + def cleanup(self) -> None: + self._model = None + self._config = None + + @abstractmethod + def map_weights(self) -> None: + """ + Maps weights from TRT-LLM to a source state dictionary (e.g., Hugging Face) + """ + + @abstractmethod + def apply_callbacks(self, module: nn.Module, module_name: str, + module_names_breakdown: list[str], + weights: dict) -> list[dict]: + """ + Applies a series of transformation functions to an internal representation + of weights or to guide the mapping process. The exact behavior might depend + on the implementation (e.g., storing callbacks to be applied later). + + Args: + module: The module to apply the callbacks to + module_name: The specific module name (e.g., 'qkv_proj', 'gate_up_proj') + module_names_breakdown: List of module path components for building full paths + weights: The weights dictionary to process + """ + + def rename_by_params_map(self, params_map: dict[str, str], + weights: dict) -> dict: + """ + Rename weight keys according to regex pattern matching. + + Args: + pattern_mapping: A dictionary mapping regex patterns to replacement strings. The key is HF name pattern, and the value is corresponding TRT-LLM name pattern. + The patterns will be used to match keys in the weights dict and replace + them according to the replacement string, which can use regex backreferences. + Example: + HF name: vision_model.encoder.layers.1.self_attn.out_proj.{weight,bias} + TRT-LLM name: vision_model.encoder.layers.1.self_attn.o_proj.{weight,bias} + Then the pattern_mapping could be: + pattern_mapping = { + r'(.*?)out_proj(.*)': r'\1o_proj\2' + } + weights: A dictionary of weights + + Returns: + A dictionary of weights with renamed keys + """ + import re + + # Create a new dictionary to store the renamed weights + renamed_weights = {} + + # Keep track of keys that have been matched by a pattern + matched_keys = set() + + # Process each key in the weights dictionary + for key in list(weights.keys()): + # Check each pattern for a match + for pattern, replacement in params_map.items(): + if re.match(pattern, key): + # Create the new key by applying the regex replacement + new_key = re.sub(pattern, replacement, key) + # Store the weight with the new key + renamed_weights[new_key] = weights[key] + matched_keys.add(key) + break + + # If the key wasn't matched by any pattern, keep it as is + if key not in matched_keys: + renamed_weights[key] = weights[key] + + return renamed_weights + + def preprocess_weights(self, weights: dict) -> dict: + """ + Preprocess weights before starting the loading process. + """ + ... + + def handle_manual_copy(self, module_name: str, module_weights: dict, n: str, + p: nn.Parameter) -> None: + p.data.copy_(module_weights[n][:]) + + def does_require_special_handling(self, module_name: str) -> bool: + return module_name in self.mapping + + def is_special_instance_module(self, module: nn.Module) -> bool: + return False + + def handle_special_instance_module(self, module: nn.Module, + module_name: str, + module_weights: dict) -> None: + raise NotImplementedError() + + @property + def skip_modules(self) -> List[str]: + return self._skip_modules + + def add_skip_modules(self, value: List[str]) -> None: + self._skip_modules.extend(value) + + def should_skip_module(self, module_name: str) -> bool: + return any(skip_module in module_name + for skip_module in self._skip_modules) + + def filter_weights(self, prefix: str, weights: dict) -> dict: + result = {} + for k, v in weights.items(): + if k.startswith(prefix): + new_k = k[len(prefix) + 1:] + result[new_k] = v + return result + + @property + def mapping(self) -> dict: + return self._mapping + + @property + def config(self) -> TConfig: + if self._config is None: + raise RuntimeError("Weight mapper is not initialized") + return self._config + + @property + def model(self) -> Union[nn.Module, DecoderModelForCausalLM]: + if self._model is None: + raise RuntimeError("Weight mapper is not initialized") + return self._model diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/__init__.py b/tensorrt_llm/_torch/models/checkpoints/hf/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py new file mode 100644 index 000000000000..d00a24114745 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py @@ -0,0 +1,75 @@ +from typing import Optional + +from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import \ + BaseCheckpointLoader +from tensorrt_llm._torch.models.checkpoints.base_config_loader import \ + BaseConfigLoader +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import \ + BaseWeightLoader +from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ + BaseWeightMapper +from tensorrt_llm._torch.models.checkpoints.hf.config_loader import \ + HfConfigLoader +from tensorrt_llm._torch.models.checkpoints.hf.weight_loader import \ + HfWeightLoader +from tensorrt_llm._torch.models.modeling_utils import register_checkpoint_loader + + +@register_checkpoint_loader("HF") +class HfCheckpointLoader(BaseCheckpointLoader): + + def __init__(self, + *, + weight_loader: Optional[BaseWeightLoader] = None, + weight_mapper: Optional[BaseWeightMapper] = None, + config_loader: Optional[BaseConfigLoader] = None): + if weight_loader is None: + self._weight_loader = self.get_default_weight_loader() + else: + self._weight_loader = weight_loader + if config_loader is None: + self._config_loader = self.get_default_config_loader() + else: + self._config_loader = config_loader + self._weight_mapper = weight_mapper + self._checkpoint_format = "HF" + + def cleanup(self) -> None: + # Clean up weight mapper first as it may hold model references + if self._weight_mapper is not None: + self._weight_mapper.cleanup() + self._weight_mapper = None + + if self._weight_loader is not None: + self._weight_loader.cleanup() + self._weight_loader = None + + if self._config_loader is not None: + self._config_loader.cleanup() + self._config_loader = None + + def get_default_weight_loader(self) -> HfWeightLoader: + return HfWeightLoader() + + def get_default_config_loader(self) -> HfConfigLoader: + return HfConfigLoader() + + @property + def weight_loader(self) -> BaseWeightLoader: + return self._weight_loader + + @property + def weight_mapper(self) -> Optional[BaseWeightMapper]: + return self._weight_mapper + + @weight_mapper.setter + def weight_mapper(self, value: BaseWeightMapper) -> None: + self._weight_mapper = value + + @property + def config_loader(self) -> Optional[BaseConfigLoader]: + return self._config_loader + + @property + def checkpoint_format(self) -> str: + return self._checkpoint_format diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/config_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/config_loader.py new file mode 100644 index 000000000000..a82824db473c --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/hf/config_loader.py @@ -0,0 +1,11 @@ +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.checkpoints.base_config_loader import \ + BaseConfigLoader +from tensorrt_llm._torch.models.modeling_utils import register_config_loader + + +@register_config_loader("HF") +class HfConfigLoader(BaseConfigLoader): + + def load(self, checkpoint_dir: str, **kwargs) -> ModelConfig: + return ModelConfig.from_pretrained(checkpoint_dir, **kwargs) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/gemma3_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/gemma3_weight_mapper.py new file mode 100644 index 000000000000..a8d31d6526d9 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/hf/gemma3_weight_mapper.py @@ -0,0 +1,35 @@ +from torch import nn + +from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import \ + HfWeightMapper +from tensorrt_llm._torch.models.modeling_utils import register_mapper + + +@register_mapper("HF", "Gemma3ForCausalLM") +@register_mapper("HF", "Gemma3ForConditionalGeneration") +class Gemma3HfWeightMapper(HfWeightMapper): + + def should_skip_module(self, module_name: str) -> bool: + if self.model.config.tie_word_embeddings and module_name.startswith( + "lm_head"): + return True + + # Skip loading weights for embedding and lm_head if LoRA is enabled and has custom values + if hasattr(self.model, "model") and hasattr( + self.model.model, 'has_custom_embed_tokens' + ) and self.model.model.has_custom_embed_tokens and module_name == "model.embed_tokens": + return True + if hasattr( + self.model, 'has_custom_lm_head' + ) and self.model.has_custom_lm_head and module_name == "lm_head": + return True + + return any(skip_module in module_name + for skip_module in self._skip_modules) + + def handle_manual_copy(self, module_name: str, module_weights: dict, n: str, + p: nn.Parameter) -> None: + if 'norm' in module_name: + p.data.copy_(module_weights[n][:] + 1) + else: + super().handle_manual_copy(module_name, module_weights, n, p) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/llama4_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/llama4_weight_mapper.py new file mode 100644 index 000000000000..4889e02ccaf1 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/hf/llama4_weight_mapper.py @@ -0,0 +1,22 @@ +from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import \ + HfWeightMapper +from tensorrt_llm._torch.models.modeling_utils import register_mapper + + +@register_mapper("HF", "Llama4ForConditionalGeneration") +class Llama4HfWeightMapper(HfWeightMapper): + """ + Weight mapper for Llama4ForConditionalGeneration that handles the + 'language_model.' prefix removal from weight keys. + """ + + def filter_weights(self, prefix: str, weights: dict) -> dict: + transformed_weights = {} + for key, value in weights.items(): + if key.startswith("language_model."): + new_key = key[len("language_model."):] + transformed_weights[new_key] = value + else: + transformed_weights[key] = value + + return super().filter_weights(prefix, transformed_weights) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/mixtral_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/mixtral_weight_mapper.py new file mode 100644 index 000000000000..6bdd31383aeb --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/hf/mixtral_weight_mapper.py @@ -0,0 +1,26 @@ +from torch import nn + +from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ + BaseWeightMapper +from tensorrt_llm._torch.models.modeling_utils import register_mapper + + +@register_mapper("HF", "MixtralForCausalLM") +class MixtralHfWeightMapper(BaseWeightMapper): + + def map_weights(self) -> None: + self.mapping.update({ + 'qkv_proj': ['q_proj', 'k_proj', 'v_proj'], + }) + + def apply_callbacks(self, module: nn.Module, module_name: str, + module_names_breakdown: list[str], + weights: dict) -> list[dict]: + module_weights = [] + + for new_name in self.mapping[module_name]: + fw = self.filter_weights( + '.'.join(module_names_breakdown + [new_name]), weights) + module_weights.append(fw) + + return module_weights diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py new file mode 100644 index 000000000000..1b2ef8cf81b3 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py @@ -0,0 +1,99 @@ +import torch + +from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import \ + HfWeightMapper +from tensorrt_llm._torch.models.modeling_nemotron_h import split +from tensorrt_llm._torch.models.modeling_utils import register_mapper + + +@register_mapper("HF", "NemotronHForCausalLM") +class NemotronHHfWeightMapper(HfWeightMapper): + + def preprocess_weights(self, weights: dict) -> dict: + config = self.config.pretrained_config + tp_size = self.config.mapping.tp_size + tp_rank = self.config.mapping.tp_rank + d_inner = config.hidden_size * config.expand + n_groups = config.n_groups + d_state = config.ssm_state_size + nheads = d_inner // config.mamba_head_dim + + new_weights = {} + for name, _ in weights.items(): + key = name + + # change backbone root name to model + if "backbone" in key: + key = key.replace("backbone", "model") + + # change embedding layer to embed_token + if "embeddings" in key: + key = key.replace("embeddings", "embed_tokens") + + if "A_log" in key: + key = key.replace("A_log", "A") + + if "_scale" in key and weights[name].dim() == 0: + new_weights[key] = weights[name] + elif "A" in key: + w = split(weights[name], tp_size, tp_rank) + w = w.to(torch.float32) + w = -torch.exp(w) + new_weights[key] = w + elif "D" in key: + w = split(weights[name], tp_size, tp_rank) + w = w.to(torch.float32) + new_weights[key] = w + elif "dt_bias" in key: + w = split(weights[name], tp_size, tp_rank) + w = w.to(torch.float32) + new_weights[key] = w + elif "mixer.in_proj" in key: + w = weights[name] + in_proj_z, in_proj_x, in_proj_b, in_proj_c, in_proj_dt = torch.split( + w, [ + d_inner, d_inner, n_groups * d_state, + n_groups * d_state, nheads + ], + dim=0) + + w = [] + for rank in range(tp_size): + in_proj_z_rank = split(in_proj_z, tp_size, rank) + in_proj_x_rank = split(in_proj_x, tp_size, rank) + in_proj_b_rank = split(in_proj_b, tp_size, rank) + in_proj_c_rank = split(in_proj_c, tp_size, rank) + in_proj_dt_rank = split(in_proj_dt, tp_size, rank) + y = torch.concat([ + in_proj_z_rank, in_proj_x_rank, in_proj_b_rank, + in_proj_c_rank, in_proj_dt_rank + ]) + w.append(y) + + w = torch.concat(w).contiguous() + new_weights[key] = w + elif "conv1d" in key: + w = weights[name] + # removing dim(1) because we are using Linear to store conv1d weights + if "weight" in key: + w = w.squeeze(1) + + conv_x, conv_b, conv_c = torch.split( + w, [d_inner, n_groups * d_state, n_groups * d_state], dim=0) + + w = [] + for rank in range(tp_size): + conv_x_rank = split(conv_x, tp_size, rank) + conv_b_rank = split(conv_b, tp_size, rank) + conv_c_rank = split(conv_c, tp_size, rank) + y = torch.concat([conv_x_rank, conv_b_rank, conv_c_rank]) + w.append(y) + w = torch.concat(w).contiguous() + new_weights[key] = w + elif "mixer.norm.weight" in key: + w = split(weights[name], tp_size, tp_rank) + new_weights[key] = w + else: + new_weights[key] = weights[name] + + return new_weights diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/qwen2_moe_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/qwen2_moe_weight_mapper.py new file mode 100644 index 000000000000..41b40042c762 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen2_moe_weight_mapper.py @@ -0,0 +1,26 @@ +from torch import nn + +from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import \ + HfWeightMapper +from tensorrt_llm._torch.models.modeling_utils import register_mapper +from tensorrt_llm._torch.modules.fused_moe.interface import MoE + + +@register_mapper("HF", "Qwen2MoeForCausalLM") +class Qwen2MoeHfWeightMapper(HfWeightMapper): + + def is_special_instance_module(self, module: nn.Module) -> bool: + return isinstance(module, MoE) + + def handle_special_instance_module(self, module: nn.Module, + module_name: str, + module_weights: dict) -> None: + if isinstance(module, MoE): + updated_module_weights = {} + for weight_name, weight_value in module_weights.items(): + new_weight_name = weight_name.replace( + "gate_proj", "w1").replace("up_proj", + "w3").replace("down_proj", "w2") + updated_module_weights[new_weight_name] = weight_value + del module_weights + module.load_weights(weights=[updated_module_weights]) 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 new file mode 100644 index 000000000000..1d0763dc5a90 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_moe_weight_mapper.py @@ -0,0 +1,39 @@ +from torch import nn + +from tensorrt_llm._torch.models.checkpoints.hf.qwen2_moe_weight_mapper import \ + Qwen2MoeHfWeightMapper +from tensorrt_llm._torch.models.modeling_utils import register_mapper + + +@register_mapper("HF", "Qwen3MoeForCausalLM") +class Qwen3MoeHfWeightMapper(Qwen2MoeHfWeightMapper): + + def should_skip_module(self, module_name: str) -> bool: + if module_name.startswith("draft_model"): + return True + return super().should_skip_module(module_name) + + def _duplicate_kv_weights(self, module: nn.Module, new_name: str, + weights: dict): + tensors_to_duplicate = ["weight", "bias"] + if module.quant_config.quant_mode.has_nvfp4(): + tensors_to_duplicate.append("weight_scale") + if module.quant_config.quant_mode.has_fp8_block_scales(): + tensors_to_duplicate.append("weight_scale_inv") + + if new_name in ['k_proj', 'v_proj']: + num_kv_heads_list = [self._num_kv_heads + ] * len(weights) if isinstance( + self._num_kv_heads, + int) else self._num_kv_heads + processed_weights = { + k: + self._duplicate_kv(weight=v[:], + num_kv_heads=num_kv_heads_list[i], + tensor_parallel_size=self._tp_size) + if k in tensors_to_duplicate else v + for i, (k, v) in enumerate(weights.items()) + } + return processed_weights + + return weights diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py new file mode 100644 index 000000000000..1277e25b42f7 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -0,0 +1,123 @@ +import glob +import multiprocessing +import os +from typing import Any, List + +import psutil +import safetensors +import torch +import tqdm + +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import \ + BaseWeightLoader +from tensorrt_llm._torch.models.modeling_utils import ( + register_checkpoint_weight_loader, run_concurrently) +from tensorrt_llm._utils import local_mpi_rank, local_mpi_size +from tensorrt_llm.logger import logger + + +@register_checkpoint_weight_loader("HF") +class HfWeightLoader(BaseWeightLoader): + """ + Loads weights from SafeTensors/bin/pth files. + """ + + def load_weights(self, checkpoint_dir: str) -> dict[str, Any]: + weight_files = glob.glob(f"{checkpoint_dir}/*.safetensors") + if weight_files: + # Prefetch the weight files to CPU memory if the size is less than 90% of the available memory. + # This is a heuristic to avoid prefetching files that are too large and causing file cache thrashing. + prefetch_size = sum(os.path.getsize(file) for file in weight_files) + # If the layer number is overridden, it indicates that only a subset of layers are loaded. + # Prefetching all layers is unnecessary. + num_layers = int(os.environ.get("TLLM_OVERRIDE_LAYER_NUM", "0")) + enable_prefetch = prefetch_size < psutil.virtual_memory( + ).available * 0.9 and num_layers == 0 + if enable_prefetch: + logger.info( + f"Prefetching {prefetch_size / (1024**3):.2f}GB checkpoint files." + ) + self.prefetch_files(weight_files) + + return self._load_weights_in_parallel( + weight_files, self._load_safetensors_file, + "Loading safetensors weights in parallel") + + weight_files = glob.glob(f"{checkpoint_dir}/*.bin") + if not weight_files: + weight_files = glob.glob(f"{checkpoint_dir}/*.pth") + + if weight_files: + return self._load_weights_in_parallel( + weight_files, self._load_bin_or_path_file, + "Loading bin weights in parallel") + + raise RuntimeError(f"No weight files found in {checkpoint_dir}.") + + def _load_weights_in_parallel(self, weight_files: List[str], load_func, + description: str) -> dict[str, Any]: + """ + Load weight files in parallel using the specified loading function. + + Args: + weight_files: List of weight file paths + load_func: Function to load individual weight files + description: Description for the progress bar + + Returns: + Dictionary containing all loaded weights + """ + weights = {} + pbar = tqdm.tqdm(total=len(weight_files), desc=description) + + # Note that the function is called with a tuple of arguments, hence we need to wrap the arguments in a tuple via [(w,) for w in weight_files] + # specifically the comma right after the w is important to make it a tuple. + run_concurrently(load_func, [(w, ) for w in weight_files], + reduce_func=weights.update, + pbar=pbar) + + return weights + + @staticmethod + def _load_safetensors_file(file): + return safetensors.torch.load_file(file) + + @staticmethod + def _load_bin_or_path_file(file): + try: + part_weights = torch.load(file, + weights_only=True, + map_location='cpu', + mmap=True) + except Exception: + logger.warning( + f"Failed to load {file} with mmap=True, fallback to mmap=False") + part_weights = torch.load(file, + weights_only=True, + map_location='cpu', + mmap=False) + finally: + return part_weights + + def _prefetch_one_file(self, file_name): + if os.path.exists(file_name): + logger.info(f"Prefetching {file_name} to memory...") + with open(file_name, 'rb') as f: + f.read() + logger.info(f"Finished prefetching {file_name}.") + + def prefetch_files(self, file_names: List[str]): + """ + Prefetch safetensors files to memory so that the weight loading will be much faster. + When multiple ranks run in parallel, each rank will prefetch some files. + """ + # Find out the files to prefetch for the current rank. + # Each rank loads files with indices local_rank, local_rank + local_mpi_size, local_rank + 2*local_mpi_size, etc. + local_file_names = file_names[local_mpi_rank()::local_mpi_size()] + if len(local_file_names) == 0: + return + + max_processes = min(multiprocessing.cpu_count() * 2, 16, + len(local_file_names)) + with multiprocessing.Pool(processes=max_processes) as pool: + pool.map(self._prefetch_one_file, local_file_names) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py new file mode 100644 index 000000000000..7df6faa613ad --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py @@ -0,0 +1,101 @@ +import torch +from torch import nn + +from tensorrt_llm._torch.models.modeling_utils import register_mapper + +from ..base_weight_mapper import BaseWeightMapper + + +@register_mapper("HF") +class HfWeightMapper(BaseWeightMapper): + + def __init__(self): + super().__init__() + self._callbacks = [ + self._duplicate_kv_weights, + ] + + def map_weights(self) -> None: + self.mapping.update({ + 'qkv_proj': ['q_proj', 'k_proj', 'v_proj'], + 'gate_up_proj': ['gate_proj', 'up_proj'] + }) + + def apply_callbacks(self, module: nn.Module, module_name: str, + module_names_breakdown: list[str], + weights: dict) -> list[dict]: + module_weights = [] + + for new_name in self._mapping[module_name]: + fw = self.filter_weights( + '.'.join(module_names_breakdown + [new_name]), weights) + for callback in self._callbacks: + fw = callback(module, new_name, fw) + module_weights.append(fw) + + return module_weights + + def should_skip_module(self, module_name: str) -> bool: + if self.model.config.tie_word_embeddings and module_name.startswith( + "lm_head"): + return True + + # Skip loading weights for embedding and lm_head if LoRA is enabled and has custom values + if hasattr(self.model, "model") and hasattr( + self.model.model, 'has_custom_embed_tokens' + ) and self.model.model.has_custom_embed_tokens and module_name == "model.embed_tokens": + return True + if hasattr( + self.model, 'has_custom_lm_head' + ) and self.model.has_custom_lm_head and module_name == "lm_head": + return True + + # WAR: better solution is that llama has its own load_weights function. + if module_name.split('.')[-1] == 'next_layer_layernorm': + return True + + return super().should_skip_module(module_name) + + def _duplicate_kv_weights(self, module: nn.Module, new_name: str, + weights: dict): + if new_name in ['k_proj', 'v_proj']: + num_kv_heads_list = [self._num_kv_heads + ] * len(weights) if isinstance( + self._num_kv_heads, + int) else self._num_kv_heads + processed_weights = { + k: + self._duplicate_kv(weight=v[:], + num_kv_heads=num_kv_heads_list[i], + tensor_parallel_size=self._tp_size) + if k in ["weight", "bias"] else v + for i, (k, v) in enumerate(weights.items()) + } + return processed_weights + + return weights + + def _duplicate_kv(self, weight: torch.Tensor, num_kv_heads: int, + tensor_parallel_size: int): + + if num_kv_heads >= tensor_parallel_size: + assert num_kv_heads % tensor_parallel_size == 0 + return weight + + assert tensor_parallel_size % num_kv_heads == 0 + reps = tensor_parallel_size // num_kv_heads + + # bias + if weight.ndim == 1: + return weight.repeat_interleave(reps) + + # weight and scale + assert weight.shape[0] % num_kv_heads == 0 + size_per_kv_head = weight.shape[0] // num_kv_heads + weight = weight.reshape(num_kv_heads, size_per_kv_head, + -1)[:, + None, :, :].expand(num_kv_heads, reps, + size_per_kv_head, + weight.shape[1]) + return weight.reshape(num_kv_heads * reps * size_per_kv_head, + -1).clone().detach() diff --git a/tensorrt_llm/_torch/models/modeling_clip.py b/tensorrt_llm/_torch/models/modeling_clip.py index 546375720bf4..da2688f1e934 100644 --- a/tensorrt_llm/_torch/models/modeling_clip.py +++ b/tensorrt_llm/_torch/models/modeling_clip.py @@ -202,7 +202,7 @@ def prepare_attn_metadata(self, batch_size): request_ids=request_ids, prompt_lens=prompt_lens, ) - attn_metadata.max_seq_len = seq_len * batch_size + attn_metadata.max_seq_len = seq_len attn_metadata.prepare() return attn_metadata diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 1834e7b14767..c8523deea2e1 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -38,6 +38,7 @@ from tqdm import tqdm from transformers import PretrainedConfig +from tensorrt_llm._ipc_utils import can_access_peer from tensorrt_llm._utils import get_sm_version from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.llmapi.utils import enable_llm_debug @@ -53,8 +54,8 @@ from ..modules.attention import MLA from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import (CutlassFusedMoE, DeepSeekV3MoeRoutingMethod, - TRTLLMGenFusedMoE, WideEPMoE, create_moe, +from ..modules.fused_moe import (DeepSeekV3MoeRoutingMethod, TRTLLMGenFusedMoE, + create_moe, moe_load_balancer_set_repeated_for_next_layer) from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear, TensorParallelMode, WeightsLoadingConfig @@ -62,8 +63,7 @@ from ..modules.rms_norm import RMSNorm from ..peft.lora.layer import LoraLayer from ..speculative import MTPEagleWorker, MTPSpecMetadata, MTPWorker -from ..utils import (AuxStreamType, EventType, Fp4QuantizedTensor, - disable_fp4_allgather) +from ..utils import AuxStreamType, EventType, Fp4QuantizedTensor from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, EagerFusionConfig, filter_weights, register_auto_model) @@ -512,22 +512,11 @@ def compute_routed_output(self, hidden_states, hidden_states_fp4, # max-throughput use_dp_padding = False if self.use_dp and self.mapping.tp_size > 1: - # FP4 all_gather moves this bf16 allgather in to after topk and fp4 quantization - # to reduce allreduce BW - if (disable_fp4_allgather() - and not self.experts.enable_alltoall) or isinstance( - self.experts, TRTLLMGenFusedMoE): + if isinstance(self.experts, TRTLLMGenFusedMoE): hidden_states = allgather(hidden_states, self.mapping, dim=0, sizes=all_rank_num_tokens) - elif not isinstance(self.experts, (CutlassFusedMoE, WideEPMoE)) or ( - not self.experts.has_fp8_qdq and self.experts.has_nvfp4): - # Use padding when not using the cutlass path or when x_sf in self.experts is not None - use_dp_padding = True - hidden_states = torch.nn.functional.pad( - hidden_states, - (0, 0, 0, all_rank_max_num_tokens - hidden_states.shape[0])) router_logits = self.gate(hidden_states) @@ -614,6 +603,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], self.enable_attention_dp = mapping.enable_attention_dp self.mlp_tp_size = mapping.tp_size + self.is_p2p_supported = can_access_peer(mapping) self.fusion_config = EagerFusionConfig() self.enable_fusion = os.environ.get( @@ -803,12 +793,12 @@ def _run_MoE(hidden_states, hidden_states_fp4, do_finalize): hidden_states, residual = self.post_attention_layernorm( hidden_states, residual) - # Note: this fusion pattern is only supported for TRTLLM-nvfp4 backend now - do_finalize = not (hidden_states.shape[0] - <= self.moe_allreduce.max_token - and self.fusion_config.POST_MOE_FUSION - and self.model_config.moe_backend == 'TRTLLM' - and self.mlp.experts.has_nvfp4) + # Note: this fusion pattern is only supported for single-node TRTLLM-nvfp4 backend now + do_finalize = self.mapping.is_multi_node() or ( + not (hidden_states.shape[0] <= self.moe_allreduce.max_token + and self.fusion_config.POST_MOE_FUSION + and self.model_config.moe_backend == "TRTLLM" + and self.mlp.experts.has_nvfp4 and self.is_p2p_supported)) hidden_states = _run_MoE(hidden_states, hidden_states_fp4=None, diff --git a/tensorrt_llm/_torch/models/modeling_gemma3.py b/tensorrt_llm/_torch/models/modeling_gemma3.py index 460af83ddbff..db672279646d 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma3.py +++ b/tensorrt_llm/_torch/models/modeling_gemma3.py @@ -3,10 +3,11 @@ import torch from torch import nn -from tqdm import tqdm from transformers import Gemma3TextConfig from transformers.activations import ACT2FN +from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ + BaseWeightMapper from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType from tensorrt_llm.mapping import Mapping @@ -23,7 +24,6 @@ from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, - duplicate_kv_weight, filter_weights, register_auto_model) @@ -486,61 +486,5 @@ def forward( return_context_logits, ) - # This is a modified version of the load_weights function in modeling_utils.py with the - # minor change for Gemma3 RMSNorm. - def load_weights(self, weights: Dict): - tp_size = self.model_config.mapping.tp_size - num_kv_heads = self.config.num_key_value_heads - - params_map = { - 'qkv_proj': ['q_proj', 'k_proj', 'v_proj'], - 'gate_up_proj': ['gate_proj', 'up_proj'] - } - - for name, module in tqdm(list(self.named_modules()), - desc="Loading weights"): - if len(module._parameters) > 0: - # skip load weights if tie word embeddings is enabled and layer is lm_head - if self.config.tie_word_embeddings and name.startswith( - "lm_head"): - continue - - # Skip loading weights for embedding and lm_head if LoRA is enabled. - if hasattr( - self.model_config, 'lora_config' - ) and self.model_config.lora_config is not None and len( - self.model_config.lora_config.lora_dir) == 1 and ( - name == "model.embed_tokens" or name == "lm_head"): - continue - - names = name.split('.') - if names[-1] in params_map: - module_weights = [] - for new_name in params_map[names[-1]]: - fw = filter_weights('.'.join(names[:-1] + [new_name]), - weights) - if new_name in ['k_proj', 'v_proj']: - fw = { - k: - duplicate_kv_weight( - weight=v[:], - num_kv_heads=num_kv_heads, - tensor_parallel_size=tp_size) - if k in ["weight", "bias"] else v - for k, v in fw.items() - } - - module_weights.append(fw) - module.load_weights(weights=module_weights) - else: - module_weights = filter_weights(name, weights) - if hasattr(module, 'load_weights'): - module.load_weights(weights=[module_weights]) - else: - for n, p in module._parameters.items(): - if p is not None: - # Gemma3 RMSNorm uses +1 just like LayerNorm-1P. - if 'norm' in names[-1]: - p.data.copy_(module_weights[n][:] + 1) - else: - p.data.copy_(module_weights[n][:]) + def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper): + super().load_weights(weights, weight_mapper) diff --git a/tensorrt_llm/_torch/models/modeling_gemma3vl.py b/tensorrt_llm/_torch/models/modeling_gemma3vl.py index 44a70254ad8a..07fb5b5417bb 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma3vl.py +++ b/tensorrt_llm/_torch/models/modeling_gemma3vl.py @@ -1,3 +1,4 @@ +import copy import dataclasses import os from typing import List, Optional, Tuple @@ -7,6 +8,9 @@ from transformers.modeling_utils import no_init_weights from transformers.models.gemma3.modeling_gemma3 import Gemma3MultiModalProjector +from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ + BaseWeightMapper + from ..._utils import nvtx_range from ...inputs import (ExtraProcessedInputs, InputProcessor, TextPrompt, register_input_processor) @@ -45,18 +49,12 @@ def _preprocess(self, inputs): raise KeyError("Expected image data in multimodal data for Gemma3.") images = mm_data.get("image") - if images and len(images) != 1: - raise ValueError( - f"Expected at most one image for processing, got {len(images)}." - ) - - image = images[0] if images else None do_rescale = self.processor.image_processor.do_rescale - if isinstance(image, torch.Tensor): + if images is not None and isinstance(images[0], torch.Tensor): do_rescale = False processor_output = self.processor( text=text_prompt, - images=image, + images=images, do_rescale=do_rescale, return_tensors="pt", device=self.device).to(dtype=torch.bfloat16) @@ -104,13 +102,14 @@ def __init__(self, model_config: ModelConfig[Gemma3Config]): dtype=torch.int32, device=self._device) - self.model_config = model_config + model_config_cp = copy.deepcopy(model_config) + self.model_config = model_config_cp - llm_model_config = self.get_sub_model_config(model_config, + llm_model_config = self.get_sub_model_config(model_config_cp, "text_config") self.llm = Gemma3ForCausalLM(llm_model_config) - vision_model_config = self.get_sub_model_config(model_config, + vision_model_config = self.get_sub_model_config(model_config_cp, "vision_config") self.siglip_tower = SiglipVisionModel(vision_model_config, use_post_layernorm=True) @@ -147,9 +146,9 @@ def get_sub_model_config( sub_model_config.pretrained_config.torch_dtype = model_config.pretrained_config.torch_dtype return sub_model_config - def load_weights(self, weights): + def load_weights(self, weights, weight_mapper: BaseWeightMapper): llm_weights = filter_weights("language_model", weights) - self.llm.load_weights(llm_weights) + self.llm.load_weights(llm_weights, weight_mapper) vit_weights = filter_weights("vision_tower", weights) self.siglip_tower.load_weights(vit_weights) diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index 1c17eeb5a8e8..33dddfc784c4 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/tensorrt_llm/_torch/models/modeling_llama.py @@ -11,6 +11,8 @@ from tensorrt_llm._torch.distributed import (AllReduce, AllReduceFusionOp, AllReduceParams, MoEAllReduce) +from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ + BaseWeightMapper from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.logger import logger from tensorrt_llm.lora_manager import HfLoraLoader @@ -303,13 +305,6 @@ def __init__( def compute_routed_output(self, hidden_states, all_rank_num_tokens, all_rank_max_num_tokens, cutlass_min_latency_mode): - use_dp_padding = False - if self.enable_attention_dp and self.mapping.tp_size > 1: - # Use padding here to keep the behavior unchanged - use_dp_padding = True - hidden_states = torch.nn.functional.pad( - hidden_states, - (0, 0, 0, all_rank_max_num_tokens - hidden_states.shape[0])) router_logits = self.router(hidden_states) routed_output = self.experts( hidden_states, @@ -317,8 +312,7 @@ def compute_routed_output(self, hidden_states, all_rank_num_tokens, do_finalize=not cutlass_min_latency_mode, all_rank_num_tokens=all_rank_num_tokens, all_rank_max_num_tokens=all_rank_max_num_tokens, - use_dp_padding=use_dp_padding, - ) + use_dp_padding=False) return routed_output def forward( @@ -709,11 +703,13 @@ def __init__(self, model_config: ModelConfig[LlamaConfig]): model_config, 'lora_config') and model_config.lora_config is not None and len( model_config.lora_config.lora_dir) == 1: - lora_loader = HfLoraLoader(model_config.lora_config.lora_dir) - if lora_loader.vocab_size != 0 and lora_loader.embed_tokens is not None: - vocab_size = lora_loader.vocab_size - weight = lora_loader.embed_tokens - self.has_custom_embed_tokens = True + # Only check for custom vocab in HF LoRA, not NeMo + if model_config.lora_config.lora_ckpt_source == "hf": + lora_loader = HfLoraLoader(model_config.lora_config.lora_dir) + if lora_loader.vocab_size != 0 and lora_loader.embed_tokens is not None: + vocab_size = lora_loader.vocab_size + weight = lora_loader.embed_tokens + self.has_custom_embed_tokens = True if self.model_config.mapping.enable_attention_dp: self.embed_tokens = Embedding( @@ -917,16 +913,8 @@ def infer_max_seq_len(self): return super().infer_max_seq_len() - def load_weights(self, weights: Dict): - new_weights = {} - for key, tensor in weights.items(): - if key.startswith("language_model."): - new_key = key[len("language_model."):] - new_weights[new_key] = tensor - else: - new_weights[key] = tensor - - super().load_weights(new_weights) + def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper): + super().load_weights(weights, weight_mapper) for idx, layer in enumerate( self.model.layers[:self.config.num_hidden_layers]): diff --git a/tensorrt_llm/_torch/models/modeling_mistral.py b/tensorrt_llm/_torch/models/modeling_mistral.py index 594ba4a56cf9..785b93fdb67f 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/tensorrt_llm/_torch/models/modeling_mistral.py @@ -3,6 +3,7 @@ from typing import Any, Dict, List, Optional, Tuple import torch +import torchvision from torch import nn from transformers import (AutoProcessor, AutoTokenizer, Mistral3Config, MistralConfig, PretrainedConfig, PreTrainedModel) @@ -226,7 +227,6 @@ def __init__( self.model_config = model_config self.tokenizer = tokenizer - self._device = "cuda" self._processor = AutoProcessor.from_pretrained(model_path, use_fast=False) @@ -256,7 +256,6 @@ def __call__( if pixel_values is not None: # We have no use for the `attention_mask`. processed.pop("attention_mask") - processed = processed.to(self._device) # NOTE: `processed` is a dict-like object, but not actually a dict. extra_processed_inputs = { "multimodal_data": { @@ -296,6 +295,8 @@ def __init__( llm_model_config = self._get_sub_model_config(model_config, "text_config") + # This is necessary for the auto weight mapper to figure out what it needs. + llm_model_config.pretrained_config.architectures = config.architectures self.llm = MistralForCausalLM(llm_model_config) self._device = "cuda" @@ -345,7 +346,6 @@ def forward( attn_metadata: AttentionMetadata, input_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, return_context_logits: bool = False, **kwargs, ) -> torch.Tensor: @@ -361,20 +361,26 @@ def forward( raise RuntimeError( f"Number of multimodal tensors ({multimodal_params_len}) should be equal to number of " f"context requests ({num_context_requests}) in the batch.") - # NOTES: - # 1. the pixel values in `multimodal_data["image"]` might vary in (height, width) between - # images, making them unsafe to batch in general. The input processor also cannot produce - # them in a batch, since it is always called with a single input - otherwise, we would - # have been able to naturally leverage the padding / resizing capabilities of the underlying - # `PixtralProcessor`. - # 2. After each `pixel_values` tensor has gone through the vision tower's `patch_conv` layer, - # they are divided into patches that are then concatenated in order to treat them as a - # single "sequence" in the vision tower's attention layers, so some form of batching still - # happens in the vision tower. - image_features = [ - self._get_image_features(**x.multimodal_data["image"]) + pixel_values = [ + x.multimodal_data["image"]["pixel_values"] for x in multimodal_params ] + image_sizes = [ + x.multimodal_data["image"]["image_sizes"] + for x in multimodal_params + ] + if not (len(pixel_values) == len(image_sizes) == + multimodal_params_len): + raise ValueError( + f"Expected as many `pixel_values` ({len(pixel_values)}) and " + f"`image_sizes` ({len(image_sizes)}) as number of multimodal parameters " + f"({multimodal_params_len}).") + batched_pixel_values, batched_image_sizes = self._batch_pixel_values( + pixel_values=pixel_values, image_sizes=image_sizes) + image_features = [ + self._get_image_features(pixel_values=batched_pixel_values, + image_sizes=batched_image_sizes) + ] input_ids, inputs_embeds = fuse_input_embeds( embedding_layer=self.llm.model.embed_tokens, @@ -427,6 +433,31 @@ def _get_image_features( image_sizes) return image_features + # Original HF implementation: + # https://github.com/huggingface/transformers/blob/v4.51.3/src/transformers/models/pixtral/ + # image_processing_pixtral.py#L276 + # We switch to using torchvision's padding functionality since it supports torch tensors + # (the transformers one expected numpy arrays). + @staticmethod + @torch.inference_mode() + def _batch_pixel_values( + pixel_values: List[torch.Tensor], + image_sizes: List[torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + batched_image_sizes = torch.cat(image_sizes) + max_shape = batched_image_sizes.max(dim=0).values + pixel_values = [ + torchvision.transforms.v2.functional.pad( + image, + # Per torchvision docs, this should be in LTRB order if it's a sequence of 4 numbers. + padding=[0, 0, max_shape[1] - size[1], max_shape[0] - size[0]], + # Values extracted from HF implementation. + fill=0.0, + padding_mode="constant", + ) for image, size in zip(pixel_values, batched_image_sizes) + ] + return torch.cat(pixel_values), batched_image_sizes + # Original implementation: # https://github.com/huggingface/transformers/blob/v4.51.3/src/transformers/models/mistral3/modeling_mistral3.py#L66 diff --git a/tensorrt_llm/_torch/models/modeling_mixtral.py b/tensorrt_llm/_torch/models/modeling_mixtral.py index bb1459b4ab1c..e16b82020bd7 100644 --- a/tensorrt_llm/_torch/models/modeling_mixtral.py +++ b/tensorrt_llm/_torch/models/modeling_mixtral.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional +from typing import Optional import torch from torch import nn @@ -16,7 +16,7 @@ from ..modules.linear import Linear from ..modules.rms_norm import RMSNorm from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, - filter_weights, register_auto_model) + register_auto_model) class MixtralMoE(nn.Module): @@ -62,20 +62,13 @@ def forward( ) -> torch.Tensor: all_rank_num_tokens = attn_metadata.all_rank_num_tokens all_rank_max_num_tokens = attn_metadata.all_rank_max_num_tokens - use_dp_padding = False - if self.enable_attention_dp and len(all_rank_num_tokens) > 1: - # Use padding here to keep the behavior unchanged - use_dp_padding = True - hidden_states = torch.nn.functional.pad( - hidden_states, - (0, 0, 0, all_rank_max_num_tokens - hidden_states.shape[0])) router_logits = self.gate(hidden_states) final_hidden_states = self.experts( hidden_states, router_logits, all_rank_num_tokens=all_rank_num_tokens, all_rank_max_num_tokens=all_rank_max_num_tokens, - use_dp_padding=use_dp_padding) + use_dp_padding=False) return final_hidden_states @@ -215,27 +208,3 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]): config=model_config, hidden_size=model_config.pretrained_config.hidden_size, vocab_size=model_config.pretrained_config.vocab_size) - - def load_weights(self, weights: Dict): - - params_map = { - 'qkv_proj': ['q_proj', 'k_proj', 'v_proj'], - } - - for name, module in self.named_modules(): - if len(module._parameters) > 0: - names = name.split('.') - if names[-1] in params_map: - module_weights = [] - for new_name in params_map[names[-1]]: - module_weights.append( - filter_weights('.'.join(names[:-1] + [new_name]), - weights)) - module.load_weights(weights=module_weights) - else: - module_weights = filter_weights(name, weights) - if hasattr(module, 'load_weights'): - module.load_weights(weights=[module_weights]) - else: - for n, p in module.named_parameters(): - p.data.copy_(module_weights[n][:]) diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py index b23f5e733d23..d6387f819084 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py @@ -26,6 +26,83 @@ from torchvision.transforms import Normalize, Resize, ToTensor from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm.inputs.multimodal import MultimodalParams +from tensorrt_llm.logger import logger + + +def find_uncached_mm_embeds( + mm_embeds: List[torch.Tensor], + multimodal_params: List[MultimodalParams]) -> torch.Tensor: + """ + Find the uncached multimodal mm_embeds from multimodal_params for each batch. + Args: + - mm_embeds: List[torch.Tensor] + - multimodal_params: List[MultimodalParams] + Returns: + - sliced_mm_embeds: List[torch.Tensor] + When kv_cache reuse is disabled or model not enabled/support kv_cache reuse, return the full mm_embeds. + Note: + - Current implementation assumes chunk prefill is disabled. To support chunk prefill, we might need to slightly modify the logic (see TODO below). + """ + # Current support two batching modes: + # 1. Pre-concatenated mm_embeds for each batch, i.e., len(mm_embeds) == 1 + # 2. Individual mm_embeds for each multimodal param, i.e., len(mm_embeds) == len(multimodal_params) + if len(mm_embeds) > 1 and len(mm_embeds) != len(multimodal_params): + raise ValueError( + f"Number of mm_embeds ({len(mm_embeds)}) does not match number of multimodal params ({len(multimodal_params)})." + ) + + if not multimodal_params or multimodal_params[0].multimodal_runtime is None: + # No slicing, return the full mm_embeds + return mm_embeds + + total_cached_mm_tokens = sum([ + param.multimodal_runtime.num_cached_mm_tokens + for param in multimodal_params + ]) + if total_cached_mm_tokens == 0: + # No cached tokens, return the full mm_embeds + # TODO: support chunk prefill for multimodal, then we need to extract full mm_embeds for each CHUNK + logger.debug( + "No multimodal cached tokens can be reused, return the full mm_embeds" + ) + return mm_embeds + + if total_cached_mm_tokens == sum([ + param.multimodal_runtime.total_mm_tokens + for param in multimodal_params + ]): + # All tokens are cached, return empty list + logger.debug( + "All multimodal tokens cached, skipping vision encoder forward") + return [] + + # Partial caching, return the sliced mm_embeds + current_pos = 0 + slices = [] + for param in multimodal_params: + runtime = param.multimodal_runtime + slices.append((current_pos + runtime.num_cached_mm_tokens, + current_pos + runtime.total_mm_tokens)) + if len(mm_embeds + ) == 1: # pre-concatenated mm_embeds, need global offset + current_pos += runtime.total_mm_tokens + + sliced_mm_embeds = [] + if len(mm_embeds) == 1: + for start, end in slices: + sliced_mm_embeds.append(mm_embeds[0][start:end]) + else: # slice each mm_embeds individually + for i, (start, end) in enumerate(slices): + sliced_mm_embeds.append(mm_embeds[i][start:end]) + + if len(mm_embeds) == 1: + sliced_mm_embeds = [torch.cat(sliced_mm_embeds, dim=0)] + + logger.debug( + f"Partial caching, return sliced_mm_embeds: {sliced_mm_embeds[0].shape}" + ) + return sliced_mm_embeds def fuse_input_embeds( @@ -64,10 +141,17 @@ def fuse_input_embeds( mm_token_mask = input_ids >= vocab_size text_token_mask = input_ids < vocab_size else: + mm_token_ids = mm_token_ids.to(input_ids.device) mm_token_mask = torch.isin(input_ids, mm_token_ids) text_token_mask = ~mm_token_mask text_token_indices = torch.where(text_token_mask)[0] mm_token_indices = torch.where(mm_token_mask)[0] + if len(mm_token_indices) != mm_embed.shape[0]: + raise ValueError( + f"Multimodal token count mismatch: found {len(mm_token_indices)} image tokens in input_ids " + f"but received {mm_embed.shape[0]} image embeddings. " + "This is likely due to KV cache reuse, chunk prefill, or other optimizations that " + "cause token count mismatches within the inference batch.") text_embed = embedding_layer(input_ids[text_token_indices]) input_embeds = torch.empty(input_ids.shape[0], diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 0e9f9a033612..e19a8dc6ea9e 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -13,13 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, Optional +from typing import Optional import torch from torch import nn from torch.nn import functional as F from transformers import AutoConfig, PretrainedConfig +from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ + BaseWeightMapper from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata from ..attention_backend import AttentionMetadata @@ -254,94 +256,9 @@ def __init__( vocab_size=model_config.pretrained_config.vocab_size, ) - def load_weights(self, weights: Dict): - config = self.model_config.pretrained_config - tp_size = self.model_config.mapping.tp_size - tp_rank = self.model_config.mapping.tp_rank - d_inner = config.hidden_size * config.expand - n_groups = config.n_groups - d_state = config.ssm_state_size - nheads = d_inner // config.mamba_head_dim - - new_weights = {} - for name, params in weights.items(): - key = name - - # change backbone root name to model - if "backbone" in key: - key = key.replace("backbone", "model") - - # change embedding layer to embed_token - if "embeddings" in key: - key = key.replace("embeddings", "embed_tokens") - - if "A_log" in key: - key = key.replace("A_log", "A") - - if "_scale" in key and weights[name].dim() == 0: - new_weights[key] = weights[name] - elif "A" in key: - w = split(weights[name], tp_size, tp_rank) - w = w.to(torch.float32) - w = -torch.exp(w) - new_weights[key] = w - elif "D" in key: - w = split(weights[name], tp_size, tp_rank) - w = w.to(torch.float32) - new_weights[key] = w - elif "dt_bias" in key: - w = split(weights[name], tp_size, tp_rank) - w = w.to(torch.float32) - new_weights[key] = w - elif "mixer.in_proj" in key: - w = weights[name] - in_proj_z, in_proj_x, in_proj_b, in_proj_c, in_proj_dt = torch.split( - w, [ - d_inner, d_inner, n_groups * d_state, - n_groups * d_state, nheads - ], - dim=0) - - w = [] - for rank in range(tp_size): - in_proj_z_rank = split(in_proj_z, tp_size, rank) - in_proj_x_rank = split(in_proj_x, tp_size, rank) - in_proj_b_rank = split(in_proj_b, tp_size, rank) - in_proj_c_rank = split(in_proj_c, tp_size, rank) - in_proj_dt_rank = split(in_proj_dt, tp_size, rank) - y = torch.concat([ - in_proj_z_rank, in_proj_x_rank, in_proj_b_rank, - in_proj_c_rank, in_proj_dt_rank - ]) - w.append(y) - - w = torch.concat(w).contiguous() - new_weights[key] = w - elif "conv1d" in key: - w = weights[name] - # removing dim(1) because we are using Linear to store conv1d weights - if "weight" in key: - w = w.squeeze(1) - - conv_x, conv_b, conv_c = torch.split( - w, [d_inner, n_groups * d_state, n_groups * d_state], dim=0) - - w = [] - for rank in range(tp_size): - conv_x_rank = split(conv_x, tp_size, rank) - conv_b_rank = split(conv_b, tp_size, rank) - conv_c_rank = split(conv_c, tp_size, rank) - y = torch.concat([conv_x_rank, conv_b_rank, conv_c_rank]) - w.append(y) - w = torch.concat(w).contiguous() - new_weights[key] = w - elif "mixer.norm.weight" in key: - w = split(weights[name], tp_size, tp_rank) - new_weights[key] = w - else: - new_weights[key] = weights[name] - - super().load_weights(new_weights) + def load_weights(self, weights: dict, weight_mapper: BaseWeightMapper): + new_weights = weight_mapper.preprocess_weights(weights) + super().load_weights(new_weights, weight_mapper) AutoConfig.register(NemotronHConfig.model_type, NemotronHConfig) diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_nas.py b/tensorrt_llm/_torch/models/modeling_nemotron_nas.py index 146d13f16f1e..3ab1cdb37ca9 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_nas.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_nas.py @@ -192,11 +192,13 @@ def __init__(self, model_config): model_config, 'lora_config') and model_config.lora_config is not None and len( model_config.lora_config.lora_dir) == 1: - lora_loader = HfLoraLoader(model_config.lora_config.lora_dir) - if lora_loader.vocab_size != 0 and lora_loader.embed_tokens is not None: - vocab_size = lora_loader.vocab_size - weight = lora_loader.embed_tokens - self.has_custom_embed_tokens = True + # Only check for custom vocab in HF LoRA, not NeMo + if model_config.lora_config.lora_ckpt_source == "hf": + lora_loader = HfLoraLoader(model_config.lora_config.lora_dir) + if lora_loader.vocab_size != 0 and lora_loader.embed_tokens is not None: + vocab_size = lora_loader.vocab_size + weight = lora_loader.embed_tokens + self.has_custom_embed_tokens = True self.embed_tokens = Embedding( vocab_size, diff --git a/tensorrt_llm/_torch/models/modeling_phi3.py b/tensorrt_llm/_torch/models/modeling_phi3.py new file mode 100644 index 000000000000..5e4221dd713a --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_phi3.py @@ -0,0 +1,249 @@ +from typing import Optional, Tuple + +import torch +from torch import nn +from tqdm import tqdm +from transformers import Phi3Config + +from tensorrt_llm._torch.attention_backend import AttentionMetadata +from tensorrt_llm._torch.attention_backend.interface import ( + PositionalEmbeddingParams, RopeParams) +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_utils import (DecoderModel, + DecoderModelForCausalLM, + register_auto_model) +from tensorrt_llm._torch.modules.attention import Attention +from tensorrt_llm._torch.modules.decoder_layer import DecoderLayer +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm.functional import PositionEmbeddingType + + +class Phi3Attention(Attention): + + def __init__( + self, + model_config: ModelConfig[Phi3Config], + layer_idx: Optional[int] = None, + ): + config = model_config.pretrained_config + + rope_params = RopeParams.from_config(config) + super().__init__( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + max_position_embeddings=config.max_position_embeddings, + bias=config.attention_bias, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=rope_params, + ), + layer_idx=layer_idx, + dtype=config.torch_dtype, + config=model_config, + ) + + +class Phi3DecoderLayer(DecoderLayer): + + def __init__( + self, + model_config: ModelConfig[Phi3Config], + layer_idx: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + super().__init__() + config = model_config.pretrained_config + self.layer_idx = layer_idx + + self.self_attn = Phi3Attention(model_config, layer_idx=layer_idx) + + self.mlp = GatedMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + bias=False, + dtype=config.torch_dtype, + config=model_config, + ) + + self.input_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + self.post_attention_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + + def forward( + self, + position_ids: torch.LongTensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + residual: Optional[torch.Tensor], + lora_params=None, + **kwargs, + ) -> torch.Tensor: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm( + hidden_states, residual) + + # Self Attention + hidden_states = self.self_attn( + position_ids=None, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + lora_params=lora_params, + **kwargs, + ) + + # Fully connected + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual) + hidden_states = self.mlp(hidden_states, **kwargs) + return hidden_states, residual + + +class Phi3Model(DecoderModel): + + def __init__(self, model_config: ModelConfig[Phi3Config]): + super().__init__(model_config) + config = self.model_config.pretrained_config + self.padding_idx = config.pad_token_id + + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + dtype=config.torch_dtype, + mapping=model_config.mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + ) + self.layers = nn.ModuleList([ + Phi3DecoderLayer( + model_config, + layer_idx, + ) for layer_idx in range(config.num_hidden_layers) + ]) + self.norm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + lora_params=None, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + hidden_states = inputs_embeds + + residual = None + for decoder_layer in self.layers: + hidden_states, residual = decoder_layer( + hidden_states=hidden_states, + position_ids=position_ids, + residual=residual, + attn_metadata=attn_metadata, + lora_params=lora_params, + ) + + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +@register_auto_model("Phi3ForCausalLM") +class Phi3ForCausalLM(DecoderModelForCausalLM[Phi3Model, Phi3Config]): + + def __init__( + self, + model_config: ModelConfig[Phi3Config], + ): + super().__init__(Phi3Model(model_config), + config=model_config, + hidden_size=model_config.pretrained_config.hidden_size, + vocab_size=model_config.pretrained_config.vocab_size) + + def load_weights(self, weights: dict): + self.model_config.mapping.tp_size + hidden_size = self.config.hidden_size + num_heads = self.config.num_attention_heads + num_kv_heads = self.config.num_key_value_heads + head_dim = hidden_size // num_heads + + def filter_weights(prefix: str, weights: dict): + result = {} + for k, v in weights.items(): + if k.startswith(prefix): + new_k = k[len(prefix) + 1:] + result[new_k] = v + return result + + for name, module in tqdm(list(self.named_modules()), + desc="Loading weights"): + if len(module._parameters) > 0: + # skip load weights if tie word embeddings is enabled and layer is lm_head + if self.config.tie_word_embeddings and name.startswith( + 'lm_head'): + continue + + module_weights = filter_weights(name, weights) + if hasattr(module, 'load_weights'): + if "self_attn.qkv_proj" in name: + # The weights need to be split correctly before sharding to support tp_size >1. + qkv_weight = module_weights['weight'][:] + q_weight = qkv_weight[:hidden_size, :] + k_weight = qkv_weight[hidden_size:hidden_size + + num_kv_heads * head_dim, :] + v_weight = qkv_weight[hidden_size + + num_kv_heads * head_dim:, :] + module.load_weights(weights=[ + { + 'weight': q_weight + }, + { + 'weight': k_weight + }, + { + 'weight': v_weight + }, + ]) + elif "mlp.gate_up_proj" in name: + # The weights need to be split correctly before sharding to support tp_size >1. + intermediate_size = self.config.intermediate_size + gate_up_weight = module_weights['weight'][:] + gate_weight = gate_up_weight[:intermediate_size, :] + up_weight = gate_up_weight[intermediate_size:, :] + module.load_weights(weights=[ + { + 'weight': gate_weight + }, + { + 'weight': up_weight + }, + ]) + else: + module.load_weights(weights=[module_weights]) + else: + for n, p in module._parameters.items(): + if p is not None: + p.data.copy_(module_weights[n][:]) diff --git a/tensorrt_llm/_torch/models/modeling_phi4mm.py b/tensorrt_llm/_torch/models/modeling_phi4mm.py new file mode 100644 index 000000000000..8c8982f6e0be --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_phi4mm.py @@ -0,0 +1,286 @@ +# Plan for phi4-mm model support. +# (done) step 1: support legacy inference pipeline for phi4-mm model. +# (todo) step 2: refactor the inference pipeline to use AGGREGATE mode (https://github.com/NVIDIA/TensorRT-LLM/pull/5522). + +import copy +from typing import List, Optional, Tuple + +import torch +import transformers +from PIL import Image + +from ...executor.request import LoRARequest +from ...inputs import (ExtraProcessedInputs, InputProcessor, TextPrompt, + register_input_processor) +from ...logger import logger +from ...lora_manager import LoraConfig +from ...sampling_params import SamplingParams +from ..attention_backend import AttentionMetadata +from ..model_config import ModelConfig +from .modeling_auto import AutoModelForCausalLM +from .modeling_multimodal_utils import fuse_input_embeds +from .modeling_utils import register_auto_model + +# Special tokens +_IMAGE_SPECIAL_TOKEN_ID = 200010 # '<|endoftext10|>' +_AUDIO_SPECIAL_TOKEN_ID = 200011 # '<|endoftext11|>' + + +# Create a PreTrainedModel class for transformers=4.53.1 upgrade. +# Core idea is to provide `prepare_inputs_for_generation` method from `GenerationMixin`. +class NewPreTrainedModel(transformers.modeling_utils.PreTrainedModel, + transformers.generation.GenerationMixin): + pass + + +class Phi4MMInputProcessor(InputProcessor): + + def __init__(self, + model_path: str, + model_config: transformers.PretrainedConfig, + tokenizer: transformers.AutoTokenizer, + trust_remote_code: bool = True): + assert trust_remote_code, "trust_remote_code must be True for Phi4MM" + + self.model_config = model_config + self.device = 'cuda' + + self.tokenizer = tokenizer + self.use_fast = True + if self.tokenizer is None: + self.tokenizer = transformers.AutoTokenizer.from_pretrained( + model_path, + trust_remote_code=trust_remote_code, + use_fast=self.use_fast) + + self.processor = transformers.AutoProcessor.from_pretrained( + model_path, + trust_remote_code=trust_remote_code, + use_fast=self.use_fast) + + # Build pure-pytorch model architecture for multimodal encoder. + # Model weights are also loaded here. + OldPreTrainedModel = transformers.modeling_utils.PreTrainedModel + transformers.modeling_utils.PreTrainedModel = NewPreTrainedModel + # TODO: Make separate Phi4VisionEncoder and Phi4AudioEncoder, and move them to LLM-side. + ref_phi4mm_model = transformers.AutoModelForCausalLM.from_pretrained( + model_path, + trust_remote_code=True, + # Flash_attn_2 only supports bf16 or fp16 and set in HF config. + torch_dtype='auto', + _attn_implementation='flash_attention_2', + ).eval() + transformers.modeling_utils.PreTrainedModel = OldPreTrainedModel + self.phi4mm_modal_encoder = ref_phi4mm_model.model.embed_tokens_extend.to( + self.device) + # Required by Phi4MMImageAudioEmbedding. + # See link: https://huggingface.co/microsoft/Phi-4-multimodal-instruct/blob/main/modeling_phi4mm.py#L701 + self.phi4mm_wte = ref_phi4mm_model.model.embed_tokens.to(self.device) + + @torch.inference_mode() + def __call__( + self, inputs: TextPrompt, sampling_params: SamplingParams + ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: + text_prompt, mm_data, mm_processor_kwargs = inputs.get("prompt"), \ + inputs.get("multi_modal_data", {}), inputs.get("mm_processor_kwargs", {}) + images = mm_data.get("image", None) + audios = mm_data.get("audio", None) + + if images is not None: + if isinstance(images[0], torch.Tensor): + # Convert normalized tensors (0-1) to PIL images (0-255). + images = [ + Image.fromarray((image.permute(1, 2, 0) * 255).to( + torch.uint8).cpu().numpy()) for image in images + ] + + # Preprocessing for multimodal data. + inputs = self.processor(text=[text_prompt], + images=images, + audios=audios, + return_tensors='pt').to(self.device) + + # Set audio_projection_mode according to the modality. + # Ref: https://huggingface.co/microsoft/Phi-4-multimodal-instruct/blob/main/modeling_phi4mm.py#L2103 + if images is not None: + audio_projection_mode = 'vision' + elif audios is not None: + audio_projection_mode = 'speech' + else: + audio_projection_mode = 'speech' + + # Processing with Phi4MMImageAudioEmbedding. + mm_features = self.phi4mm_modal_encoder( + input_ids=inputs['input_ids'], + input_embeds=None, + input_image_embeds=inputs['input_image_embeds'], + input_audio_embeds=inputs['input_audio_embeds'], + image_sizes=inputs['image_sizes'], + image_attention_mask=inputs['image_attention_mask'], + audio_embed_sizes=inputs['audio_embed_sizes'], + audio_attention_mask=inputs['audio_attention_mask'], + audio_projection_mode=audio_projection_mode, + wte=self.phi4mm_wte, + ) + + # Postprocessing to get multimodal-only embeddings. + image_token_mask = inputs['input_ids'] == _IMAGE_SPECIAL_TOKEN_ID + audio_token_mask = inputs['input_ids'] == _AUDIO_SPECIAL_TOKEN_ID + mm_token_mask = image_token_mask | audio_token_mask + mm_features = mm_features[mm_token_mask] + + multimodal_data = {} + multimodal_data["multimodal_embedding"] = mm_features + + return inputs['input_ids'][0].to(torch.int32).tolist(), { + "multimodal_data": multimodal_data, + } + + +@register_auto_model("Phi4MMForCausalLM") +@register_input_processor(Phi4MMInputProcessor, model_type="phi4mm") +class Phi4MMForCausalLM(transformers.PreTrainedModel): + + _supports_flash_attn_2 = True + MM_TOKEN_IDS = torch.tensor( + [_IMAGE_SPECIAL_TOKEN_ID, _AUDIO_SPECIAL_TOKEN_ID]) + + def __init__(self, model_config: ModelConfig): + + config = model_config.pretrained_config + super().__init__(config) + + self.model_config = model_config + if hasattr(self, "llm"): + return + + # We use Phi3ForCausalLM as the language model. + llm_model_config = copy.deepcopy(model_config) + llm_model_config.pretrained_config.architectures = ["Phi3ForCausalLM"] + # Only build the language model architecture without loading weights. + self.llm = AutoModelForCausalLM.from_config(llm_model_config) + + self.vocab_size = config.vocab_size + self.model_dtype = getattr(config, "torch_dtype", torch.float16) + logger.info(f"{self.dtype=} {self.model_dtype=}") + self.post_config() + self.is_loaded = True + + def load_weights(self, weights): + # Filter out non-language model weights. + weights = { + k: v + for k, v in weights.items() + if not k.startswith('model.embed_tokens_extend') + } + # Filter out LoRA weights. + # LoRA weights will be loaded by LoraManager. + weights = {k: v for k, v in weights.items() if '.lora_' not in k} + # Rename base layer weights. + updated_weights = {} + for k in weights.keys(): + if 'base_layer.weight' in k: + new_k = k.replace('base_layer.weight', 'weight') + updated_weights[new_k] = weights[k] + else: + updated_weights[k] = weights[k] + weights = updated_weights + + self.llm.load_weights(weights) + + def infer_max_seq_len(self) -> int: + return self.llm.infer_max_seq_len() + + def post_config(self): + # use llm.config as config for pytorch model engine + self.config = self.llm.config + self.model_config.pretrained_config = self.llm.config + + @torch.inference_mode() + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + input_embeds: Optional[torch.Tensor] = None, + return_context_logits: bool = False, + **kwargs, + ) -> torch.Tensor: + """ + VLM forward logic with inflight batching support. + """ + num_context_requests, num_generation_requests = attn_metadata.num_contexts, attn_metadata.num_generations + logger.debug( + f"num_context_requests: {num_context_requests}, num_generation_requests: {num_generation_requests}" + ) + + multimodal_params = kwargs.get("multimodal_params", []) + mm_embedding = [ + multimodal_param.multimodal_data["multimodal_embedding"] + for multimodal_param in multimodal_params + ] + input_ids, input_embeds = fuse_input_embeds( + self.llm.model.embed_tokens, + input_ids, + mm_embedding, + mm_token_ids=self.MM_TOKEN_IDS, + ) + + output_prob = self.llm.forward( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=input_embeds, + return_context_logits=return_context_logits, + lora_params=kwargs.get("lora_params", None), + ) + + logger.debug(f'output shape: {output_prob.shape}') + return output_prob + + @staticmethod + def lora_config(model_dir: str): + _lora_config = LoraConfig( + lora_dir=[ + f"{model_dir}/vision-lora", + f"{model_dir}/speech-lora", + ], + lora_target_modules=[ + "attn_qkv", + "attn_dense", + "mlp_h_to_4h", + "mlp_4h_to_h", + ], + trtllm_modules_to_hf_modules={ + "attn_qkv": "qkv_proj", + "attn_dense": "o_proj", + "mlp_h_to_4h": "gate_up_proj", + "mlp_4h_to_h": "down_proj", + }, + max_lora_rank=320, # Max rank for Phi4MM. + ) + return _lora_config + + @staticmethod + def lora_request(num_requests: int, modality: str, base_model_dir: str): + # Prepare LoRA requests for different modalities. + # Ref: https://huggingface.co/microsoft/Phi-4-multimodal-instruct/blob/main/modeling_phi4mm.py#L2103 + lora_request = None + if modality == "image" or modality == "image_audio": + lora_request = [ + LoRARequest( + lora_name=f"vision-lora-{i}", + lora_int_id=i, + lora_path=f"{base_model_dir}/vision-lora", + ) for i in range(num_requests) + ] + elif modality == "audio": + lora_request = [ + LoRARequest( + lora_name=f"speech-lora-{i}", + lora_int_id=i, + lora_path=f"{base_model_dir}/speech-lora", + ) for i in range(num_requests) + ] + + return lora_request diff --git a/tensorrt_llm/_torch/models/modeling_pixtral.py b/tensorrt_llm/_torch/models/modeling_pixtral.py index b5f18b0a356e..273ff0a5040b 100644 --- a/tensorrt_llm/_torch/models/modeling_pixtral.py +++ b/tensorrt_llm/_torch/models/modeling_pixtral.py @@ -106,11 +106,18 @@ def forward( class PixtralTransformer(torch.nn.Module): def __init__(self, config: model_config_lib.ModelConfig[transformers.PixtralVisionConfig]): super().__init__() + tp_size = config.mapping.tp_size + num_heads = config.pretrained_config.num_attention_heads + if (num_heads % tp_size) > 0: + raise ValueError(f"{tp_size=} must divide {num_heads=}.") + num_heads //= tp_size + + self._head_dim = config.pretrained_config.head_dim + self._num_heads = num_heads + self.layers = torch.nn.ModuleList() for i in range(config.pretrained_config.num_hidden_layers): self.layers.append(PixtralAttentionLayer(config=config, layer_idx=i)) - self._head_dim = config.pretrained_config.head_dim - self._num_heads = config.pretrained_config.num_attention_heads def forward( self, @@ -165,12 +172,6 @@ def __init__( self, model_config: model_config_lib.ModelConfig[transformers.PixtralVisionConfig] ): super().__init__() - tp_size = model_config.mapping.tp_size - # TODO: implement support for `tp_size > 1`. - if tp_size > 1: - raise NotImplementedError( - f"Mistral3VLM does not support `mapping.tp_size > 1` yet (got {tp_size})." - ) # Both the below are needed in order to use `_load_weights_impl`. self.model_config = model_config self.config: transformers.PixtralVisionConfig = model_config.pretrained_config @@ -204,12 +205,14 @@ def forward( ): with torch.autocast(device_type="cuda", dtype=self.config.torch_dtype): patch_embeds = self.patch_conv(pixel_values) + patch_embeds_list = [ embed[..., : (size[0] // self._patch_size), : (size[1] // self._patch_size)] for embed, size in zip(patch_embeds, image_sizes) ] - patch_embeds = torch.cat([p.flatten(1).T for p in patch_embeds_list], dim=0) + flattened_embeds = [p.flatten(1).T for p in patch_embeds_list] + patch_embeds = torch.cat(flattened_embeds, dim=0) patch_embeds = self.ln_pre(patch_embeds) position_ids = transformers.models.pixtral.modeling_pixtral.position_ids_in_meshgrid( @@ -218,10 +221,8 @@ def forward( position_embeddings = self._patch_positional_embedding(patch_embeds, position_ids) attn_metadata = self._prepare_attn_metadata( - # The `torch.cat` that creates the `patch_embeds` flattens the conv features from multiple - # images into a single sequence - hence why we hardcode the batch size to 1 here. - batch_size=1, - seq_len=position_ids.size(0), + batch_size=pixel_values.size(0), + seq_lengths=[x.size(0) for x in flattened_embeds], ) out = self.transformer( patch_embeds, @@ -235,19 +236,18 @@ def forward( def load_weights(self, weights): modeling_utils._load_weights_impl(self, weights) - def _prepare_attn_metadata(self, batch_size: int, seq_len: int): + def _prepare_attn_metadata(self, batch_size: int, seq_lengths: List[int]): request_ids = list(range(1, batch_size + 1)) - prompt_lens = [seq_len] * batch_size attn_metadata = self._metadata_cls( - seq_lens=torch.tensor([seq_len] * batch_size, dtype=torch.int), + seq_lens=torch.tensor(seq_lengths, dtype=torch.int), num_contexts=batch_size, max_num_requests=batch_size, - max_num_tokens=seq_len * batch_size, + max_num_tokens=sum(seq_lengths), kv_cache_manager=None, request_ids=request_ids, - prompt_lens=prompt_lens, + prompt_lens=seq_lengths, ) - attn_metadata.max_seq_len = seq_len * batch_size + attn_metadata.max_seq_len = max(seq_lengths) attn_metadata.prepare() return attn_metadata diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index 2d63a4bbf92b..3371bb6fc550 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -18,7 +18,8 @@ from ..attention_backend import AttentionMetadata from ..model_config import ModelConfig from .modeling_auto import AutoModelForCausalLM -from .modeling_multimodal_utils import fuse_input_embeds +from .modeling_multimodal_utils import (find_uncached_mm_embeds, + fuse_input_embeds) from .modeling_utils import register_auto_model DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1' @@ -33,9 +34,7 @@ def __init__(self, trust_remote_code: bool = True): self.model_config = model_config self.tokenizer = tokenizer - # TODO: change to True and also change the according test result - self.use_fast = False - self.device = 'cuda' + self.use_fast = True self.processor = AutoProcessor.from_pretrained( model_path, use_fast=self.use_fast, @@ -225,7 +224,7 @@ def _post_init_(self): self.model_config.num_attention_heads), theta=float(self.model_config.rope_theta), scale_type=RotaryScalingType.mrope) - self.rotary_cos_sin = torch.from_numpy(rotary_cos_sin).to(self.device) + self.rotary_cos_sin = torch.from_numpy(rotary_cos_sin) self.rotary_cos_sin = self.rotary_cos_sin.reshape( self.model_config.max_position_embeddings, int(self.model_config.hidden_size / @@ -343,7 +342,7 @@ def __call__( inputs.get("multi_modal_data", {}), inputs.get("mm_processor_kwargs", {}) processed_inputs = self._preprocess(text_prompt, mm_data, - mm_processor_kwargs).to(self.device) + mm_processor_kwargs) if not mm_data: fused_input_ids = processed_inputs['input_ids'] @@ -601,6 +600,8 @@ def forward( mrope_config = self._parse_and_concat_mrope_config( multimodal_params, num_context_requests, num_generation_requests) + mm_embeds = find_uncached_mm_embeds( + mm_embeds, multimodal_params[:num_context_requests]) if 'mrope_position_deltas' in kwargs: mrope_config['mrope_position_deltas'] = kwargs[ diff --git a/tensorrt_llm/_torch/models/modeling_qwen3.py b/tensorrt_llm/_torch/models/modeling_qwen3.py index 26353acdb04b..8635e510f423 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3.py @@ -16,8 +16,9 @@ from ..modules.linear import TensorParallelMode from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm -from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, - register_auto_model) +from ..speculative import SpecMetadata +from .modeling_speculative import SpecDecOneEngineForCausalLM +from .modeling_utils import DecoderModel, register_auto_model class Qwen3Attention(Attention): @@ -148,6 +149,7 @@ def forward( attn_metadata: AttentionMetadata, residual: Optional[torch.Tensor], mrope_config: Optional[Tuple[torch.Tensor, int]] = None, + spec_metadata: Optional[SpecMetadata] = None, **kwargs, ) -> torch.Tensor: if residual is None: @@ -171,6 +173,10 @@ def forward( hidden_states, residual) hidden_states = self.mlp(hidden_states) + if spec_metadata is not None: + spec_metadata.maybe_capture_hidden_states(self.layer_idx, + hidden_states, residual) + return hidden_states, residual @@ -207,6 +213,7 @@ def forward( position_ids: Optional[torch.IntTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, mrope_config: Optional[Tuple[torch.Tensor, int]] = None, + spec_metadata: Optional[SpecMetadata] = None, **kwargs, ) -> torch.Tensor: if (input_ids is None) ^ (inputs_embeds is not None): @@ -227,6 +234,7 @@ def forward( attn_metadata=attn_metadata, residual=residual, mrope_config=mrope_config, + spec_metadata=spec_metadata, ) hidden_states, _ = self.norm(hidden_states, residual) @@ -234,7 +242,7 @@ def forward( @register_auto_model("Qwen3ForCausalLM") -class Qwen3ForCausalLM(DecoderModelForCausalLM[Qwen3Model, Qwen3Config]): +class Qwen3ForCausalLM(SpecDecOneEngineForCausalLM[Qwen3Model, Qwen3Config]): def __init__( self, @@ -242,33 +250,5 @@ def __init__( ): super().__init__( Qwen3Model(model_config), - config=model_config, - hidden_size=model_config.pretrained_config.hidden_size, - vocab_size=model_config.pretrained_config.vocab_size, - ) - - # NOTE: Qwen2-VL needs special mrope_config so adding separate forward() function to accept 'mrope_config'. - def forward( - self, - attn_metadata: AttentionMetadata, - input_ids: torch.IntTensor = None, - position_ids: Optional[torch.IntTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - return_context_logits: bool = False, - mrope_config: Optional[dict] = None, - **kwargs, - ) -> torch.Tensor: - output = self.model( - input_ids=input_ids, - attn_metadata=attn_metadata, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - mrope_config=mrope_config, - ) - - return self.logits_processor.forward( - output, - self.lm_head, - attn_metadata, - return_context_logits, + model_config, ) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py index d01bce4ded67..2d447dd527b4 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py @@ -3,29 +3,28 @@ import torch from torch import nn -from tqdm import tqdm from transformers import Qwen3MoeConfig +from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ + BaseWeightMapper + from ..attention_backend import AttentionMetadata from ..distributed import (AllReduce, AllReduceFusionOp, AllReduceParams, MoEAllReduce, MoEAllReduceParams, allgather) from ..model_config import ModelConfig from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import (BaseMoeRoutingMethod, CutlassFusedMoE, MoE, +from ..modules.fused_moe import (BaseMoeRoutingMethod, RenormalizeMoeRoutingMethod, RenormalizeNaiveMoeRoutingMethod, RoutingMethodType, TRTLLMGenFusedMoE, - WideEPMoE, create_moe) + create_moe) from ..modules.linear import TensorParallelMode from ..modules.rms_norm import RMSNorm from ..speculative import SpecMetadata -from ..utils import disable_fp4_allgather from .modeling_qwen3 import Qwen3Attention from .modeling_speculative import SpecDecOneEngineForCausalLM -from .modeling_utils import (DecoderModel, EagerFusionConfig, - duplicate_kv_weight, filter_weights, - register_auto_model) +from .modeling_utils import DecoderModel, EagerFusionConfig, register_auto_model class Qwen3Gate(nn.Module): @@ -133,22 +132,11 @@ def forward( assert not self.enable_attention_dp if self.enable_attention_dp and self.mapping.tp_size > 1: - # FP4 all_gather moves this bf16 allgather in to after topk and fp4 quantization - # to reduce allreduce BW - if (disable_fp4_allgather() - and not self.experts.enable_alltoall) or isinstance( - self.experts, TRTLLMGenFusedMoE): + if isinstance(self.experts, TRTLLMGenFusedMoE): hidden_states = allgather(hidden_states, self.mapping, dim=0, sizes=all_rank_num_tokens) - elif not isinstance(self.experts, (CutlassFusedMoE, WideEPMoE)) or ( - not self.experts.has_fp8_qdq and self.experts.has_nvfp4): - # Use padding when not using the cutlass path or when x_sf in self.experts is not None - use_dp_padding = True - hidden_states = torch.nn.functional.pad( - hidden_states, - (0, 0, 0, all_rank_max_num_tokens - hidden_states.shape[0])) router_logits = self.gate(hidden_states) final_hidden_states = self.experts( @@ -321,6 +309,13 @@ def __init__(self, model_config: ModelConfig[Qwen3MoeConfig]): super().__init__(model_config) config = self.model_config self.aux_stream = torch.cuda.Stream() + self.preload_weight_modules = [] + if config.moe_backend == "TRTLLM": + self.preload_weight_modules = [ + "experts", + "routing_method", + "all_reduce", + ] if model_config.mapping.enable_attention_dp: # When attention_dp is enabled, we cannot do all_reduce since @@ -393,68 +388,11 @@ def __init__( Qwen3MoEModel(model_config), model_config, ) + self.preload_weight_modules = self.model.preload_weight_modules + + def load_weights(self, weights: dict, weight_mapper: BaseWeightMapper): + super().load_weights(weights, weight_mapper) - def load_weights(self, weights: Dict): - tp_size = self.model_config.mapping.tp_size - enable_attention_dp = self.model_config.mapping.enable_attention_dp - - num_kv_heads = self.config.num_key_value_heads - - params_map = { - "qkv_proj": ["q_proj", "k_proj", "v_proj"], - "gate_up_proj": ["gate_proj", "up_proj"] - } - for name, module in tqdm(list(self.named_modules()), - desc="Loading weights"): - if len(module._parameters) > 0: - # skip load weights if tie word embeddings is enabled and layer is lm_head - if self.config.tie_word_embeddings and name.startswith( - "lm_head") or name.startswith("draft_model"): - continue - - names = name.split(".") - if names[-1] in params_map: - module_weights = [] - for new_name in params_map[names[-1]]: - fw = filter_weights(".".join(names[:-1] + [new_name]), - weights) - tensors_need_duplication = ["weight", "bias"] - if module.quant_config.quant_mode.has_nvfp4(): - tensors_need_duplication.append("weight_scale") - if module.quant_config.quant_mode.has_fp8_block_scales( - ): - tensors_need_duplication.append("weight_scale_inv") - if new_name in ["k_proj", "v_proj"]: - fw = { - k: (duplicate_kv_weight( - weight=v[:], - num_kv_heads=num_kv_heads, - tensor_parallel_size=tp_size - if not enable_attention_dp else 1) - if k in tensors_need_duplication else v) - for k, v in fw.items() - } - module_weights.append(fw) - module.load_weights(weights=module_weights) - else: - module_weights = filter_weights(name, weights) - if isinstance(module, MoE): - updated_module_weights = {} - for weight_name, weight_value in module_weights.items(): - new_weight_name = (weight_name.replace( - "gate_proj", - "w1").replace("up_proj", - "w3").replace("down_proj", "w2")) - updated_module_weights[ - new_weight_name] = weight_value - del module_weights - module.load_weights(weights=[updated_module_weights]) - elif hasattr(module, "load_weights"): - module.load_weights(weights=[module_weights]) - else: - for n, p in module._parameters.items(): - if p is not None: - p.data.copy_(module_weights[n][:]) for idx, layer in enumerate( self.model.layers[:self.config.num_hidden_layers]): if idx == self.config.num_hidden_layers - 1: diff --git a/tensorrt_llm/_torch/models/modeling_qwen_moe.py b/tensorrt_llm/_torch/models/modeling_qwen_moe.py index 0ee6d1ac55d4..2bbf9b80d546 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen_moe.py @@ -1,9 +1,8 @@ -from typing import Dict, Optional +from typing import Optional import torch import torch.nn.functional as F from torch import nn -from tqdm import tqdm from transformers import Qwen2MoeConfig from tensorrt_llm.functional import PositionEmbeddingType @@ -14,12 +13,11 @@ from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import DefaultMoeRoutingMethod, MoE, create_moe +from ..modules.fused_moe import DefaultMoeRoutingMethod, create_moe from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear, TensorParallelMode from ..modules.rms_norm import RMSNorm from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, - duplicate_kv_weight, filter_weights, register_auto_model) @@ -255,57 +253,3 @@ def __init__( config=model_config, hidden_size=model_config.pretrained_config.hidden_size, vocab_size=model_config.pretrained_config.vocab_size) - - def load_weights(self, weights: Dict): - tp_size = self.model_config.mapping.tp_size - num_kv_heads = self.config.num_key_value_heads - - params_map = { - 'qkv_proj': ['q_proj', 'k_proj', 'v_proj'], - 'gate_up_proj': ['gate_proj', 'up_proj'] - } - for name, module in tqdm(list(self.named_modules()), - desc="Loading weights"): - if len(module._parameters) > 0: - # skip load weights if tie word embeddings is enabled and layer is lm_head - if self.config.tie_word_embeddings and name.startswith( - "lm_head"): - continue - - names = name.split('.') - if names[-1] in params_map: - module_weights = [] - for new_name in params_map[names[-1]]: - fw = filter_weights('.'.join(names[:-1] + [new_name]), - weights) - if new_name in ['k_proj', 'v_proj']: - fw = { - k: - duplicate_kv_weight( - weight=v[:], - num_kv_heads=num_kv_heads, - tensor_parallel_size=tp_size) - if k in ["weight", "bias"] else v - for k, v in fw.items() - } - module_weights.append(fw) - module.load_weights(weights=module_weights) - else: - module_weights = filter_weights(name, weights) - if isinstance(module, MoE): - updated_module_weights = {} - for weight_name, weight_value in module_weights.items(): - new_weight_name = weight_name.replace( - "gate_proj", - "w1").replace("up_proj", - "w3").replace("down_proj", "w2") - updated_module_weights[ - new_weight_name] = weight_value - del module_weights - module.load_weights(weights=[updated_module_weights]) - elif hasattr(module, 'load_weights'): - module.load_weights(weights=[module_weights]) - else: - for n, p in module._parameters.items(): - if p is not None: - p.data.copy_(module_weights[n][:]) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index c6f05219775a..ee178b48f145 100644 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -4,6 +4,8 @@ from torch import nn from transformers import LlamaConfig +from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ + BaseWeightMapper from tensorrt_llm.functional import PositionEmbeddingType from ..attention_backend import AttentionMetadata @@ -264,7 +266,7 @@ def forward( return_context_logits, ) - def load_weights(self, weights: Dict): + def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper): new_weights = {} for k, v in weights.items(): if 'lm_head' not in k: @@ -274,9 +276,12 @@ def load_weights(self, weights: Dict): new_k = k new_weights[new_k] = v if self.load_lm_head_from_target: - super().load_weights(new_weights, skip_modules=['lm_head']) + super().load_weights(weights=new_weights, + weight_mapper=weight_mapper, + skip_modules=['lm_head']) else: - super().load_weights(new_weights) + super().load_weights(weights=new_weights, + weight_mapper=weight_mapper) def load_weights_from_target_model(self, target_model: torch.nn.Module) -> None: @@ -401,9 +406,16 @@ def forward( return logits - def load_weights(self, weights: Dict): - super().load_weights(weights, skip_modules=["draft_model"]) - - def load_draft_weights(self, weights: Dict): - self.draft_model.load_weights(weights) + def load_weights(self, + weights: Dict, + weight_mapper: Optional[BaseWeightMapper] = None): + super().load_weights(weights=weights, + weight_mapper=weight_mapper, + skip_modules=["draft_model"]) + + def load_draft_weights(self, + weights: Dict, + weight_mapper: Optional[BaseWeightMapper] = None): + self.draft_model.load_weights(weights=weights, + weight_mapper=weight_mapper) self.draft_model.load_weights_from_target_model(self) diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 1dac009f5c1c..020762d8927b 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -364,11 +364,13 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], if (hasattr(config, 'lora_config') and config.lora_config is not None and len(config.lora_config.lora_dir) == 1): - lora_loader = HfLoraLoader(config.lora_config.lora_dir) - if lora_loader.lm_head is not None and lora_loader.vocab_size != 0: - weight = lora_loader.lm_head - self.has_custom_lm_head = True - vocab_size = lora_loader.vocab_size + # Only check for custom lm_head in HF LoRA, not NeMo + if config.lora_config.lora_ckpt_source == "hf": + lora_loader = HfLoraLoader(config.lora_config.lora_dir) + if lora_loader.lm_head is not None and lora_loader.vocab_size != 0: + weight = lora_loader.lm_head + self.has_custom_lm_head = True + vocab_size = lora_loader.vocab_size self.lm_head = LMHead( vocab_size, @@ -524,12 +526,25 @@ def forward( return_context_logits, ) - def load_weights(self, weights: Dict, skip_modules: List[str] = []): + def load_weights(self, + weights: Dict, + weight_mapper: Optional["BaseWeightMapper"] = None, + skip_modules: List[str] = []): + # TODO smor- this solution is a temporary solution to load weights while we are still using + # the old checkpoint format loading process. Once checkpoint format is unified + # this method will be removed. preload_weight_modules = getattr(self, "preload_weight_modules", None) - _load_weights_impl(self, - weights, - skip_modules, - preload_weight_modules=preload_weight_modules) + if weight_mapper is None: + _load_weights_impl(self, + weights, + skip_modules, + preload_weight_modules=preload_weight_modules) + else: + _load_weights_impl_v2(self, + weights, + weight_mapper, + skip_modules, + preload_weight_modules=preload_weight_modules) def infer_max_seq_len(self) -> int: # Modified from tensorrt_llm/builder.py _init_max_seq_len @@ -558,6 +573,10 @@ def infer_max_seq_len(self) -> int: MODEL_CLASS_MAPPING = {} +MODEL_CLASS_MAPPER_MAPPING = {} +MODEL_CLASS_CHECKPOINT_WEIGHT_LOADER_DEFAULT_MAPPING = {} +MODEL_CLASS_CONFIG_LOADER_DEFAULT_MAPPING = {} +CHECKPOINT_LOADER_FORMAT_DEFAULT_MAPPING = {} def register_auto_model(name: str): @@ -569,6 +588,59 @@ def decorator(cls): return decorator +def register_mapper(format: str, name: Optional[str] = None): + + def decorator(cls): + if name is not None: + # set cls for model name and format pair + MODEL_CLASS_MAPPER_MAPPING[f'{name}_{format}'] = cls + else: + # resort to the default per format + MODEL_CLASS_MAPPER_MAPPING[format] = cls + return cls + + return decorator + + +def register_checkpoint_weight_loader(name: str): + + def decorator(cls): + MODEL_CLASS_CHECKPOINT_WEIGHT_LOADER_DEFAULT_MAPPING[name] = cls + return cls + + return decorator + + +def register_checkpoint_loader(name: str): + + def decorator(cls): + CHECKPOINT_LOADER_FORMAT_DEFAULT_MAPPING[name] = cls + return cls + + return decorator + + +def register_config_loader(name: str): + + def decorator(cls): + MODEL_CLASS_CONFIG_LOADER_DEFAULT_MAPPING[name] = cls + return cls + + return decorator + + +def get_checkpoint_weight_loader(name: str) -> Type["BaseWeightLoader"]: + if name not in MODEL_CLASS_CHECKPOINT_WEIGHT_LOADER_DEFAULT_MAPPING: + raise ValueError(f"Default checkpoint weight loader {name} not found.") + return MODEL_CLASS_CHECKPOINT_WEIGHT_LOADER_DEFAULT_MAPPING[name] + + +def get_config_loader(name: str) -> Type["BaseConfigLoader"]: + if name not in MODEL_CLASS_CONFIG_LOADER_DEFAULT_MAPPING: + raise ValueError(f"Default config loader {name} not found.") + return MODEL_CLASS_CONFIG_LOADER_DEFAULT_MAPPING[name] + + def get_model_architecture( model_config: TConfig) -> Tuple[Type[nn.Module], str]: cls = None @@ -587,7 +659,6 @@ def get_model_architecture( def rename_weights_with_regex(pattern_mapping: Dict[str, str], weights: Dict): """ Rename weight keys according to regex pattern matching. - Args: pattern_mapping: A dictionary mapping regex patterns to replacement strings. The key is HF name pattern, and the value is corresponding TRT-LLM name pattern. The patterns will be used to match keys in the weights dict and replace @@ -600,7 +671,6 @@ def rename_weights_with_regex(pattern_mapping: Dict[str, str], weights: Dict): r'(.*?)out_proj(.*)': r'\1o_proj\2' } weights: A dictionary of weights - Returns: A dictionary of weights with renamed keys """ @@ -683,6 +753,9 @@ def _load_weights_impl(model: Union[nn.Module, DecoderModelForCausalLM], preload_weight_modules: Optional[List[str]] = None): # TODO: remove preload_weight_modules - it is a workaround for min-latency llama4 model loading where # we need some order in the module loading. Once this is resolved, we can remove this workaround. + # TODO smor- this method is here as a temporary solution to load weights. + # Once checkpoint format is unified, this method will be removed. + if not hasattr(model, 'model_config') or not isinstance( model.model_config, ModelConfig): raise ValueError("model must have a model_config attribute") @@ -784,3 +857,71 @@ def load_single_module(name, module): args_list = [(name, module) for name, module in model.named_modules() if name not in serial_load_modules] run_concurrently(load_single_module, args_list, pbar=pbar) + + +def _load_weights_impl_v2(model: Union[nn.Module, DecoderModelForCausalLM], + weights: Dict, + weight_mapper: "BaseWeightMapper", + skip_modules: List[str] = [], + params_map: Optional[Dict[str, str]] = None, + preload_weight_modules: Optional[List[str]] = None): + # TODO: remove preload_weight_modules - it is a workaround for min-latency llama4 and Qwen3 model loading where + # we need some order in the module loading. Once this is resolved, we can remove this workaround. + weight_mapper.add_skip_modules(skip_modules) + if params_map is not None: + weights = weight_mapper.rename_by_params_map(params_map, weights) + logger.info(f"Renamed weights with params_map: {params_map}") + + def load_single_module(name, module): + if len(module._parameters) > 0: + if weight_mapper.should_skip_module(name): + return + + names = name.split('.') + module_names_breakdown, module_name = names[:-1], names[-1] + + if weight_mapper.does_require_special_handling(module_name): + module_weights = weight_mapper.apply_callbacks( + module, module_name, module_names_breakdown, weights) + module.load_weights(weights=module_weights) + else: + module_weights = weight_mapper.filter_weights(name, weights) + if weight_mapper.is_special_instance_module(module): + weight_mapper.handle_special_instance_module( + module, module_name, module_weights) + + elif hasattr(module, 'load_weights'): + module.load_weights(weights=[module_weights]) + else: + for n, p in module._parameters.items(): + if p is not None: + weight_mapper.handle_manual_copy( + module_name, module_weights, n, p) + + if os.environ.get("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", + False) in ["True", "true", "1", "yes", "y"]: + for name, module in tqdm(list(model.named_modules()), + desc="Loading weights"): + load_single_module(name, module) + else: + all_modules = dict(model.named_modules()) + serial_load_modules = [] + if preload_weight_modules is not None: + for module in preload_weight_modules: + serial_load_modules.extend([ + name for name in all_modules.keys() if name.endswith(module) + ]) + logger.info(f"Serial load modules: {serial_load_modules}") + pbar = tqdm(serial_load_modules, desc="Loading weights serially") + for module in serial_load_modules: + # logger.info(f"Loading weights for {module} in serial") + load_single_module(module, all_modules[module]) + pbar.update(1) + del all_modules[module] + pbar.close() + + pbar = tqdm(list(model.named_modules()), + desc="Loading weights concurrently") + args_list = [(name, module) for name, module in model.named_modules() + if name not in serial_load_modules] + run_concurrently(load_single_module, args_list, pbar=pbar) diff --git a/tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py b/tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py index 62146d9295fc..385a5ec4b911 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py +++ b/tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py @@ -59,7 +59,7 @@ def reserve(self, hidden_size: int, hidden_dtype: torch.dtype): def dispatch(self, x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], topk_idx: torch.Tensor, topk_weights: torch.Tensor, - num_experts: int) -> \ + num_experts: int, global_expert_id_offset: int) -> \ Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], torch.Tensor, torch.Tensor, List, Tuple]: # NOTES: an optional `previous_event` means a CUDA event captured that you want to make it as a dependency # of the dispatch kernel, it may be useful with communication-computation overlap. For more information, please @@ -76,7 +76,8 @@ def dispatch(self, x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event = \ self.buffer.dispatch(x, topk_idx=topk_idx, topk_weights=topk_weights, num_tokens_per_rank=num_tokens_per_rank, num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, - is_token_in_rank=is_token_in_rank, num_tokens_per_expert=num_tokens_per_expert) + is_token_in_rank=is_token_in_rank, num_tokens_per_expert=num_tokens_per_expert, + global_expert_id_offset=global_expert_id_offset) assert event.event is None # For event management, please refer to the docs of the `EventOverlap` class @@ -99,7 +100,7 @@ class VariableLengthLowLatencyBuffer: def __init__(self, mapping: Mapping): self.comm = mpi_comm().Split(mapping.pp_rank, mapping.moe_ep_rank) self.buffer = None - self.num_max_dispatch_tokens_per_rank = None + self.num_experts = None def __del__(self): self.comm.Free() @@ -119,6 +120,7 @@ def reserve(self, num_max_dispatch_tokens_per_rank: int, hidden_size: int, allow_nvlink_for_low_latency_mode = (os.environ.get( "TRTLLM_DEEP_EP_DISABLE_P2P_FOR_LOW_LATENCY_MODE", "0") == "0") + assert self.num_experts is None or self.num_experts == num_experts # Allocate a buffer if not existed or not enough buffer size if self.buffer is None or self.buffer.num_rdma_bytes < num_rdma_bytes: # NOTES: for best performance, the QP number **must** be equal to the number of the local experts @@ -132,17 +134,13 @@ def reserve(self, num_max_dispatch_tokens_per_rank: int, hidden_size: int, allow_nvlink_for_low_latency_mode= allow_nvlink_for_low_latency_mode, comm=self.comm) + self.num_experts = num_experts def low_latency_dispatch(self, hidden_states: torch.Tensor, topk_idx: torch.Tensor, num_max_dispatch_tokens_per_rank: int, num_experts: int): - if self.num_max_dispatch_tokens_per_rank is None: - self.num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank - if num_max_dispatch_tokens_per_rank != self.num_max_dispatch_tokens_per_rank: - raise NotImplementedError( - "There are issues if `low_latency_dispatch` calls use different `num_max_dispatch_tokens_per_rank` values" - ) + assert num_experts == self.num_experts # Do MoE dispatch, compatible with CUDA graph (but you may restore some buffer status once you replay) recv_hidden_states, recv_expert_count, handle, event, hook = \ diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py index c0a2c4fbecbf..025b112034da 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -4,8 +4,7 @@ from ...distributed import allgather, reducescatter from ...model_config import ModelConfig -from ...utils import (EventType, Fp4QuantizedTensor, ceil_div, - disable_fp4_allgather, swizzle_sf) +from ...utils import EventType, Fp4QuantizedTensor, ceil_div, swizzle_sf from .interface import MoE from .quantization import (DeepSeekFP8BlockScalesFusedMoEMethod, FP8QDQFusedMoEMethod, MoEWeightLoadingMode, @@ -220,9 +219,7 @@ def forward_chunk( # TODO: remove this once we have correct fusedmoe kernel ready token_final_scales = None - use_allgather = self.use_dp and self.parallel_size > 1 and not disable_fp4_allgather( - ) - + run_post_quant_allgather = self.use_dp and self.parallel_size > 1 # quantize inputs use_deepseek_fp8_block_scale = False use_w4a8_group_scaling = False @@ -238,7 +235,7 @@ def forward_chunk( use_w4a8_group_scaling = True weight_dtype = torch.quint4x2 elif self.has_nvfp4: - if use_allgather: + if run_post_quant_allgather: if isinstance(x, Fp4QuantizedTensor): assert not x.is_sf_swizzled, "Fp4QuantizedTensor should not be swizzled before communication" x_row = x.shape[0] @@ -249,28 +246,26 @@ def forward_chunk( x_row = x.shape[0] x_col = x.shape[1] x, x_sf = torch.ops.trtllm.fp4_quantize( - x, - self.fc31_input_scale, - self.scaling_vector_size, - sfUseUE8M0=False, - swizzedLayout=False) - x_sf = x_sf.view( - x_row, ceil_div(x_col, self.scaling_vector_size)) + x, self.fc31_input_scale, self.scaling_vector_size, + False, False) else: if not isinstance(x, Fp4QuantizedTensor): x, x_sf = torch.ops.trtllm.fp4_quantize( - x, - self.fc31_input_scale, - self.scaling_vector_size, - sfUseUE8M0=False, - swizzedLayout=True) + x, self.fc31_input_scale, self.scaling_vector_size, + False, True) else: raise ValueError( f"unsupported quantization mode: {self.quant_config.quant_mode}" ) # gather inputs for attention dp - if use_allgather: + if run_post_quant_allgather: + if x_sf is not None: + x_sf = x_sf.view(x_row, ceil_div(x_col, + self.scaling_vector_size)) + assert len( + x_sf.shape + ) == 2, "The hidden states scaling factor should be 2D tensor before allgather" x, x_sf, token_selected_experts, token_final_scales = allgather( [x, x_sf, token_selected_experts, token_final_scales], self.mapping, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index b5f93ab2500c..94e082a6670c 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -4,7 +4,7 @@ from ...distributed.ops import reducescatter from ...model_config import ModelConfig -from ...utils import Fp4QuantizedTensor, next_positive_power_of_2 +from ...utils import Fp4QuantizedTensor from .interface import MoE, MoEWeightLoadingMode from .quantization import (DeepSeekFP8BlockScalesFusedMoEMethod, NVFP4TRTLLMGenFusedMoEMethod) @@ -91,19 +91,6 @@ def __init__( def _check_configs(self): assert self.has_deepseek_fp8_block_scales or self.has_nvfp4, "TRTLLMGenFusedMoE only supports fp8_block_scaling and nvfp4 dtypes." - def _get_tile_tokens_dim(self, x: torch.Tensor): - top_k = self.routing_method.top_k - # Number of tokens in the input tensor. - num_tokens = x.shape[0] - # Guess tokens per expert assuming perfect expert distribution first. - num_tokens_per_expert = (num_tokens * top_k) // self.num_experts - # And pad the number to the next power of 2. - tile_tokens_dim = next_positive_power_of_2(num_tokens_per_expert) - # Cap to 8-64 tokens per CTA tile as it's the range supported by the kernel. - tile_tokens_dim = min(max(tile_tokens_dim, 8), 64) - - return tile_tokens_dim - def _get_quant_method(self): if self.quant_config is not None: if self.quant_config.layer_quant_mode.has_fp8_block_scales(): @@ -204,7 +191,6 @@ def forward( slot_start, # local_expert_start; use ep_rank if stride!=1 self.expert_size_per_partition, # local_expert_size routed_scaling_factor, - self._get_tile_tokens_dim(x), self.routing_method.routing_method_type, ) elif self.has_nvfp4: @@ -240,7 +226,6 @@ def forward( slot_start, # local_expert_start; use ep_rank if stride!=1 self.expert_size_per_partition, # local_expert_size routed_scaling_factor, - self._get_tile_tokens_dim(x), self.routing_method.routing_method_type, do_finalize=do_finalize, ) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py index 9290aae30290..81778c28544d 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py @@ -192,8 +192,12 @@ def __init__( model_config.mapping) self.deep_ep_buffer.reserve(hidden_size, dtype) elif self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: - self.deep_ep_max_num_tokens = min(model_config.max_num_tokens, - self.moe_max_num_tokens) + self.deep_ep_max_num_tokens = int( + os.environ.get( + "TRTLLM_DEEP_EP_TOKEN_LIMIT", + str( + min(model_config.max_num_tokens, + self.moe_max_num_tokens)))) self.deep_ep_buffer = buffer_pool.get_low_latency_buffer( model_config.mapping) self.deep_ep_buffer.reserve(self.deep_ep_max_num_tokens, @@ -274,6 +278,23 @@ def enable_alltoall(self): """ return self.alltoall_method_type != AlltoallMethodType.NotEnabled + def calculate_num_chunks(self, all_rank_num_tokens: List[int]) -> int: + num_rows = sum(all_rank_num_tokens) + return (num_rows + self.moe_max_num_tokens - + 1) // self.moe_max_num_tokens + + def can_use_alltoall(self, all_rank_num_tokens, all_rank_max_num_tokens): + # Disable alltoall when chunking is used + if self.calculate_num_chunks(all_rank_num_tokens) > 1: + return False + + # For DeepEPLowLatency, check if tokens exceed the threshold + if (self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency + and all_rank_max_num_tokens > self.deep_ep_max_num_tokens): + return False + + return self.enable_alltoall + def _get_quant_method(self): if self.quant_config is not None and self.quant_config.layer_quant_mode.has_any_quant( exclude_kv_cache=True): @@ -316,11 +337,12 @@ def dummy_allreduce(self): def reducescatter_or_allreduce( self, inputs, + use_all_to_all: bool, all_rank_num_tokens: Optional[List[int]] = None, use_dp_padding: Optional[bool] = None, ): outputs = inputs - if not self.enable_alltoall: + if not use_all_to_all: if self.enable_dummy_allreduce: self.dummy_allreduce() outputs = reducescatter( @@ -334,6 +356,7 @@ def forward_chunk( self, x: Union[torch.Tensor, Fp4QuantizedTensor], router_logits: torch.Tensor, + use_all_to_all: bool, output_dtype: Optional[torch.dtype] = None, all_rank_num_tokens: Optional[List[int]] = None, all_rank_max_num_tokens: Optional[int] = None, @@ -382,7 +405,7 @@ def forward_chunk( ) and is_first_call: self.layer_load_balancer.maybe_cudagraph_done_wait() - use_allgather = not self.enable_alltoall + use_allgather = not use_all_to_all loadbalancer_local_statistic_info = None gathered_loadbalancer_local_statistic_info = None @@ -391,7 +414,7 @@ def forward_chunk( token_selected_slots = token_selected_experts else: if not self.layer_load_balancer.is_static_routing( - ) and self.enable_alltoall: + ) and use_all_to_all: self.layer_load_balancer.local_statistic( token_selected_experts, is_first_stage=is_first_call, @@ -400,7 +423,7 @@ def forward_chunk( token_selected_experts, self.use_dp) if not self.layer_load_balancer.is_static_routing(): # split into two part to get possible overlap with load balancer routing - if self.enable_alltoall: + if use_all_to_all: if is_last_call: loadbalancer_local_statistic_info = self.layer_load_balancer.get_local_statistic_tensor( ) @@ -412,7 +435,9 @@ def forward_chunk( ExpertStatistic.set_layer(self.layer_idx) ExpertStatistic.maybe_add_info(self.num_slots, token_selected_slots) - if self.enable_alltoall: + # If alltoall is disabled, we need also disable use_postquant_alltoall + use_postquant_alltoall = self.use_postquant_alltoall and use_all_to_all + if use_all_to_all: if self.alltoall_method_type == AlltoallMethodType.MNNVL: if self.enable_dummy_allreduce: self.dummy_allreduce() @@ -423,24 +448,27 @@ def forward_chunk( x, token_selected_slots, token_final_scales, + use_postquant_alltoall, loadbalancer_local_statistic_info) elif self.alltoall_method_type == AlltoallMethodType.DeepEP: - if not self.use_postquant_alltoall: + if not use_postquant_alltoall: x, recv_topk_idx, token_final_scales, num_recv_tokens_per_expert_list, deep_ep_handle = \ - self.deep_ep_buffer.dispatch(x, token_selected_slots.to(torch.int64), token_final_scales, self.num_slots) + self.deep_ep_buffer.dispatch(x, token_selected_slots, token_final_scales, self.num_slots, + self.expert_size_per_partition * self.mapping.moe_ep_rank) + padded, x, _, token_selected_slots, token_final_scales = self.pad_empty_recv_tensors( + x, None, recv_topk_idx, token_final_scales) elif self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: - if not self.use_postquant_alltoall: - deep_ep_topk_idx = token_selected_slots.to(torch.int64) + if not use_postquant_alltoall: + deep_ep_topk_idx = token_selected_slots deep_ep_topk_weights = token_final_scales + assert all_rank_max_num_tokens <= self.deep_ep_max_num_tokens x, recv_expert_count, deep_ep_handle = \ - self.deep_ep_buffer.low_latency_dispatch(x, deep_ep_topk_idx, self.deep_ep_max_num_tokens, self.num_slots) - # x shape: [#local experts, EP size * deep_ep_max_num_tokens, hidden_size] + self.deep_ep_buffer.low_latency_dispatch(x, deep_ep_topk_idx, all_rank_max_num_tokens, self.num_slots) + # x shape: [#local experts, EP size * all_rank_max_num_tokens, hidden_size] # recv_expert_count shape: [#local experts] # Adapter between `torch.ops.trtllm.fused_moe` and DeepEP # TODO: remove the adapter by changing `torch.ops.trtllm.fused_moe` API - x = x[:, :self.mapping.moe_ep_size * - all_rank_max_num_tokens] mask = torch.arange( x.shape[1], dtype=torch.int32, device=x.device).expand( x.shape[0], @@ -460,8 +488,6 @@ def forward_chunk( token_selected_slots, dtype=token_final_scales.dtype) x_sf = None - x_is_sf_swizzled = x.is_sf_swizzled if isinstance( - x, Fp4QuantizedTensor) else False x_row = x.shape[0] x_col = x.shape[1] if self.has_any_quant: @@ -469,7 +495,7 @@ def forward_chunk( x, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( x, self.fc31_input_dequant) elif self.has_nvfp4: - if use_allgather or self.use_postquant_alltoall: + if use_allgather or use_postquant_alltoall: if isinstance(x, Fp4QuantizedTensor): if use_allgather: assert not x.is_sf_swizzled, "Fp4QuantizedTensor should not be swizzled before allgather" @@ -479,7 +505,6 @@ def forward_chunk( x_col = x.shape[1] * 2 else: # for both postquant alltoall and allgather, we need non swizzle layout - needed_sf_swizzle = False x_row = x.shape[0] x_col = x.shape[1] x, x_sf = torch.ops.trtllm.fp4_quantize( @@ -487,10 +512,8 @@ def forward_chunk( self.fc31_input_scale, self.scaling_vector_size, sfUseUE8M0=False, - swizzedLayout=needed_sf_swizzle) - if self.use_postquant_alltoall: - x_sf = x_sf.view((x_row, -1)) - x_is_sf_swizzled = needed_sf_swizzle + swizzedLayout=False) + x_sf = x_sf.view((x_row, -1)) elif self.has_deepseek_fp8_block_scales: use_deepseek_fp8_block_scale = True @@ -520,12 +543,11 @@ def forward_chunk( x_row = x.shape[0] # Fp4 gemm has extra scaling factor if x_sf is not None: - assert not x_is_sf_swizzled, "Fp4QuantizedTensor should not be swizzled before allgather" x_sf = swizzle_sf(x_sf, x_row, x_col, self.scaling_vector_size) if self.layer_load_balancer and not self.layer_load_balancer.is_static_routing( ): - if self.enable_alltoall: + if use_all_to_all: if is_last_call: gathered_loadbalancer_local_statistic_info = gathered_loadbalancer_local_statistic_info.view( (self.mapping.moe_ep_size, self.num_experts)) @@ -545,9 +567,7 @@ def forward_chunk( cluster_rank = self.cluster_rank quant_scales = self.quant_scales - if self.use_postquant_alltoall: - if x_sf is not None and self.has_nvfp4: - assert not x_is_sf_swizzled, "Fp4 scaling factor should not be swizzled before Alltoall" + if use_postquant_alltoall: if self.alltoall_method_type == AlltoallMethodType.MNNVL: x, x_sf = self.alltoall_postquant_dispatch( x, x_sf, alltoall_info) @@ -558,14 +578,17 @@ def forward_chunk( x_sf_dtype = x_sf.dtype x_sf = x_sf.view(torch.float32) (x, x_sf), recv_topk_idx, token_final_scales, num_recv_tokens_per_expert_list, deep_ep_handle = \ - self.deep_ep_buffer.dispatch((x, x_sf), token_selected_slots.to(torch.int64), token_final_scales, self.num_slots) + self.deep_ep_buffer.dispatch((x, x_sf), token_selected_slots, token_final_scales, self.num_slots, + self.expert_size_per_partition * self.mapping.moe_ep_rank) + padded, x, x_sf, token_selected_slots, token_final_scales = self.pad_empty_recv_tensors( + x, x_sf, recv_topk_idx, token_final_scales) if x_sf is not None: x_sf = x_sf.view(x_sf_dtype) if self.has_nvfp4: x_sf = swizzle_sf(x_sf, x.shape[0], x.shape[1] * 2, self.scaling_vector_size) elif self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: - assert x_sf is not None and self.has_nvfp4 and not x_is_sf_swizzled + assert x_sf is not None and self.has_nvfp4 token_num = x_row hidden_size = x_col assert hidden_size % 32 == 0 @@ -587,28 +610,16 @@ def forward_chunk( fp4_packed_tensor[:, x.shape[1]:x.shape[1] + x_sf.shape[1]] = x_sf - deep_ep_topk_idx = token_selected_slots.to(torch.int64) + deep_ep_topk_idx = token_selected_slots deep_ep_topk_weights = token_final_scales - # Each LL combine/dispatch kernel call requires that the `dispatch_rdma_recv_count_buffer` be properly cleaned. - # However, the offset of this buffer within the entire RDMA buffer changes according to the hidden size. - # Therefore, if the hidden size for the next LL dispatch/combine call is different from the current kernel call, manual cleaning is necessary. - if packed_hidden_size != hidden_size: - self.deep_ep_buffer.clean_low_latency_buffer( - self.deep_ep_max_num_tokens, packed_hidden_size, - self.num_slots) + + assert all_rank_max_num_tokens <= self.deep_ep_max_num_tokens fp4_packed_tensor, recv_expert_count, deep_ep_handle = \ - self.deep_ep_buffer.low_latency_dispatch(fp4_packed_tensor, deep_ep_topk_idx, self.deep_ep_max_num_tokens, self.num_slots) - if packed_hidden_size != hidden_size: - self.deep_ep_buffer.clean_low_latency_buffer( - self.deep_ep_max_num_tokens, hidden_size, - self.num_slots) + self.deep_ep_buffer.low_latency_dispatch(fp4_packed_tensor, deep_ep_topk_idx, all_rank_max_num_tokens, self.num_slots) deep_ep_handle = list(deep_ep_handle) deep_ep_handle[3] = hidden_size deep_ep_handle = tuple(deep_ep_handle) - fp4_packed_tensor = fp4_packed_tensor[:, :self.mapping. - moe_ep_size * - all_rank_max_num_tokens] assert fp4_packed_tensor.ndim == 3 and fp4_packed_tensor.shape[ 2] == packed_hidden_size x_sf = fp4_packed_tensor[:, :, x.shape[1]:x.shape[1] + @@ -636,29 +647,6 @@ def forward_chunk( f"Not available alltoall method type: {self.alltoall_method_type!r}" ) - if self.enable_alltoall: - # Adapter between `torch.ops.trtllm.fused_moe` and DeepEP - # TODO: remove the adapter by changing APIs - if self.alltoall_method_type == AlltoallMethodType.DeepEP: - token_selected_slots = recv_topk_idx.to(torch.int32) - mask = token_selected_slots == -1 - token_selected_slots += self.expert_size_per_partition * self.mapping.moe_ep_rank - token_selected_slots[mask] = self.num_slots - num_recv_token_is_zero = x.shape[0] == 0 - if x.shape[0] == 0: - x = torch.zeros((1, x.shape[1]), - dtype=x.dtype, - device=x.device) - token_selected_slots = torch.full( - (1, token_selected_slots.shape[1]), - self.num_slots, - dtype=token_selected_slots.dtype, - device=token_selected_slots.device) - token_final_scales = torch.ones( - (1, token_final_scales.shape[1]), - dtype=token_final_scales.dtype, - device=token_final_scales.device) - final_hidden_states = torch.ops.trtllm.fused_moe( x, token_selected_slots, @@ -676,7 +664,7 @@ def forward_chunk( ep_rank=ep_rank, cluster_size=cluster_size, cluster_rank=cluster_rank, - enable_alltoall=self.enable_alltoall, + enable_alltoall=use_all_to_all, use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale, use_w4a8_group_scaling=use_w4a8_group_scaling, min_latency_mode=False, @@ -691,36 +679,22 @@ def forward_chunk( # Otherwise, the output should be unpacked as a single tensor. final_hidden_states = final_hidden_states[0] - if self.enable_alltoall: + if use_all_to_all: if self.alltoall_method_type == AlltoallMethodType.MNNVL: if self.enable_dummy_allreduce: self.dummy_allreduce() final_hidden_states = self.alltoall_combine( final_hidden_states, alltoall_info, token_count) elif self.alltoall_method_type == AlltoallMethodType.DeepEP: - if num_recv_token_is_zero: - final_hidden_states = final_hidden_states[:0] + final_hidden_states = self.unpad_tensors( + padded, final_hidden_states) final_hidden_states = self.deep_ep_buffer.combine( final_hidden_states, deep_ep_handle) elif self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: num_tokens_per_expert_for_fused_moe = self.mapping.moe_ep_size * all_rank_max_num_tokens - num_tokens_per_expert_for_deep_ep = self.deep_ep_max_num_tokens * self.mapping.moe_ep_size final_hidden_states = final_hidden_states.view( self.expert_size_per_partition, num_tokens_per_expert_for_fused_moe, self.hidden_size) - if num_tokens_per_expert_for_deep_ep != num_tokens_per_expert_for_fused_moe: - # Adapter between fused_moe num_tokens and DeepEP num_tokens - # This adapter can be removed if fused_moe accepts DeepEP num_tokens without overhead - final_hidden_states_for_fused_moe = final_hidden_states - final_hidden_states = torch.empty( - self.expert_size_per_partition, - self.deep_ep_max_num_tokens * self.mapping.moe_ep_size, - self.hidden_size, - dtype=final_hidden_states.dtype, - device=final_hidden_states.device) - final_hidden_states[:, : - num_tokens_per_expert_for_fused_moe] = final_hidden_states_for_fused_moe - del final_hidden_states_for_fused_moe # Release memory final_hidden_states = self.deep_ep_buffer.low_latency_combine( final_hidden_states, deep_ep_topk_idx, deep_ep_topk_weights, deep_ep_handle) @@ -747,11 +721,11 @@ def forward( ) -> torch.Tensor: assert all_rank_num_tokens is not None assert use_dp_padding is not None - num_rows = sum(all_rank_num_tokens) # in case of num_rows is larger than max_chunk_size, we need to split the input into multiple chunks - num_chunks = (num_rows + self.moe_max_num_tokens - - 1) // self.moe_max_num_tokens + num_chunks = self.calculate_num_chunks(all_rank_num_tokens) + use_all_to_all = self.can_use_alltoall(all_rank_num_tokens, + all_rank_max_num_tokens) if use_dp_padding: all_rank_num_tokens_padded = [all_rank_max_num_tokens @@ -764,6 +738,7 @@ def forward( outputs = self.forward_chunk( x, router_logits, + use_all_to_all, output_dtype, all_rank_num_tokens=all_rank_num_tokens_padded, all_rank_max_num_tokens=all_rank_max_num_tokens, @@ -771,6 +746,7 @@ def forward( repeating_info=(is_first_call, is_last_call)) outputs = self.reducescatter_or_allreduce( outputs, + use_all_to_all, all_rank_num_tokens=all_rank_num_tokens_padded, use_dp_padding=use_dp_padding) else: @@ -792,7 +768,7 @@ def split_chunk(split_token_num: int, split_num_chunks: int): all_rank_max_num_tokens_list = split_chunk(all_rank_max_num_tokens, num_chunks) chunk_size_list = all_rank_chunk_size_list[self.rank] - if self.enable_alltoall: + if use_all_to_all: all_rank_num_tokens_list = [[ 1 if val == 0 else val for val in val_list ] for val_list in all_rank_num_tokens_list] @@ -804,7 +780,7 @@ def split_chunk(split_token_num: int, split_num_chunks: int): x_list = x.split(chunk_size_list) router_logits_list = router_logits.split(chunk_size_list) - if not self.enable_alltoall: + if not use_all_to_all: self.event_dict[EventType.Main].record() with torch.cuda.stream(self.aux_stream): self.event_dict[EventType.Main].wait() @@ -815,12 +791,13 @@ def split_chunk(split_token_num: int, split_num_chunks: int): zip(x_list, router_logits_list)): is_first_call = idx_chunk == 0 and self.repeat_idx == 0 is_last_call = idx_chunk == num_chunks - 1 and self.repeat_idx == self.repeat_count - 1 - if not self.enable_alltoall: + if not use_all_to_all: if idx_chunk % 2 == 0: with torch.cuda.stream(self.aux_stream): outputs = self.forward_chunk( x, router_logits, + use_all_to_all, all_rank_num_tokens=all_rank_num_tokens_list[ idx_chunk], all_rank_max_num_tokens= @@ -830,6 +807,7 @@ def split_chunk(split_token_num: int, split_num_chunks: int): if idx_chunk > 0: outputs_list[-1] = self.reducescatter_or_allreduce( outputs_list[-1], + use_all_to_all, all_rank_num_tokens=all_rank_num_tokens_list[ idx_chunk - 1], use_dp_padding=use_dp_padding) @@ -837,6 +815,7 @@ def split_chunk(split_token_num: int, split_num_chunks: int): outputs = self.forward_chunk( x, router_logits, + use_all_to_all, all_rank_num_tokens=all_rank_num_tokens_list[ idx_chunk], all_rank_max_num_tokens=all_rank_max_num_tokens_list[ @@ -846,6 +825,7 @@ def split_chunk(split_token_num: int, split_num_chunks: int): with torch.cuda.stream(self.aux_stream): outputs_list[-1] = self.reducescatter_or_allreduce( outputs_list[-1], + use_all_to_all, all_rank_num_tokens=all_rank_num_tokens_list[ idx_chunk - 1], use_dp_padding=use_dp_padding) @@ -853,22 +833,25 @@ def split_chunk(split_token_num: int, split_num_chunks: int): outputs = self.forward_chunk( x, router_logits, + use_all_to_all, all_rank_num_tokens=all_rank_num_tokens_list[idx_chunk], all_rank_max_num_tokens=all_rank_max_num_tokens_list[ idx_chunk], repeating_info=(is_first_call, is_last_call)) outputs_list.append(outputs) - if not self.enable_alltoall: + if not use_all_to_all: if num_chunks % 2 == 0: outputs_list[-1] = self.reducescatter_or_allreduce( outputs_list[-1], + use_all_to_all, all_rank_num_tokens=all_rank_num_tokens_list[-1], use_dp_padding=use_dp_padding) else: with torch.cuda.stream(self.aux_stream): outputs_list[-1] = self.reducescatter_or_allreduce( outputs_list[-1], + use_all_to_all, all_rank_num_tokens=all_rank_num_tokens_list[-1], use_dp_padding=use_dp_padding) with torch.cuda.stream(self.aux_stream): @@ -883,7 +866,7 @@ def split_chunk(split_token_num: int, split_num_chunks: int): def alltoall_prepare_maybe_dispatch( self, all_rank_max_num_tokens: int, x: torch.Tensor, token_selected_slots: torch.Tensor, - token_final_scales: torch.Tensor, + token_final_scales: torch.Tensor, use_postquant_alltoall: bool, local_statistic_tensor: Optional[torch.Tensor]): top_k = self.routing_method.experts_per_token @@ -929,7 +912,7 @@ def alltoall_prepare_maybe_dispatch( gathered_token_final_scales, all_rank_max_num_tokens, self.num_slots, top_k, self.ep_rank, self.ep_size) - if not self.use_postquant_alltoall: + if not use_postquant_alltoall: assert not isinstance( x, Fp4QuantizedTensor ), "pre-quant alltoall doesn't support fp4 tensor" @@ -972,6 +955,40 @@ def alltoall_combine(self, final_hidden_states: torch.Tensor, return final_hidden_states + def pad_empty_recv_tensors( + self, x: torch.Tensor, x_sf: Optional[torch.Tensor], + recv_topk_idx: torch.Tensor, token_final_scales: torch.Tensor + ) -> Tuple[bool, torch.Tensor, Optional[torch.Tensor], torch.Tensor, + torch.Tensor]: + """ + Pad the output of DeepEP `dispatch` if the output length is zero. + We can remove the adapter if both `fused_moe` op and `swizzle_sf` + accept zero-length inputs. + """ + if x.shape[0] == 0: + padded = True + x = torch.zeros((1, x.shape[1]), dtype=x.dtype, device=x.device) + if x_sf is not None: + x_sf = torch.zeros((1, x_sf.shape[1]), + dtype=x_sf.dtype, + device=x_sf.device) + recv_topk_idx = torch.full((1, recv_topk_idx.shape[1]), + self.num_slots, + dtype=recv_topk_idx.dtype, + device=recv_topk_idx.device) + token_final_scales = torch.ones((1, token_final_scales.shape[1]), + dtype=token_final_scales.dtype, + device=token_final_scales.device) + else: + padded = False + return padded, x, x_sf, recv_topk_idx, token_final_scales + + def unpad_tensors(self, padded: bool, + final_hidden_states: torch.Tensor) -> torch.Tensor: + if padded: + final_hidden_states = final_hidden_states[:0] + return final_hidden_states + def register_parameter_weight_slot_fn(self, weight_name: str, local_slot_id: int): assert hasattr( diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 2a28f019033c..f957712e3e5a 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -1096,10 +1096,13 @@ def load_expert_w3_w1_weight_scale_nvfp4( orig_shape = dst_w3_w1_weight_scale.shape - dst_w3_w1_weight_scale.copy_( - torch.ops.trtllm.nvfp4_block_scale_interleave( - dst_w3_w1_weight_scale.view(float4_sf_dtype)).view( - self.block_scales_dtype).reshape(orig_shape)) + dst_w3_w1_weight_scale_interleaved = torch.ops.trtllm.nvfp4_block_scale_interleave( + dst_w3_w1_weight_scale.view(float4_sf_dtype)).view( + self.block_scales_dtype).reshape(orig_shape) + + torch.cuda.synchronize() + + dst_w3_w1_weight_scale.copy_(dst_w3_w1_weight_scale_interleaved) def load_expert_w2_weight_scale_nvfp4(self, module: torch.nn.Module, w2_weight_scale: torch.Tensor, @@ -1113,10 +1116,13 @@ def load_expert_w2_weight_scale_nvfp4(self, module: torch.nn.Module, orig_shape = dst_w2_weight_scale.shape - dst_w2_weight_scale.copy_( - torch.ops.trtllm.nvfp4_block_scale_interleave( - dst_w2_weight_scale.view(float4_sf_dtype)).view( - self.block_scales_dtype).reshape(orig_shape)) + dst_w2_weight_scale_interleaved = torch.ops.trtllm.nvfp4_block_scale_interleave( + dst_w2_weight_scale.view(float4_sf_dtype)).view( + self.block_scales_dtype).reshape(orig_shape) + + torch.cuda.synchronize() + + dst_w2_weight_scale.copy_(dst_w2_weight_scale_interleaved) class NVFP4TRTLLMGenFusedMoEMethod(NVFP4FusedMoEMethod): diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index ca9cb6501d09..1ef5be24c8b5 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -47,6 +47,12 @@ class TensorParallelMode(str, enum.Enum): def split_dim(cls, mode): return 1 if mode == cls.ROW else 0 + # Helper to shard the corresponding per-channel activation scales + # Which shard along the dimension orthogonal to the weights + @classmethod + def flip(cls, mode): + return cls.ROW if mode == cls.COLUMN else cls.COLUMN + def load_weight_shard( weight, @@ -110,12 +116,14 @@ def load_weights_vanilla_helper(module: Linear, weights: List[Dict]): weight = load_weight_shard(weights[0]['weight'], module.tp_size, module.tp_rank, module.tp_mode, device) - if module.has_w4a16_awq: + if module.has_weight_only_quant: # NOTE: without the preprocess during the runtime, the gemm output nan's. in order to use the preprocess_weights_for_mixed_gemm # we need to cast the weight to int8 first. + activation_dtype = torch.float8_e4m3fn if module.has_w4a8_awq else torch.float16 + weight_dtype, _ = get_weight_dtype_and_id(module) weight = preprocess_weights_for_mixed_gemm( - weight.T.to(torch.int8).contiguous().cpu(), torch.quint4x2, - torch.float16).cuda().contiguous() + weight.T.to(torch.int8).contiguous().cpu(), weight_dtype, + activation_dtype).cuda().contiguous() copy_weight(module.weight, weight) @@ -169,6 +177,27 @@ def load_weights_fused_gate_up_helper( return (gate_weight, up_weight) +def get_weight_dtype_and_id(module: Linear) -> tuple[torch.dtype, int]: + """ + Get weight dtype and weight_id for weight only quantization mode. + + Returns: + tuple[torch.dtype, int]: (weight_dtype, weight_id) where: + - weight_dtype: torch.int8 for INT8 weights, torch.quint4x2 for INT4 weights + - weight_id: 1 for INT8, 2 for INT4 (used for weight packing) + """ + assert module.quant_config is not None and module.quant_config.layer_quant_mode.is_weight_only( + ), "This function should only be called when the module has weight-only quantization enabled." + + if module.quant_config.layer_quant_mode.is_int8_weight_only(): + return torch.int8, 1 + elif module.quant_config.layer_quant_mode.is_int4_weight_only(): + return torch.quint4x2, 2 + else: + raise ValueError( + f"Unsupported quant_mode: {module.quant_config.layer_quant_mode}") + + class LinearMethodBase(ABC): """ Base class for all linear methods. @@ -562,7 +591,8 @@ def load_weights_vanilla(self, module: Linear, weights: List[Dict]) -> None: scale_name = self._get_scale_name(weights) weight_scale = load_weight_shard(weights[0][scale_name], module.tp_size, - module.tp_rank, module.tp_mode) + module.tp_rank, + module.tp_mode).squeeze() copy_weight(module.weight_scale, weight_scale) if "input_scale" in weights[0]: copy_weight(module.input_scale, weights[0]["input_scale"]) @@ -582,7 +612,8 @@ def load_weights_fused_qkv_linear(self, module: Linear, module.tp_rank, module.tp_mode) v_scale = load_weight_shard(weights[2][scale_name], module.tp_size, module.tp_rank, module.tp_mode) - fused_fp8_block_scale = torch.cat((q_scale, k_scale, v_scale)) + fused_fp8_block_scale = torch.cat((q_scale, k_scale, v_scale)).squeeze() + copy_weight(module.weight_scale, fused_fp8_block_scale) def load_weights_fused_gate_up_linear(self, module: Linear, @@ -597,7 +628,7 @@ def load_weights_fused_gate_up_linear(self, module: Linear, module.tp_rank, module.tp_mode) right_scale = load_weight_shard(weights[1][scale_name], module.tp_size, module.tp_rank, module.tp_mode) - fused_scale = torch.cat([left_scale, right_scale], dim=0) + fused_scale = torch.cat([left_scale, right_scale], dim=0).squeeze() copy_weight(module.weight_scale, fused_scale) @@ -873,6 +904,122 @@ def load_weights_fused_gate_up_linear(self, module: Linear, copy_weight(module.weight_scale, weight_scale) +class WeightOnlyQuantLinearMethod(LinearMethodBase): + + def create_weights(self, module: Linear, in_features: int, + out_features: int, bias: bool, + dtype: torch.dtype) -> None: + + _, weight_id = get_weight_dtype_and_id(module) + + # Quantized weights (int4 weights are packed into int8) + module.weight = Parameter(torch.empty( + (in_features, out_features // weight_id), dtype=torch.int8), + requires_grad=False) + + module.weight_scale = Parameter(torch.empty((out_features), + dtype=dtype), + requires_grad=False) + + if bias: + module.bias = Parameter(torch.empty((out_features), dtype=dtype), + requires_grad=False) + else: + module.register_parameter("bias", None) + + def apply(self, module: Linear, input: torch.Tensor, + bias: Optional[torch.Tensor]) -> torch.Tensor: + + weight_dtype, _ = get_weight_dtype_and_id(module) + bias = bias.contiguous() if bias is not None else None + + output = torch.ops.trtllm.weight_only_quant_gemm( + input, module.weight, weight_dtype, module.weight_scale, + module.dtype) + + return output + + def load_weight_scales( + self, + weights: List[Dict], + tp_size: int = 1, + tp_rank: int = 0, + tp_mode: Optional[TensorParallelMode] = None) -> List[torch.Tensor]: + device = torch.device("cuda") + q_weight_scale = load_weight_shard(weights[0]['weight_scale'], + tp_size, + tp_rank, + tp_mode, + device=device) + k_weight_scale = load_weight_shard(weights[1]['weight_scale'], + tp_size, + tp_rank, + tp_mode, + device=device) + v_weight_scale = load_weight_shard(weights[2]['weight_scale'], + tp_size, + tp_rank, + tp_mode, + device=device) + weight_scales = [q_weight_scale, k_weight_scale, v_weight_scale] + + return weight_scales + + def load_weights_vanilla(self, module: Linear, weights: List[Dict]) -> None: + load_weights_vanilla_helper(module, weights) + + device = torch.device('cuda') + weight_scale = load_weight_shard(weights[0]['weight_scale'], + module.tp_size, module.tp_rank, + module.tp_mode, device) + + copy_weight(module.weight_scale, weight_scale) + + def load_weights_fused_qkv_linear(self, module: Linear, + weights: List[Dict]) -> None: + q_weight, k_weight, v_weight = load_weights_fused_qkv_helper( + module, weights) + + fused_weight = torch.cat((q_weight, k_weight, v_weight)) + + weight_dtype, _ = get_weight_dtype_and_id(module) + fused_weight = preprocess_weights_for_mixed_gemm( + fused_weight.to(torch.int8).T.contiguous().cpu(), weight_dtype, + torch.float16).cuda().contiguous() + + copy_weight(module.weight, fused_weight) + + weight_scales = self.load_weight_scales(weights) + + # Create concatenated weight scale tensor + cat_weight_scale = torch.cat(weight_scales, dim=0) + copy_weight(module.weight_scale, cat_weight_scale) + + def load_weights_fused_gate_up_linear(self, module: Linear, + weights: List[Dict]) -> None: + device = torch.device('cuda') + weight_dtype, _ = get_weight_dtype_and_id(module) + gate_weight, up_weight = load_weights_fused_gate_up_helper( + module, weights) + + fused_weight = torch.cat((gate_weight, up_weight)) + + fused_weight = preprocess_weights_for_mixed_gemm( + fused_weight.to(torch.int8).T.contiguous().cpu(), weight_dtype, + torch.float16).cuda().contiguous() + + copy_weight(module.weight, fused_weight) + + left_scale = load_weight_shard(weights[0]['weight_scale'], + module.tp_size, module.tp_rank, + module.tp_mode, device).contiguous() + right_scale = load_weight_shard(weights[1]['weight_scale'], + module.tp_size, module.tp_rank, + module.tp_mode, device).contiguous() + fused_scale = torch.cat([left_scale, right_scale], dim=0) + copy_weight(module.weight_scale, fused_scale) + + class W4A16_AWQ_LinearMethod(LinearMethodBase): def create_weights(self, module: Linear, in_features: int, @@ -892,7 +1039,7 @@ def create_weights(self, module: Linear, in_features: int, f"for INT4 per-group quantization scale dimensions.") module.weight_scale = Parameter(torch.empty( - (out_features, in_features // group_size), dtype=dtype), + (in_features // group_size, out_features), dtype=dtype), requires_grad=False) # NOTE: Not in all linear we have this tensor - pre_quant_scale is computed as an average and merged with the # LayerNorm for QKV and Gate/Up projection layers when possible. we can see the tensor only for o_proj and down_proj @@ -908,19 +1055,19 @@ def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]) -> torch.Tensor: if module.pre_quant_scale is not None: - pre_quant_scale = module.pre_quant_scale.repeat(input.shape[0], 1) - input = torch.mul(input, pre_quant_scale) + input = input * module.pre_quant_scale bias = bias.contiguous() if bias is not None else None - output = torch.ops.trtllm.w4a16_gemm(input.to( - module.dtype).contiguous(), - module.weight, - module.weight_scale.T.contiguous(), - module.quant_config.group_size, - module.quant_config.has_zero_point, - bias, - zeros=None) + output = torch.ops.trtllm.finegrained_mixed_dtype_gemm( + input=input.to(module.dtype).contiguous(), + weight=module.weight, + scales=module.weight_scale, + group_size=module.quant_config.group_size, + has_zero_point=module.quant_config.has_zero_point, + output_dtype=module.dtype or input.dtype, + bias=bias, + zeros=None) return output def load_weight_scales( @@ -953,9 +1100,16 @@ def load_weights_vanilla(self, module: Linear, weights: List[Dict]) -> None: load_weights_vanilla_helper(module, weights) device = torch.device('cuda') - pre_quant_scale = load_weight_shard(weights[0]['pre_quant_scale'], - module.tp_size, module.tp_rank, - module.tp_mode, device) + + pre_quant_scale = load_weight_shard( + weights[0]["pre_quant_scale"], + module.tp_size, + module.tp_rank, + # pre_quant_scale applies to activation as opposed to weight, so flip tp_mode the other way around + TensorParallelMode.flip(module.tp_mode), + device, + ) + module.pre_quant_scale = Parameter( torch.ones((module.in_features, ), dtype=pre_quant_scale.dtype), requires_grad=False).to(device=device) @@ -965,7 +1119,7 @@ def load_weights_vanilla(self, module: Linear, weights: List[Dict]) -> None: module.tp_mode, device) copy_weight(module.pre_quant_scale, pre_quant_scale) - copy_weight(module.weight_scale, weight_scale) + copy_weight(module.weight_scale, weight_scale.T.contiguous()) def load_weights_fused_qkv_linear(self, module: Linear, weights: List[Dict]) -> None: @@ -982,7 +1136,7 @@ def load_weights_fused_qkv_linear(self, module: Linear, weight_scales = self.load_weight_scales(weights) # Create concatenated weight scale tensor - cat_weight_scale = torch.cat(weight_scales, dim=0) + cat_weight_scale = torch.cat(weight_scales, dim=0).T.contiguous() copy_weight(module.weight_scale, cat_weight_scale) def load_weights_fused_gate_up_linear(self, module: Linear, @@ -1004,8 +1158,248 @@ def load_weights_fused_gate_up_linear(self, module: Linear, right_scale = load_weight_shard(weights[1]['weight_scale'], module.tp_size, module.tp_rank, module.tp_mode, device).contiguous() - fused_scale = torch.cat([left_scale, right_scale], dim=0) + fused_scale = torch.cat([left_scale, right_scale], dim=0).T.contiguous() + copy_weight(module.weight_scale, fused_scale) + + +class W4A8_AWQ_LinearMethod(LinearMethodBase): + + def create_weights(self, module: Linear, in_features: int, + out_features: int, bias: bool, dtype: torch.dtype): + # Quantized weights + module.weight = Parameter(torch.empty( + (in_features, out_features // 2), + dtype=torch.int8, + ), + requires_grad=False) + + group_size = module.quant_config.group_size + if in_features % group_size != 0: + raise ValueError( + f"in_features ({module.in_features}) must be divisible by group_size ({group_size}) " + f"for INT4 per-group quantization scale dimensions.") + + # NOTE: for FP8 activation, scales needs to be float16 + module.weight_scale = Parameter(torch.empty( + (in_features // group_size, out_features), dtype=torch.float16), + requires_grad=False) + + # Similar to W4A16 AWQ, not all linears will have this tensor + module.pre_quant_scale = None + + module.input_scale = Parameter(torch.tensor(1., dtype=torch.float32), + requires_grad=False) + module.inv_input_scale = Parameter(torch.tensor(1., + dtype=torch.float32), + requires_grad=False) + + module.alpha = Parameter(torch.empty([1], dtype=torch.float32), + requires_grad=False) + + if bias: + module.bias = Parameter(torch.empty((out_features), dtype=dtype), + requires_grad=False) + else: + module.register_parameter("bias", None) + + def apply(self, module: Linear, input: torch.Tensor, + bias: Optional[torch.Tensor]): + """ + modelopt flow for w4a8_awq: + 1. multiply pre_quant_scale to input + 2. quantize input to fp8 using input_scale + 3. unpack_weights and multiply by weight_scales (int4 -> fp16) + 4. divied by weight_scale_2 (fp16 -> fp8 to allow gemm in fp8). + 5. apply gemm in fp8. + 6. rescale using alpha which is input_scale * weight_scale_2 + """ + if module.pre_quant_scale is not None: + input = input * module.pre_quant_scale + + if input.dtype == torch.float8_e4m3fn: + quantized_input = input + else: + quantized_input, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( + input, (module.input_scale)) + + bias = bias.contiguous() if bias is not None else None + + output = torch.ops.trtllm.finegrained_mixed_dtype_gemm( + input=quantized_input.contiguous(), + weight=module.weight, + scales=module.weight_scale, + group_size=module.quant_config.group_size, + has_zero_point=module.quant_config.has_zero_point, + output_dtype=module.dtype + or input.dtype, # NOTE: output_dtype can only be bf16/fp16 for W4A8 + alpha=module.alpha.item(), + bias=bias, + zeros=None) + + return output + + def load_weight_scales_w4a8(self, + weights: List[Dict], + tp_size: int = 1, + tp_rank: int = 0, + tp_mode: Optional[TensorParallelMode] = None): + # For concatenated weights (qkv_proj / up_gate_proj), the global scaling factors and input scaling factors should be shared. + input_scale = None + weight_scale_2 = None + weight_scale = [] + + device = torch.device("cuda") + + for w in weights: + if "input_scale" in w: + if input_scale is None: + input_scale = w["input_scale"][...] + else: + assert input_scale == w["input_scale"][ + ...], "The input_scale should be same for all the weights" + if "weight_scale" in w: + ws = load_weight_shard(w["weight_scale"], + tp_size, + tp_rank, + tp_mode, + device=device) + + weight_scale.append(ws.to(torch.float16)) + if "weight_scale_2" in w: + if weight_scale_2 is None: + weight_scale_2 = w["weight_scale_2"][...] + else: + assert weight_scale_2 == w["weight_scale_2"][ + ...], "The weight_scale_2 should be same for all the weights" + + # Compute scaling factor and alpha required by GEMM kernels (rescale the gemm output in fp8) + alpha = (input_scale.float() * weight_scale_2.float()) + + return input_scale, weight_scale, alpha, weight_scale_2 + + def load_weights_vanilla(self, module: Linear, weights: List[Dict]): + load_weights_vanilla_helper(module, weights) + + device = torch.device('cuda') + pre_quant_scale = load_weight_shard( + weights[0]["pre_quant_scale"], + module.tp_size, + module.tp_rank, + # pre_quant_scale applies to activation as opposed to weight, so flip tp_mode the other way around + TensorParallelMode.flip(module.tp_mode), + device, + ) + + assert pre_quant_scale.dtype == module.dtype + + module.pre_quant_scale = Parameter( + torch.empty((module.in_features, ), dtype=pre_quant_scale.dtype), + requires_grad=False).to(device=device) + + copy_weight(module.pre_quant_scale, pre_quant_scale) + + input_scale, weight_scale, alpha, weight_scale_2 = self.load_weight_scales_w4a8( + weights=weights, + tp_size=module.tp_size, + tp_rank=module.tp_rank, + tp_mode=module.tp_mode) + + assert len(weight_scale) == 1, "there should be only one weight scale" + + weight_scale = (weight_scale[0].T / weight_scale_2).contiguous() + + copy_weight(module.weight_scale, weight_scale) + copy_weight(module.input_scale, input_scale) + copy_weight(module.alpha, alpha) + + module.inv_input_scale.data = 1.0 / module.input_scale + + def load_weights_fused_qkv_linear(self, module: Linear, + weights: List[Dict]): + + q_weight, k_weight, v_weight = load_weights_fused_qkv_helper( + module, weights) + + fused_weight = torch.cat((q_weight, k_weight, v_weight)) + fused_weight = preprocess_weights_for_mixed_gemm( + fused_weight.to(torch.int8).T.contiguous().cpu(), torch.quint4x2, + torch.float8_e4m3fn).cuda().contiguous() + + copy_weight(module.weight, fused_weight) + + input_scale, weight_scales, alpha, weight_scale_2 = self.load_weight_scales_w4a8( + weights=weights, + tp_size=module.tp_size, + tp_rank=module.tp_rank, + tp_mode=module.tp_mode) + + # Create concatenated weight scale tensor + cat_weight_scale = (torch.cat(weight_scales, dim=0).T / + weight_scale_2).contiguous() + copy_weight(module.weight_scale, cat_weight_scale) + copy_weight(module.input_scale, input_scale) + copy_weight(module.alpha, alpha) + + # NOTE: pre_quant_scale is the same for q,k,v since modelopt checks which layer shared the same input and create an avg pre_quant_scale + # Usually when modelopt exports the quantized model, pre_quant_Scale is fused in the layer norm (this case relevant if fused is disabled - modelopt internal) + if "pre_quant_scale" in weights[0].keys(): + + pre_quant_scale = load_weight_shard( + weights[0]["pre_quant_scale"], + module.tp_size, + module.tp_rank, + # pre_quant_scale applies to activation as opposed to weight, so flip tp_mode the other way around + TensorParallelMode.flip(module.tp_mode), + torch.device('cuda'), + ) + + module.pre_quant_scale = Parameter( + torch.ones((module.in_features, ), dtype=pre_quant_scale.dtype), + requires_grad=False).to(device=torch.device('cuda')) + + copy_weight(module.pre_quant_scale, pre_quant_scale) + + def load_weights_fused_gate_up_linear(self, module: Linear, + weights: List[Dict]): + + gate_weight, up_weight = load_weights_fused_gate_up_helper( + module, weights) + + fused_weight = torch.cat((gate_weight, up_weight)) + fused_weight = preprocess_weights_for_mixed_gemm( + fused_weight.to(torch.int8).T.contiguous().cpu(), torch.quint4x2, + torch.float8_e4m3fn).cuda().contiguous() + + copy_weight(module.weight, fused_weight) + + input_scale, weight_scale, alpha, weight_scale_2 = self.load_weight_scales_w4a8( + weights=weights, + tp_size=module.tp_size, + tp_rank=module.tp_rank, + tp_mode=module.tp_mode) + + fused_scale = (torch.cat(weight_scale, dim=0).T / + weight_scale_2).contiguous() copy_weight(module.weight_scale, fused_scale) + copy_weight(module.input_scale, input_scale) + copy_weight(module.alpha, alpha) + + if "pre_quant_scale" in weights[0].keys(): + pre_quant_scale = load_weight_shard( + weights[0]["pre_quant_scale"], + module.tp_size, + module.tp_rank, + # pre_quant_scale applies to activation as opposed to weight, so flip tp_mode the other way around + TensorParallelMode.flip(module.tp_mode), + torch.device('cuda'), + ) + + # NOTE:Create this tensor in load_weights, since not all layer have this tensor and memory is not allocated for it (same as W4A16) + module.pre_quant_scale = Parameter( + torch.ones((module.in_features, ), dtype=pre_quant_scale.dtype), + requires_grad=False).to(device=torch.device('cuda')) + + copy_weight(module.pre_quant_scale, pre_quant_scale) def get_quant_method(quant_config: Optional[QuantConfig] = None): @@ -1022,9 +1416,15 @@ def get_quant_method(quant_config: Optional[QuantConfig] = None): return NVFP4LinearMethod() if quant_config.layer_quant_mode.has_w4a8_mxfp4_fp8(): return W4A8MXFP4FP8LinearMethod() + if quant_config.layer_quant_mode.is_weight_only( + ) and not quant_config.layer_quant_mode.has_per_group_scaling(): + return WeightOnlyQuantLinearMethod() if quant_config.layer_quant_mode.is_int4_weight_only_per_group( ) and quant_config.quant_algo == QuantAlgo.W4A16_AWQ: return W4A16_AWQ_LinearMethod() + if quant_config.layer_quant_mode.is_int4_weight_only_per_group( + ) and quant_config.quant_algo == QuantAlgo.W4A8_AWQ: + return W4A8_AWQ_LinearMethod() raise ValueError(f'unsupported quant mode: {quant_config.quant_mode}') @@ -1143,12 +1543,24 @@ def has_nvfp4(self): return self.quant_config is not None and self.quant_config.layer_quant_mode.has_nvfp4( ) + @property + def has_weight_only_quant(self): + assert self._weights_created + return self.quant_config is not None and self.quant_config.layer_quant_mode.is_weight_only( + ) + @property def has_w4a16_awq(self): assert self._weights_created return self.quant_config is not None and self.quant_config.layer_quant_mode.is_int4_weight_only_per_group( ) and self.quant_config.quant_algo == QuantAlgo.W4A16_AWQ + @property + def has_w4a8_awq(self): + assert self._weights_created + return self.quant_config is not None and self.quant_config.layer_quant_mode.is_int4_weight_only_per_group( + ) and self.quant_config.quant_algo == QuantAlgo.W4A8_AWQ + def apply_linear(self, input, bias, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 979bc83f218e..4754e693fc57 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -14,13 +14,14 @@ from tensorrt_llm.logger import logger from tensorrt_llm.lora_manager import (LoraConfig, get_default_trtllm_modules_to_hf_modules, - load_torch_hf_lora) + load_torch_lora) from tensorrt_llm.mapping import Mapping from ..model_config import ModelConfig -from ..speculative import get_spec_decoder +from ..speculative import get_num_extra_kv_tokens, get_spec_decoder from .config import PyTorchConfig from .config_utils import is_mla, is_nemotron_hybrid +from .guided_decoder import GuidedDecoder from .kv_cache_transceiver import AttentionTypeCpp, create_kv_cache_transceiver from .llm_request import ExecutorResponse from .model_engine import PyTorchModelEngine @@ -163,7 +164,7 @@ def _get_token_num_for_estimation(self) -> int: if spec_cfg is not None: num_extra_tokens_per_seq += spec_cfg.max_draft_len - num_extra_tokens_per_seq += spec_cfg.num_extra_kv_tokens + num_extra_tokens_per_seq += get_num_extra_kv_tokens(spec_cfg) for req in self._dummy_reqs: num_req_tokens = len(req.input_token_ids) + num_extra_tokens_per_seq # Requests cannot share KV cache blocks. Round up to nearest integer multiple of block size. @@ -414,19 +415,12 @@ def create_py_executor_instance( start_worker, sampler, drafter, + guided_decoder: Optional[GuidedDecoder] = None, lora_config: Optional[LoraConfig] = None, garbage_collection_gen0_threshold: Optional[int] = None) -> PyExecutor: kv_cache_manager = resources.get(ResourceManagerType.KV_CACHE_MANAGER, None) spec_config = model_engine.spec_config - if mapping.is_last_pp_rank( - ) and executor_config.guided_decoding_config is not None: - if spec_config is not None: - raise ValueError( - "Guided decoding is not supported with speculative decoding.") - if not pytorch_backend_config.disable_overlap_scheduler: - raise ValueError( - "Guided decoding is not supported with overlap scheduler.") logger.info( f"max_seq_len={executor_config.max_seq_len}, max_num_requests={executor_config.max_batch_size}, max_num_tokens={executor_config.max_num_tokens}, max_batch_size={executor_config.max_batch_size}" @@ -438,11 +432,13 @@ def create_py_executor_instance( f"Cannot overwrite existing resource manager {key}.") resources[key] = value + peft_cache_manager = None if lora_config is not None: from tensorrt_llm.bindings import LoraModule if len(lora_config.lora_dir) == 1: - load_torch_hf_lora(lora_config) + # Route to appropriate loader based on checkpoint source + load_torch_lora(lora_config) else: assert len(lora_config.lora_target_modules ) >= 1, "Expecting at least one lora target module" @@ -455,12 +451,25 @@ def create_py_executor_instance( num_experts = _try_infer_num_experts(model_engine.model.model_config) + num_attn_layers = model_binding_config.num_attention_layers() + per_layer_kv_heads = [ + model_binding_config.num_kv_heads(i) for i in range(num_attn_layers) + ] + num_kv_attention_heads = max(per_layer_kv_heads) + if len(set(per_layer_kv_heads)) > 1: + # NOTE: This code-path is currently untested and not validated. Can fail! + # This support is tracked in TRTLLM-6561 + logger.warning( + f"Non-uniform KV heads per layer detected, using max ({num_kv_attention_heads}) for LoRA. " + "This code-path is currently untested and not validated. May fail!" + ) + lora_modules = LoraModule.create_lora_modules( lora_module_names=lora_config.lora_target_modules, hidden_size=model_binding_config.hidden_size, mlp_hidden_size=model_binding_config.mlp_hidden_size, num_attention_heads=model_binding_config.num_heads, - num_kv_attention_heads=model_binding_config.num_heads, + num_kv_attention_heads=num_kv_attention_heads, attention_head_size=model_binding_config.head_size, tp_size=mapping.tp_size, num_experts=num_experts) @@ -472,7 +481,6 @@ def create_py_executor_instance( num_lora_modules = model_engine.model.model_config.pretrained_config.num_hidden_layers * \ len(lora_config.lora_target_modules + lora_config.missing_qkv_modules) - # TODO smor- need to figure out how to set these values executor_config.peft_cache_config = trtllm.PeftCacheConfig( num_device_module_layer=max_lora_rank * num_lora_modules * lora_config.max_loras, @@ -514,6 +522,7 @@ def create_py_executor_instance( capacity_scheduler = BindCapacityScheduler( max_num_sequences, kv_cache_manager.impl if kv_cache_manager is not None else None, + peft_cache_manager.impl if peft_cache_manager is not None else None, executor_config.scheduler_config.capacity_scheduler_policy, two_step_lookahead=mapping.has_pp()) mb_scheduler = BindMicroBatchScheduler(executor_config.max_batch_size, @@ -527,7 +536,6 @@ def create_py_executor_instance( cache_transceiver_config = executor_config.cache_transceiver_config kv_cache_transceiver = create_kv_cache_transceiver( mapping, kv_cache_manager, attention_type, cache_transceiver_config) - return PyExecutor( resource_manager, scheduler, @@ -544,6 +552,7 @@ def create_py_executor_instance( if spec_config is not None else 0, kv_cache_transceiver=kv_cache_transceiver, draft_model_engine=draft_model_engine, + guided_decoder=guided_decoder, start_worker=start_worker, garbage_collection_gen0_threshold=garbage_collection_gen0_threshold) diff --git a/tensorrt_llm/_torch/pyexecutor/config.py b/tensorrt_llm/_torch/pyexecutor/config.py index b1935a512347..483d220c2e10 100644 --- a/tensorrt_llm/_torch/pyexecutor/config.py +++ b/tensorrt_llm/_torch/pyexecutor/config.py @@ -1,6 +1,8 @@ from dataclasses import dataclass, field from typing import Dict, List, Optional, Union +from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import \ + BaseCheckpointLoader from tensorrt_llm.bindings.executor import ExecutorConfig from ...builder import BuildConfig @@ -71,6 +73,7 @@ class PyTorchConfig: torch_compile_piecewise_cuda_graph: bool = False # When torch compile is enabled, userbuffers is enabled by default torch_compile_enable_userbuffers: bool = True + torch_compile_max_num_streams: int = 1 # Enable autotuner only when torch compile is enabled # TODO: after it can be work stable in warmup stage @@ -117,7 +120,9 @@ def update_executor_config( speculative_config: Optional["DecodingBaseConfig"] = None, hf_model_dir: Optional[str] = None, max_input_len: Optional[int] = None, - max_seq_len: Optional[int] = None): + max_seq_len: Optional[int] = None, + checkpoint_format: Optional[str] = None, + checkpoint_loader: Optional[BaseCheckpointLoader] = None): if backend is None: return @@ -145,3 +150,31 @@ def update_executor_config( if max_seq_len is not None: executor_config.max_seq_len = max_seq_len + + executor_config.checkpoint_loader = _construct_checkpoint_loader( + backend, checkpoint_loader, checkpoint_format) + + +def _construct_checkpoint_loader( + backend: str, checkpoint_loader: Optional[BaseCheckpointLoader], + checkpoint_format: Optional[str]) -> Optional[BaseCheckpointLoader]: + if backend == "_autodeploy": + return None + + from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import \ + BaseCheckpointLoader + from tensorrt_llm._torch.models.modeling_utils import ( + get_checkpoint_weight_loader, get_config_loader) + + if checkpoint_loader is None: + checkpoint_weight_loader = get_checkpoint_weight_loader( + checkpoint_format)() + config_loader = get_config_loader(checkpoint_format)() + + checkpoint_loader = BaseCheckpointLoader.get( + checkpoint_format=checkpoint_format, + weight_loader=checkpoint_weight_loader, + weight_mapper=None, + config_loader=config_loader) + + return checkpoint_loader diff --git a/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py b/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py new file mode 100644 index 000000000000..2ec4f3c460f1 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py @@ -0,0 +1,589 @@ +import dataclasses +import datetime +import heapq +import queue +import threading +import time +from collections import deque, namedtuple +from typing import Dict, List, Optional, Tuple + +import torch + +from tensorrt_llm._utils import nvtx_range + +from ..distributed import Distributed +from .llm_request import ExecutorRequest, executor_request_to_llm_request +from .sampler import Sampler, TorchSampler + +SHUTDOWN_REQUEST_ID = -1 + + +@dataclasses.dataclass +class RequestQueueItem: + id: int + request: Optional[ExecutorRequest] = None + is_canceled_request: bool = False + query: Optional[list] = None # only used in `StarAttention` + + @property + def is_shutdown_request(self): + return self.id == SHUTDOWN_REQUEST_ID + + @property + def is_normal_request(self): + return not (self.is_shutdown_request or self.is_canceled_request) + + +class ExecutorRequestQueue: + """Handles fetching and processing of new requests from the request queue.""" + + def __init__(self, dist: Distributed, enable_attention_dp: bool, + max_batch_size: int, max_beam_width: int, + max_num_active_requests: int, enable_iter_perf_stats: bool, + is_disaggregated: bool): + self.dist = dist + self.request_queue: queue.Queue[RequestQueueItem] = queue.Queue() + self.waiting_queue: deque[RequestQueueItem] = deque() + self.canceled_req_ids = [] + self.enable_attention_dp = enable_attention_dp + self.max_beam_width = max_beam_width + self.max_num_active_requests = max_num_active_requests + self.is_disaggregated = is_disaggregated + self.enqueue_lock = threading.Lock() + self.next_request_id = max_batch_size + self.enable_iter_perf_stats = enable_iter_perf_stats + self.start_times = {} + self.active = True + + # State tracking + self.num_fetch_requests = 0 + self.num_fetch_requests_cur_rank = 0 + self.expected_num_active_requests = 0 + self.new_active_requests_queue_latency_ms = 0 + self.is_shutdown = False + self.should_exclude_last_generation_logits = False + + def _get_from_request_queue( + self, + timeout: Optional[datetime.timedelta]) -> List[RequestQueueItem]: + + items = [] + timeout_secs = timeout.total_seconds() if timeout is not None else None + try: + if self.request_queue.empty() and (timeout_secs is None + or timeout_secs > 0): + # if queue is empty and want to wait, wait + items.append(self.request_queue.get(timeout=timeout_secs)) + else: + # if not empty or don't want to wait, just return all items in queue + while True: + queue_item = self.request_queue.get_nowait() + items.append(queue_item) + except queue.Empty: + pass + return items + + def _get_from_waiting_queue( + self, + waiting_queue: deque[RequestQueueItem], + max_req_count: int, + ) -> List[RequestQueueItem]: + """Safely extracts up to max_req_count items from a deque. + + Args: + waiting_queue: The queue to pop items from. + max_req_count: Maximum items to retrieve. Returns empty list if <=0. + + Returns: + List of retrieved items (may be shorter than max_req_count if queue empties first). + """ + # Edge case handling + if max_req_count <= 0: # Handles negative/zero counts + return [] + + items = [] + req_count = 0 + while req_count < max_req_count and waiting_queue: + items.append(waiting_queue.popleft()) + req_count += 1 + return items + + def enqueue_requests(self, requests: List[ExecutorRequest]): + req_ids = [] + try: + self.enqueue_lock.acquire() + start_time = time.time() + for request in requests: + self.start_times[self.next_request_id] = start_time + self.request_queue.put( + RequestQueueItem(self.next_request_id, request)) + req_ids.append(self.next_request_id) + self.next_request_id += 1 + finally: + self.enqueue_lock.release() + return req_ids + + def enqueue_cancel_request(self, req_id: int): + try: + self.enqueue_lock.acquire() + self.request_queue.put( + RequestQueueItem(req_id, is_canceled_request=True)) + finally: + self.enqueue_lock.release() + + def enqueue_shutdown_request(self): + try: + self.enqueue_lock.acquire() + self.request_queue.put(RequestQueueItem(SHUTDOWN_REQUEST_ID)) + self.active = False + finally: + self.enqueue_lock.release() + + def enqueue_request(self, + request: ExecutorRequest, + query: Optional[list] = None): + try: + self.enqueue_lock.acquire() + assert self.active, "PyExecutor has already been shutdown." + req_id = self.next_request_id + if self.enable_iter_perf_stats: + self.start_times[req_id] = time.time() + + if query is not None: + self.request_queue.put(RequestQueueItem(req_id, request, query)) + else: + self.request_queue.put(RequestQueueItem(req_id, request)) + self.next_request_id += 1 + finally: + self.enqueue_lock.release() + + return req_id + + def can_enqueue_request(self) -> bool: + self.enqueue_lock.acquire() + can_enqueue = self.active + self.enqueue_lock.release() + return can_enqueue and self.dist.rank == 0 + + def _fetch_and_process_requests( + self, total_num_active_requests: int, + total_max_num_active_requests: int) -> List[RequestQueueItem]: + """Common logic for fetching and processing requests from the queue.""" + # Calculate timeout + timeout = None if (total_num_active_requests == 0) and len( + self.waiting_queue) == 0 else datetime.timedelta(0) + + # Fetch requests from rank 0 + new_requests = [] + if self.dist.rank == 0: + new_requests = self._get_from_request_queue(timeout) + + # Broadcast requests and handle Python objects + new_requests, py_request_objects = self._handle_request_broadcasting( + new_requests) + + # Validate and filter requests + new_requests = self._validate_and_filter_requests(new_requests) + + # Attach Python objects to requests + if py_request_objects and (self.dist.tp_size > 1 + or self.dist.has_pp) and self.dist.rank > 0: + self._attach_py_objects_to_requests(new_requests, + py_request_objects) + + self.waiting_queue.extend(new_requests) + + new_requests = self._get_from_waiting_queue( + self.waiting_queue, + total_max_num_active_requests - total_num_active_requests) + + # Update performance metrics + if self.enable_iter_perf_stats and self.dist.rank == 0: + self._update_new_active_requests_queue_latency(new_requests) + + return new_requests + + @nvtx_range("_fetch_new_requests") + def fetch_new_requests(self, + num_active_requests: int) -> List[RequestQueueItem]: + + if self.enable_attention_dp: + return self._fetch_new_requests_attention_dp(num_active_requests) + else: + return self._fetch_new_requests_attention_tp(num_active_requests) + + def _fetch_new_requests_attention_tp( + self, num_active_requests: int) -> List[RequestQueueItem]: + """Handle standard (non-attention DP) request fetching.""" + total_num_active_requests = num_active_requests + total_max_num_active_requests = self.max_num_active_requests + + # Use common request fetching logic + new_requests = self._fetch_and_process_requests( + total_num_active_requests, total_max_num_active_requests) + + # Merge requests and add to active list + merged_requests = self._merge_requests(new_requests) + return merged_requests + + def _fetch_new_requests_attention_dp( + self, num_active_requests: int) -> List[RequestQueueItem]: + """Handle attention DP request fetching with load balancing.""" + # Get active request counts across all ranks + all_ranks_num_active_requests = [] + responses_list = self.dist.tp_allgather(num_active_requests) + for num_active_requests in responses_list: + all_ranks_num_active_requests.append(num_active_requests) + + total_num_active_requests = sum(all_ranks_num_active_requests) + total_max_num_active_requests = self.dist.tp_size * self.max_num_active_requests + + # Use common request fetching logic + new_requests = self._fetch_and_process_requests( + total_num_active_requests, total_max_num_active_requests) + + # Balance requests across ranks + num_new_requests_all_ranks = len(new_requests) + self.expected_num_active_requests = max( + (total_num_active_requests + num_new_requests_all_ranks + + self.dist.tp_size - 1) // self.dist.tp_size, + max(all_ranks_num_active_requests), + ) + + new_requests_cur_rank = self._balance_requests_across_ranks( + new_requests, all_ranks_num_active_requests) + + # Update performance metrics + if self.enable_iter_perf_stats and self.start_times: + self._update_new_active_requests_queue_latency( + new_requests_cur_rank) + + # Update counters + self.num_fetch_requests += num_new_requests_all_ranks + self.num_fetch_requests_cur_rank += len(new_requests_cur_rank) + + # Merge requests and add to active list + new_requests_cur_rank = self._merge_requests(new_requests_cur_rank) + return new_requests_cur_rank + + def _handle_request_broadcasting(self, + new_requests: List[RequestQueueItem]): + """Handle broadcasting of requests and Python objects across ranks.""" + if self.dist.rank == 0: + py_logits_post_processors = self._collect_py_objects_from_requests( + new_requests, "py_logits_post_processors") + py_multimodal_data = self._collect_py_objects_from_requests( + new_requests, "py_multimodal_data") + py_request_objects = tuple( + filter(None, [py_logits_post_processors, py_multimodal_data])) + else: + py_request_objects = None + + if self.dist.rank == 0: + # Preserve original `new_requests` on rank 0 + _ = self._broadcast_new_requests(new_requests, py_request_objects) + else: + new_requests, py_request_objects = self._broadcast_new_requests( + new_requests, py_request_objects) + + return new_requests, py_request_objects + + def _validate_and_filter_requests( + self, + new_requests: List[RequestQueueItem]) -> List[RequestQueueItem]: + """Validate and filter requests, handling shutdown signals.""" + valid_new_requests = [] + for req_item in new_requests: + if req_item.is_shutdown_request: + self.is_shutdown = True + break + elif req_item.is_canceled_request: + self.canceled_req_ids.append(req_item.id) + else: + valid_new_requests.append(req_item) + + # Check beam width validation + for req_item in valid_new_requests: + if req_item.request and hasattr(req_item.request, + 'sampling_config'): + assert req_item.request.sampling_config.beam_width == self.max_beam_width, \ + f"Request beam width {req_item.request.sampling_config.beam_width} " \ + f"is not equal to max_beam_width {self.max_beam_width}. This is not supported!" + + return valid_new_requests + + def _balance_requests_across_ranks( + self, new_requests: List[RequestQueueItem], + all_ranks_num_active_requests: List[int]) -> List[RequestQueueItem]: + """Balance requests across ranks for attention DP.""" + new_requests_cur_rank = [] + + if new_requests and self.expected_num_active_requests > all_ranks_num_active_requests[ + self.dist.tp_rank]: + # Balance context tokens across ranks using heap + HeapVal = namedtuple( + 'HeapVal', + ['num_tokens', 'num_requests', 'rank', 'request_list']) + + all_ranks_new_requests_heap = [ + HeapVal(0, self.expected_num_active_requests - val, tp_rank, []) + for tp_rank, val in enumerate(all_ranks_num_active_requests) + ] + + new_requests_cur_rank = all_ranks_new_requests_heap[ + self.dist.tp_rank].request_list + all_ranks_new_requests_heap = [ + val for val in all_ranks_new_requests_heap + if val.num_requests > 0 + ] + heapq.heapify(all_ranks_new_requests_heap) + + # Sort by token count (descending) for better load balancing + new_requests = sorted( + new_requests, + key=lambda x: len(getattr(x.request, 'input_token_ids', [])) + if x.request else 0, + reverse=True) + + # Distribute requests across ranks + for req_item in new_requests: + val = heapq.heappop(all_ranks_new_requests_heap) + token_count = len( + getattr(req_item.request, 'input_token_ids', + [])) if req_item.request else 0 + val = val._replace( + num_tokens=val.num_tokens + token_count, + num_requests=val.num_requests - 1, + ) + val.request_list.append(req_item) + if val.num_requests > 0: + heapq.heappush(all_ranks_new_requests_heap, val) + elif val.rank == self.dist.tp_rank: + break + + return new_requests_cur_rank + + def _collect_py_objects_from_requests( + self, requests: List[RequestQueueItem], + attribute_name: str) -> Optional[Tuple[str, Dict]]: + """Collect Python-only objects from requests.""" + req_id_to_obj = {} + for item in requests: + if not item.is_normal_request: + continue + if item.request: + obj = getattr(item.request, attribute_name, None) + if obj is not None: + req_id_to_obj[item.id] = obj + return None if not req_id_to_obj else (attribute_name, req_id_to_obj) + + def _broadcast_new_requests( + self, new_requests: List[RequestQueueItem], py_request_objects + ) -> Tuple[List[RequestQueueItem], Optional[Dict]]: + """Broadcast new_requests and optional Python-only metadata across pipeline stages.""" + payloads = (new_requests, py_request_objects) + + if not self.dist.has_pp: + return self.dist.broadcast(payloads, root=0) + + # Broadcast within first tp group before send/recv chain to other tp groups + if self.dist.tp_size > 1 and self.dist.is_first_pp_rank: + payloads = self.dist.tp_broadcast(payloads, root=0) + + # Tag for communication + tag = self.dist.pp_size # Use pp_size as tag to avoid conflicts + + # Send payloads + if not self.dist.is_first_pp_rank: + payloads = self.dist.recv_object(self.dist.prev_pp_rank, tag) + + if not self.dist.is_last_pp_rank: + self.dist.send_object(payloads, self.dist.next_pp_rank, tag) + + return payloads + + def _attach_py_objects_to_requests(self, requests: List[RequestQueueItem], + py_request_objects) -> None: + """Attach Python-only objects to each request.""" + for attr_name, req_obj_dict in py_request_objects: + for item in requests: + if item.request: + py_obj = req_obj_dict.get(item.id) + if py_obj is not None: + setattr(item.request, attr_name, py_obj) + + def _update_new_active_requests_queue_latency( + self, new_requests: List[RequestQueueItem]): + """Update queue latency metrics for new requests.""" + now = time.time() + for req_item in new_requests: + if req_item.id in self.start_times: + self.new_active_requests_queue_latency_ms += now - self.start_times.pop( + req_item.id) + + @nvtx_range("_merge_requests") + def _merge_requests(self, new_requests: list[RequestQueueItem]): + cp_config = self.dist.cp_config + if 'cp_type' in cp_config: + cp_type = cp_config['cp_type'] + if cp_type == 'star_attention': + return self._merge_star_attention_requests(new_requests) + elif cp_type == 'ring_attention': + raise NotImplementedError("ring attention not implemented yet") + else: + raise NotImplementedError(f'unsupport cp type {cp_type}') + else: + return [ + executor_request_to_llm_request( + req_item.id, req_item.request, + self._should_exclude_last_generation_logits()) + for req_item in new_requests + ] + + def _merge_star_attention_requests(self, + new_requests: list[RequestQueueItem]): + result = [] + for req_item in new_requests: + req_id, exe_req, query_token_ids = req_item.id, req_item.request, req_item.query + ctx_len0 = len(exe_req.input_token_ids) + ctx_blocks, position_blocks, last_block_padding_num = [ + exe_req.input_token_ids + ], [[i for i in range(ctx_len0)]], 0 + ctx_blocks, position_blocks, last_block_padding_num = self._partition_context( + exe_req.input_token_ids) + if self.dist.cp_rank == self.dist.cp_size - 1 and last_block_padding_num > 0: + ctx_blocks[-1] = ctx_blocks[-1][:-last_block_padding_num] + position_blocks[-1] = position_blocks[ + -1][:-last_block_padding_num] + #if has query + if query_token_ids: + ctx_blocks.append(query_token_ids) + position_blocks.append([ + i for i in range(ctx_len0, ctx_len0 + len(query_token_ids)) + ]) + + # insert the dummy block to align the number of ctx iterations of each rank + block_size = self.dist.cp_config['block_size'] + total_blocks = (ctx_len0 + block_size - 1) // block_size + num_blocks_per_rank = ( + total_blocks + self.dist.cp_size - + 1) // self.dist.cp_size + 1 # 1 for query block + if len(ctx_blocks) == num_blocks_per_rank: + ctx_blocks.insert(1, []) + position_blocks.insert(1, []) + elif len(ctx_blocks) == num_blocks_per_rank + 1: + # anchor + ctx_blocks + qry_block + pass + else: + print( + f'rank = {self.dist.cp_rank}, len(ctx_blocks) = {len(ctx_blocks) }, num_blocks_per_rank = {num_blocks_per_rank}' + ) + assert False, f'invalid context partition' + + # fake data for scheduler + ctx_blocks_list = [0] * (block_size + + self.dist.cp_config['cp_anchor_size']) + + req = executor_request_to_llm_request( + req_id, exe_req, self._should_exclude_last_generation_logits(), + ctx_blocks_list) + req.gen_iters = 0 + req.ctx_iters = 0 + req.ctx_blocks = ctx_blocks + req.ctx_position_blocks = position_blocks + req.query_id = query_token_ids + + result.append(req) + + return result + + def _partition_context(self, ctx_ids_list): + ctx_ids = torch.tensor(ctx_ids_list).unsqueeze(0) + ctx_len = ctx_ids.shape[-1] + block_size = self.dist.cp_config['block_size'] + if block_size is None: + block_size = ctx_len // self.dist.cp_size + anchor_block_size = self.dist.cp_config['cp_anchor_size'] + if anchor_block_size is None: + anchor_block_size = block_size + + assert anchor_block_size <= block_size, f'cp_anchor_size {anchor_block_size} should be smaller than block_size {block_size}' + padding = 0 + if ctx_len % block_size != 0: + padding = block_size - (ctx_len % block_size) + assert padding <= ctx_len, f'block size is too large for context, please set it smaller' + ctx_ids = torch.cat( + (ctx_ids, torch.zeros_like(ctx_ids)[:, :padding]), dim=-1) + position_ids = torch.arange(0, ctx_ids.shape[-1]).unsqueeze(0) + + ctx_ids_blocks = torch.tensor_split( + torch.stack(ctx_ids.split(block_size, dim=-1)), self.dist.cp_size) + position_ids_blocks = torch.tensor_split( + torch.stack(position_ids.split(block_size, dim=-1)), + self.dist.cp_size) + if self.dist.cp_rank != 0: + ctx_blocks, position_blocks = [ + ctx_ids_blocks[0][0].tolist()[0][:anchor_block_size] + ], [position_ids_blocks[0][0].tolist()[0][:anchor_block_size]] + else: + ctx_blocks, position_blocks = [], [] + + for idx in range(len(ctx_ids_blocks[self.dist.cp_rank])): + ctx_block = ctx_ids_blocks[self.dist.cp_rank][idx] + position_block = position_ids_blocks[self.dist.cp_rank][idx] + ctx_blocks.append(ctx_block.tolist()[0]) + position_blocks.append(position_block.tolist()[0]) + return ctx_blocks, position_blocks, padding + + def set_exclude_last_generation_logits(self, + disable_overlap_scheduler: bool, + sampler: Sampler) -> None: + # When overlap scheduler is enabled then when starting to handle a new prompt, + # sample_async is called twice before the first call to update_requests: + # - 1st time as a context request that handles on the 1st generated token + # - 2nd time as a generation request that handles on the 2nd generated token. + # and only after these two calls the sampler's update_request method is called. + # So in a sampler that works by the expected flow of handling the logits in + # sample_async (TorchSampler is an anomaly that instead does that on + # update_requests), every update_request doesn't handle the newest token, but one + # before it. Since all these calls work on the same request object, then its + # logits storage contains the logits of both the token update_requests should work + # on, and also its next token. Thus, excluding the last generation logits from any + # getter is required, when not using TorchSampler. + self.should_exclude_last_generation_logits = not disable_overlap_scheduler and not isinstance( + sampler, TorchSampler) + + def _should_exclude_last_generation_logits(self) -> bool: + return self.should_exclude_last_generation_logits + + def get_new_active_requests_queue_latency(self) -> float: + return self.new_active_requests_queue_latency_ms + + def get_expected_num_active_requests(self) -> int: + return self.expected_num_active_requests + + def get_request_queue_size(self) -> int: + return self.request_queue.qsize() + + def get_request_queue(self) -> queue.Queue[RequestQueueItem]: + return self.request_queue + + def get_waiting_queue(self) -> deque[RequestQueueItem]: + return self.waiting_queue + + def update_waiting_queue(self): + # Remove cancel request in the waiting queue + self.waiting_queue = deque(req for req in self.waiting_queue + if req.id not in self.canceled_req_ids) + + def get_waiting_queue_size(self) -> int: + return len(self.waiting_queue) + + def get_canceled_req_ids_size(self) -> int: + return len(self.canceled_req_ids) + + def get_canceled_req_ids(self) -> List[int]: + return self.canceled_req_ids + + def clear_canceled_req_ids(self): + self.canceled_req_ids.clear() diff --git a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py index 756c177a6ea6..f1b21339b9af 100644 --- a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py +++ b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py @@ -3,11 +3,11 @@ import torch +from ..._utils import nvtx_range from ...bindings.executor import GuidedDecodingConfig from .grammar_matcher import (GrammarMatcher, GrammarMatcherFactory, LLGuidanceMatcherFactory, XGrammarMatcherFactory) from .scheduler import ScheduledRequests -from .seq_slot_manager import SeqSlotManager class GuidedDecoder: @@ -49,12 +49,12 @@ def __init__(self, guided_decoding_config: GuidedDecodingConfig, def bitmask_size(self) -> int: return math.ceil(self.vocab_size_padded / 32) - def build(self, scheduled_requests: ScheduledRequests, - resource_manager: SeqSlotManager) -> None: + @nvtx_range("GuidedDecoder.build") + def build(self, scheduled_requests: ScheduledRequests) -> None: for llm_req in scheduled_requests.all_requests(): if llm_req.guided_decoding_params is None: continue - slot = resource_manager.slot_manager.get_slot(llm_req.request_id) + slot = llm_req.py_seq_slot if llm_req.is_context_init_state and llm_req.context_current_position == llm_req.prepopulated_prompt_len: self.grammar_matchers[ slot] = self.grammar_matcher_factory.create( @@ -75,8 +75,9 @@ def build(self, scheduled_requests: ScheduledRequests, self.bitmask[slot].copy_(self.bitmask_host[slot], non_blocking=True) + @nvtx_range("GuidedDecoder.execute") def execute(self, scheduled_requests: ScheduledRequests, - logits: torch.Tensor, resource_manager: SeqSlotManager) -> None: + logits: torch.Tensor) -> None: assert logits.size(0) == len(scheduled_requests.context_requests) + len( scheduled_requests.generation_requests) torch.cuda.current_stream().wait_stream(self._stream) @@ -88,8 +89,7 @@ def execute(self, scheduled_requests: ScheduledRequests, if llm_req.is_context_init_state and not llm_req.is_last_context_chunk: continue batched_logits.append(logits[i]) - slot = resource_manager.slot_manager.get_slot(llm_req.request_id) - batched_bitmask.append(self.bitmask[slot]) + batched_bitmask.append(self.bitmask[llm_req.py_seq_slot]) if len(batched_logits) > 0: torch.ops.trtllm.logits_bitmask(batched_logits, batched_bitmask) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index a7db4910b78c..547239b92045 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -2,6 +2,7 @@ from os import getenv import tensorrt_llm +from tensorrt_llm import logger from tensorrt_llm.bindings import WorldConfig from tensorrt_llm.bindings.executor import CacheTransceiverConfig from tensorrt_llm.mapping import Mapping @@ -10,9 +11,9 @@ from .resource_manager import KVCacheManager CacheTransceiverCpp = tensorrt_llm.bindings.internal.batch_manager.CacheTransceiver -CommTypeCpp = tensorrt_llm.bindings.internal.batch_manager.CommType AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType CacheTransBufferManagerCpp = tensorrt_llm.bindings.internal.batch_manager.CacheTransBufferManager +BackendTypeCpp = tensorrt_llm.bindings.executor.CacheTransceiverBackendType def mapping_to_world_config(mapping: Mapping) -> WorldConfig: @@ -30,23 +31,36 @@ def create_kv_cache_transceiver( mapping: Mapping, kv_cache_manager: KVCacheManager, attention_type: AttentionTypeCpp, cache_transceiver_config: CacheTransceiverConfig): - - comm_type = None - if getenv("TRTLLM_USE_UCX_KVCACHE"): - comm_type = CommTypeCpp.UCX - elif getenv("TRTLLM_USE_NIXL_KVCACHE"): - comm_type = CommTypeCpp.NIXL - elif getenv("TRTLLM_USE_MPI_KVCACHE"): - comm_type = CommTypeCpp.MPI - - cache_transceiver = None - if comm_type is not None: - cache_transceiver = BindKvCacheTransceiver(mapping, comm_type, - kv_cache_manager, - attention_type, - cache_transceiver_config) - - return cache_transceiver + if cache_transceiver_config is None or cache_transceiver_config.backend is None: + logger.info("cache_transceiver is disabled") + return None + + if cache_transceiver_config.backend == BackendTypeCpp.DEFAULT: + # When cache_transceiver_config.backend is not set, fallback to env_vars settings + # UCX is the default backend + cache_transceiver_config.backend = BackendTypeCpp.UCX + # Ordered by priority + env_vars = [("TRTLLM_USE_NIXL_KVCACHE", BackendTypeCpp.NIXL), + ("TRTLLM_USE_MPI_KVCACHE", BackendTypeCpp.MPI)] + for env_var, be_type in env_vars: + if getenv(env_var) == "1": + logger.warning( + f"{env_var}=1 is set, but it's recommended to set cache_transceiver_config.backend in yaml config" + ) + cache_transceiver_config.backend = be_type + break + + if cache_transceiver_config.backend == BackendTypeCpp.MPI: + logger.warning( + "MPI CacheTransceiver is deprecated, UCX or NIXL is recommended") + elif cache_transceiver_config.backend == BackendTypeCpp.UCX: + logger.info( + f"Using UCX kv-cache transceiver. If your devices are not in the same domain, please consider setting " + f"UCX_CUDA_IPC_ENABLE_MNNVL=n, UCX_RNDV_SCHEME=put_zcopy and/or unset UCX_NET_DEVICES upon server " + f"hangs or lower-than-expected performance.") + + return BindKvCacheTransceiver(mapping, kv_cache_manager, attention_type, + cache_transceiver_config) class KvCacheTransceiver(ABC): @@ -78,8 +92,7 @@ def check_gen_transfer_complete(self): class BindKvCacheTransceiver(KvCacheTransceiver): - def __init__(self, mapping: Mapping, comm_type: CommTypeCpp, - kv_cache_manager: KVCacheManager, + def __init__(self, mapping: Mapping, kv_cache_manager: KVCacheManager, attention_type: AttentionTypeCpp, cache_transceiver_config: CacheTransceiverConfig): world_config = mapping_to_world_config(mapping) @@ -88,7 +101,7 @@ def __init__(self, mapping: Mapping, comm_type: CommTypeCpp, tokens_per_block = kv_cache_manager.tokens_per_block dtype = kv_cache_manager.dtype - self.impl = CacheTransceiverCpp(kv_cache_manager.impl, comm_type, + self.impl = CacheTransceiverCpp(kv_cache_manager.impl, num_kv_heads_per_layer, head_dim, tokens_per_block, world_config, dtype, attention_type, @@ -120,7 +133,7 @@ def __init__(self, kv_cache_manager: KVCacheManager, max_num_tokens: int): max_num_tokens) @staticmethod - def pre_alloc_buffer_size(max_num_tokens: int, - kv_cache_size_per_token: int): + def pre_alloc_buffer_size(kv_cache_size_per_token: int, + cache_transceiver_config: CacheTransceiverConfig): return CacheTransBufferManagerCpp.pre_alloc_buffer_size( - max_num_tokens) * kv_cache_size_per_token + kv_cache_size_per_token, cache_transceiver_config) diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 461c5de941e7..7a7e4510dd0c 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -303,6 +303,7 @@ def __init__( self.py_batch_idx = None self.py_rewind_len = 0 self.py_draft_tokens = [] if self.draft_tokens is None else self.draft_tokens + self.py_last_context_chunk = (None, None) self.py_last_draft_tokens = None self.py_num_accepted_draft_tokens = 0 self.py_decoding_iter = 0 diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 42a0c0010764..2875f19b5b4f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2,31 +2,29 @@ import contextlib import functools import gc -import glob import inspect import math -import multiprocessing import os import traceback import weakref from abc import ABC, abstractmethod from contextlib import contextmanager -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, Optional, Tuple -import psutil -import safetensors import torch import torch._dynamo.config -import tqdm import tensorrt_llm.bindings.internal.userbuffers as ub +from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import \ + BaseCheckpointLoader from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors +from tensorrt_llm._torch.speculative import ( + get_num_extra_kv_tokens, update_spec_config_from_model_config) from tensorrt_llm._torch.speculative.mtp import SampleStateTensorsMTP -from tensorrt_llm._utils import (is_trace_enabled, local_mpi_rank, - local_mpi_size, nvtx_range, release_gc, +from tensorrt_llm._utils import (is_trace_enabled, nvtx_range, release_gc, torch_dtype_to_str, trace_func) -from tensorrt_llm.bindings.executor import GuidedDecodingConfig -from tensorrt_llm.inputs.multimodal import MultimodalParams +from tensorrt_llm.inputs.multimodal import (MultimodalParams, + MultimodalRuntimeData) from tensorrt_llm.logger import logger from tensorrt_llm.lora_manager import LoraConfig, LoraModelConfig from tensorrt_llm.mapping import Mapping @@ -48,7 +46,7 @@ from ..model_config import ModelConfig, MoeLoadBalancerConfig from ..models import AutoModelForCausalLM from ..models.modeling_utils import (DecoderModelForCausalLM, MetaInitMode, - run_concurrently, timing) + timing) from ..modules.fused_moe.moe_load_balancer import ( MoeLoadBalancer, MoeLoadBalancerIterContext, maybe_create_moe_load_balancer) from ..speculative import SpecMetadata, get_spec_metadata @@ -57,7 +55,6 @@ from .config import LoadFormat, PyTorchConfig from .config_utils import is_mla from .cuda_graph_runner import DecodingCUDAGraphRunner -from .guided_decoder import GuidedDecoder from .layerwise_nvtx_marker import LayerwiseNvtxMarker from .resource_manager import (BaseResourceManager, KVCacheManager, ResourceManager, ResourceManagerType) @@ -140,99 +137,6 @@ def validate_and_set_kv_cache_quant(model_config: ModelConfig, model_config.quant_config.kv_cache_quant_algo = mapped_pyt_quant -def _prefetch_one_file(file_name): - if os.path.exists(file_name): - logger.info(f"Prefetching {file_name} to memory...") - with open(file_name, 'rb') as f: - f.read() - logger.info(f"Finished prefetching {file_name}.") - - -def prefetch_files(file_names: List[str]): - """ - Prefetch safetensors files to memory so that the weight loading will be much faster. - When multiple ranks run in parallel, each rank will prefetch some files. - """ - - # Find out the files to prefetch for the current rank. - # Each rank loads files with indices local_rank, local_rank + local_mpi_size, local_rank + 2*local_mpi_size, etc. - local_file_names = file_names[local_mpi_rank()::local_mpi_size()] - if len(local_file_names) == 0: - return - - max_processes = min(multiprocessing.cpu_count() * 2, 16, - len(local_file_names)) - with multiprocessing.Pool(processes=max_processes) as pool: - pool.map(_prefetch_one_file, local_file_names) - - -def load_weights(checkpoint_dir: str): - weights = {} - weight_files = glob.glob(f"{checkpoint_dir}/*.safetensors") - if weight_files: - # Prefetch the weight files to CPU memory if the size is less than 90% of the available memory. - # This is a heuristic to avoid prefetching files that are too large and causing file cache thrashing. - prefetch_size = sum(os.path.getsize(file) for file in weight_files) - # If the layer number is overridden, it indicates that only a subset of layers are loaded. - # Prefetching all layers is unnecessary. - num_layers = int(os.environ.get("TLLM_OVERRIDE_LAYER_NUM", "0")) - enable_prefetch = prefetch_size < psutil.virtual_memory( - ).available * 0.9 and num_layers == 0 - if enable_prefetch: - logger.info( - f"Prefetching {prefetch_size / (1024**3):.2f}GB checkpoint files." - ) - prefetch_files(weight_files) - - def load_safetensors_file(file): - return safetensors.torch.load_file(file) - - pbar = tqdm.tqdm(total=len(weight_files), - desc="Loading safetensors weights in parallel") - - # Note that the function is called with a tuple of arguments, hence we need to wrap the arguments in a tuple via [(w,) for w in weight_files] - # specifically the comma right after the w is important to make it a tuple. - run_concurrently(load_safetensors_file, [(w, ) for w in weight_files], - reduce_func=weights.update, - pbar=pbar) - - return weights - - weight_files = glob.glob(f"{checkpoint_dir}/*.bin") - if not weight_files: - weight_files = glob.glob(f"{checkpoint_dir}/*.pth") - - if weight_files: - - def load_bin_or_path_file(file): - try: - part_weights = torch.load(file, - weights_only=True, - map_location='cpu', - mmap=True) - except Exception: - logger.warning( - f"Failed to load {file} with mmap=True, fallback to mmap=False" - ) - part_weights = torch.load(file, - weights_only=True, - map_location='cpu', - mmap=False) - finally: - return part_weights - - pbar = tqdm.tqdm(total=len(weight_files), - desc="Loading bin weights in parallel") - # Note that the function is called with a tuple of arguments, hence we need to wrap the arguments in a tuple via [(w,) for w in weight_files] - # specifically the comma right after the w is important to make it a tuple. - run_concurrently(load_bin_or_path_file, [(w, ) for w in weight_files], - reduce_func=weights.update, - pbar=pbar) - return weights - - raise RuntimeError(f"No weight files found in {checkpoint_dir}.") - - def initialize_dummy_weights( model: torch.nn.Module, low: float = -1e-3, @@ -346,6 +250,7 @@ def __init__( *, model_path: str, pytorch_backend_config: PyTorchConfig, + checkpoint_loader: BaseCheckpointLoader, batch_size: int = 8, max_beam_width: int = 1, max_num_tokens: int = 8192, @@ -354,7 +259,6 @@ def __init__( attn_runtime_features: Optional[AttentionRuntimeFeatures] = None, dist: Optional[MPIDist] = None, spec_config: Optional["DecodingBaseConfig"] = None, - guided_decoding_config: Optional[GuidedDecodingConfig] = None, lora_config: Optional[LoraConfig] = None, is_draft_model: bool = False, ): @@ -384,6 +288,7 @@ def __init__( self.model = self._load_model( model_path, mapping=self.mapping, + checkpoint_loader=checkpoint_loader, attn_backend=attn_backend, moe_backend=pytorch_backend_config.moe_backend, load_format=pytorch_backend_config.load_format, @@ -408,13 +313,6 @@ def __init__( self.dtype = self.model.config.torch_dtype self._init_model_capacity() - self.guided_decoder: Optional[GuidedDecoder] = None - if self.mapping.is_last_pp_rank( - ) and guided_decoding_config is not None: - self.guided_decoder = GuidedDecoder(guided_decoding_config, - self.batch_size, - self.model.vocab_size_padded) - self._torch_compile_backend = None try: @@ -428,7 +326,9 @@ def __init__( enable_piecewise_cuda_graph=pytorch_backend_config. torch_compile_piecewise_cuda_graph, cuda_graph_batch_sizes=pytorch_backend_config. - cuda_graph_batch_sizes) + cuda_graph_batch_sizes, + max_num_streams=pytorch_backend_config. + torch_compile_max_num_streams) if isinstance(self.model, DecoderModelForCausalLM): self.model.model = torch.compile( self.model.model, @@ -455,7 +355,8 @@ def __init__( if self.is_spec_decode: self.spec_metadata = None - self.spec_config.update_from_model_config(self.model.config) + update_spec_config_from_model_config(self.spec_config, + self.model.config) max_num_draft_tokens = self.spec_config.max_draft_len * batch_size self.draft_tokens_cuda = torch.empty((max_num_draft_tokens, ), dtype=torch.int, @@ -521,6 +422,17 @@ def __init__( self.lora_model_config: Optional[LoraModelConfig] = None self.cuda_graph_dummy_request = None + # Setup the local cache indirection buffer only once and reuse it. + # This way it can also be used for CUDA graphs. + if self.use_beam_search: + self.cache_indirection_attention = torch.zeros( + (self.batch_size, self.max_beam_width, self.max_seq_len + + (0 if self._disable_overlap_scheduler else 1)), + device="cuda", + dtype=torch.int32) + else: + self.cache_indirection_attention = None + def set_lora_model_config(self, lora_target_modules: list[str], trtllm_modules_to_hf_modules: dict[str, str]): self.lora_model_config = LoraModelConfig( @@ -540,6 +452,10 @@ def use_mrope(self): logger.info(f"Detected use_mrope: {use_mrope}") return use_mrope + @property + def use_beam_search(self): + return self.max_beam_width > 1 + @contextmanager def set_warmup_flag(self): self.in_warmup = True @@ -583,7 +499,9 @@ def warmup(self, resource_manager: ResourceManager) -> None: self.cuda_graph_dummy_request = None def get_cuda_graph_warmup_request(batch_size): - available_blocks = kv_cache_manager.get_num_free_blocks() + # Divide by max_beam_width to get an approximation of the number of requests that can be run in parallel. + available_blocks = kv_cache_manager.get_num_free_blocks( + ) // self.max_beam_width if available_blocks >= batch_size: result = ScheduledRequests() result.context_requests = [] @@ -594,9 +512,10 @@ def get_cuda_graph_warmup_request(batch_size): is_gen=True, max_num_draft_tokens=self.max_draft_len, use_mrope=use_mrope, - ) + max_beam_width=self.max_beam_width) + # Divide by max_beam_width to get an approximation of the number of tokens that can be added to the final request. available_tokens = kv_cache_manager.get_num_available_tokens( - self.max_draft_len) + self.max_draft_len) // self.max_beam_width # Add one dummy request with the maximum possible sequence length. # The sequence length is limited by both the max_seq_len and the number of available blocks. @@ -607,7 +526,7 @@ def get_cuda_graph_warmup_request(batch_size): is_gen=True, max_num_draft_tokens=self.max_draft_len, use_mrope=use_mrope, - )[0] + max_beam_width=self.max_beam_width)[0] # Add the longest request before all other seq_len=1 request to simulate the padding CUDA graph case. # This batch contains both the longest request and the shortest requests, # it also contains the maximum number of requests and the maximum token number, @@ -835,6 +754,7 @@ def _set_up_attn_metadata(self, kv_cache_manager: KVCacheManager): self.model.model_config.pretrained_config) and ( self.attn_runtime_features.cache_reuse or self.attn_runtime_features.chunked_prefill) + cache_indirection = self.cache_indirection_attention if self.attn_backend.Metadata is TrtllmAttentionMetadata else None if kv_cache_manager is None: return self.attn_backend.Metadata( max_num_requests=self.batch_size, @@ -844,7 +764,8 @@ def _set_up_attn_metadata(self, kv_cache_manager: KVCacheManager): mapping=self.mapping, runtime_features=self.attn_runtime_features, enable_flash_mla=self.model.model_config.enable_flash_mla, - enable_paged_context_mla=enable_paged_context_mla) + enable_paged_context_mla=enable_paged_context_mla, + cache_indirection=cache_indirection) if self.attn_metadata is not None: # This assertion can be relaxed if needed: just create a new metadata @@ -860,7 +781,9 @@ def _set_up_attn_metadata(self, kv_cache_manager: KVCacheManager): mapping=self.mapping, runtime_features=self.attn_runtime_features, enable_flash_mla=self.model.model_config.enable_flash_mla, - enable_paged_context_mla=enable_paged_context_mla) + enable_paged_context_mla=enable_paged_context_mla, + cache_indirection=cache_indirection) + return self.attn_metadata def _set_up_spec_metadata( @@ -891,7 +814,8 @@ def _get_padded_batch(self, scheduled_requests: ScheduledRequests, kv_cache_manager) -> int: can_run_cuda_graph = scheduled_requests.can_run_cuda_graph batch_size = scheduled_requests.batch_size - new_batch_size = batch_size + # The number of sequences in the batch is the number of prompts times the beam width. + new_batch_size = batch_size * self.max_beam_width if self._run_cuda_graphs and self.enable_attention_dp and self.mapping.tp_size > 1: graph_batch_size = self.dist.tp_allgather( [can_run_cuda_graph, batch_size]) @@ -927,7 +851,8 @@ def _get_padded_batch(self, scheduled_requests: ScheduledRequests, [MAX_UINT64 - 1], is_gen=True, max_num_draft_tokens=self.max_draft_len, - use_mrope=self.use_mrope)[0] + use_mrope=self.use_mrope, + max_beam_width=self.max_beam_width)[0] self.cuda_graph_dummy_request.is_cuda_graph_dummy = True scheduled_requests.generation_requests.extend( @@ -999,19 +924,21 @@ def _maybe_get_cuda_graph( if batch_size not in self._cuda_graph_batch_sizes: return None + num_sequences_in_batch = batch_size * self.max_beam_width attn_metadata = self.attn_metadata.create_cuda_graph_metadata( - batch_size, False, spec_max_draft_tokens) + num_sequences_in_batch, False, spec_max_draft_tokens) assert attn_metadata.is_cuda_graph if self.is_spec_decode: spec_metadata = self.spec_metadata.create_cuda_graph_metadata( - batch_size) + num_sequences_in_batch) spec_metadata.draft_tokens = self.draft_tokens_cuda else: spec_metadata = None self._cuda_graphs[batch_size] = DecodingCUDAGraphRunner( - batch_size, "cuda", attn_metadata, spec_metadata, self.use_mrope) + num_sequences_in_batch, "cuda", attn_metadata, spec_metadata, + self.use_mrope) return self._cuda_graphs[batch_size] def __del__(self) -> None: @@ -1023,13 +950,15 @@ def __del__(self) -> None: def _load_model(self, checkpoint_dir: str, + checkpoint_loader: BaseCheckpointLoader, load_format: LoadFormat, max_num_tokens: int, moe_max_num_tokens: Optional[int] = None, moe_load_balancer: Optional[MoeLoadBalancerConfig] = None, lora_config: Optional[LoraConfig] = None, **kwargs): - config = ModelConfig.from_pretrained( + + config = checkpoint_loader.load_config( checkpoint_dir, trust_remote_code=True, enable_min_latency=self.pytorch_backend_config.enable_min_latency, @@ -1082,20 +1011,24 @@ def init_meta_tensor(t: torch.Tensor): logger.info( f"Use {rank_model_storage / (1024**3):.2f} GB for model weights." ) - if load_format == LoadFormat.AUTO: if hasattr(model, 'llm_checkpoint_dir'): - weights = load_weights(model.llm_checkpoint_dir) + weights = checkpoint_loader.load_weights( + model.llm_checkpoint_dir) else: - weights = load_weights(checkpoint_dir) + weights = checkpoint_loader.load_weights(checkpoint_dir) - model.load_weights(weights) + weight_mapper = checkpoint_loader.get_initilized_weight_mapper( + model, config) + self._call_load_weights(model.load_weights, weights, + weight_mapper) if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights( ): - weights = load_weights( + weights = checkpoint_loader.load_weights( self.spec_config.speculative_model_dir) - model.load_draft_weights(weights) + self._call_load_weights(model.load_draft_weights, weights, + weight_mapper) elif load_format == LoadFormat.DUMMY: initialize_dummy_weights(model) @@ -1114,6 +1047,16 @@ def init_meta_tensor(t: torch.Tensor): torch.cuda.current_stream().synchronize() return model + def _call_load_weights(self, load_method, weights, weight_mapper): + # TODO smor- this is a temporary solution to load weights. + # Once checkpoint format is unified, this method will be removed. + from inspect import getfullargspec + args = getfullargspec(load_method).args + if "weight_mapper" in args: + load_method(weights, weight_mapper=weight_mapper) + else: + load_method(weights) + def _init_max_seq_len(self): inferred_max_seq_len = self.model.infer_max_seq_len() if self.max_seq_len is None: @@ -1229,8 +1172,16 @@ def _prepare_tp_inputs( num_cached_tokens_per_seq.append(past_seen_token_num) # Multimodal + # TODO: enable chunk prefill for multimodal (maybe need to pass prompt_tokens to MultimodalRuntimeData) + py_multimodal_runtime = MultimodalRuntimeData( + mm_token_lengths=request.multimodal_lengths, + mm_token_positions=request.multimodal_positions, + num_cached_tokens=past_seen_token_num + ) if request.multimodal_hashes is not None else None + multimodal_params = MultimodalParams( - multimodal_data=request.py_multimodal_data) + multimodal_data=request.py_multimodal_data, + multimodal_runtime=py_multimodal_runtime) multimodal_params.to_device("multimodal_data", "cuda", pin_memory=True) @@ -1238,7 +1189,7 @@ def _prepare_tp_inputs( if multimodal_params.has_content(): multimodal_params_list.append(multimodal_params) - request.py_batch_idx = request.seq_slot + request.py_batch_idx = request.py_seq_slot num_ctx_requests = len(scheduled_requests.context_requests) num_ctx_tokens = len(input_ids) @@ -1291,7 +1242,8 @@ def _prepare_tp_inputs( if next_draft_tokens_device is None or request.is_dummy or request.py_batch_idx is None: # get token ids, including input token ids and draft token ids. For these dummy requests, # no need to copy the token ids. - if not request.is_dummy: + if not (request.is_attention_dp_dummy + or request.is_cuda_graph_dummy): input_ids.append(request.get_last_tokens(0)) input_ids.extend(request.py_draft_tokens) draft_tokens.extend(request.py_draft_tokens) @@ -1320,11 +1272,11 @@ def _prepare_tp_inputs( num_cached_tokens_per_seq.append(past_seen_token_num) request_ids.append(request.py_request_id) # update batch index - request.py_batch_idx = request.seq_slot + request.py_batch_idx = request.py_seq_slot else: # update batch index previous_batch_idx = request.py_batch_idx - request.py_batch_idx = request.seq_slot + request.py_batch_idx = request.py_seq_slot # inputs # overlap scheduler can only support the speculative decoding # methods with a fixed number of draft tokens @@ -1378,8 +1330,8 @@ def _prepare_tp_inputs( gather_ids.append(len(position_ids) - 1) request_ids.append(request.py_request_id) - gen_request_seq_slots.append(request.seq_slot) - request.py_batch_idx = request.seq_slot + gen_request_seq_slots.append(request.py_seq_slot) + request.py_batch_idx = request.py_seq_slot previous_batch_len = len(previous_batch_indices) @@ -1394,7 +1346,6 @@ def previous_seq_slots_device(): num_tokens = len(input_ids) num_draft_tokens = len(draft_tokens) - num_requests = len(request_ids) total_num_tokens = len(position_ids) assert total_num_tokens <= self.max_num_tokens, ( "total_num_tokens should be less than or equal to max_num_tokens") @@ -1411,6 +1362,10 @@ def previous_seq_slots_device(): self.draft_tokens_cuda[:len(draft_tokens)].copy_(draft_tokens, non_blocking=True) if next_draft_tokens_device is not None: + # Initialize these two values to zeros + self.previous_pos_id_offsets_cuda *= 0 + self.previous_kv_lens_offsets_cuda *= 0 + if previous_batch_len > 0: previous_slots = previous_seq_slots_device() # previous input ids @@ -1435,24 +1390,37 @@ def previous_seq_slots_device(): pin_memory=True) self.previous_pos_indices_cuda[0:previous_batch_tokens].copy_( previous_pos_indices_host, non_blocking=True) + + # The order of requests in a batch: [context requests, generation requests] + # generation requests: ['requests that do not have previous batch', 'requests that already have previous batch', 'dummy requests'] + # 1) 'requests that do not have previous batch': disable overlap scheduler or the first step in the generation server of disaggregated serving. + # 2) 'requests that already have previous batch': previous iteration's requests. + # 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. + # Therefore, both of self.previous_pos_id_offsets_cuda and self.previous_kv_lens_offsets_cuda are also 3 segments. + # For 1) 'requests that do not have previous batch': disable overlap scheduler or the first step in the generation server of disaggregated serving. + # Set these requests' previous_pos_id_offsets and previous_kv_lens_offsets to '0' to skip the value changes in _preprocess_inputs. + # Already set to '0' during initialization. + # For 2) 'requests that already have previous batch': enable overlap scheduler. + # Set their previous_pos_id_offsets and previous_kv_lens_offsets according to new_tokens_lens_device and kv_len_offsets_device. + # For 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. + # Already set to '0' during initialization. + + num_extend_reqeust_wo_dummy = len(extend_requests) - len( + extend_dummy_requests) self.previous_pos_id_offsets_cuda[ - 0:previous_batch_tokens].copy_( + (num_extend_reqeust_wo_dummy - previous_batch_len) * + (1 + self.max_draft_len):num_extend_reqeust_wo_dummy * + (1 + self.max_draft_len)].copy_( new_tokens_lens_device[self.previous_pos_indices_cuda[ 0:previous_batch_tokens]], non_blocking=True) - self.previous_kv_lens_offsets_cuda[0:previous_batch_len].copy_( - kv_len_offsets_device[previous_slots], non_blocking=True) - # for the requests that do not have previous batch, set the previous_pos_id_offsets and - # previous_kv_lens_offsets to zeros to skip the value changes in _preprocess_inputs - self.previous_pos_id_offsets_cuda[ - previous_batch_tokens:num_requests * - (1 + self.max_draft_len)] *= 0 + self.previous_kv_lens_offsets_cuda[ - previous_batch_len:num_requests] *= 0 - else: - # change the data to zeros to skip the value changes in _preprocess_inputs - self.previous_pos_id_offsets_cuda *= 0 - self.previous_kv_lens_offsets_cuda *= 0 + num_extend_reqeust_wo_dummy - + previous_batch_len:num_extend_reqeust_wo_dummy].copy_( + kv_len_offsets_device[previous_slots], + non_blocking=True) + elif new_tokens_device is not None: seq_slots_device = previous_seq_slots_device() max_draft_len = max(draft_lens) @@ -1463,6 +1431,16 @@ def previous_seq_slots_device(): previous_batch_len * self.max_beam_width].copy_( new_tokens.flatten(), non_blocking=True) + if (not self._disable_overlap_scheduler + and next_draft_tokens_device is None + and len(extend_requests) > 0): + # During warmup, for those generation requests, we don't have previous tensors, + # so we need to set the previous_pos_id_offsets and previous_kv_lens_offsets to zeros + # to skip the value changes in _preprocess_inputs. Otherwise, there will be illegal memory access + # when writing key/values to the KV cache. + self.previous_pos_id_offsets_cuda *= 0 + self.previous_kv_lens_offsets_cuda *= 0 + position_ids = torch.tensor(position_ids, dtype=torch.int, pin_memory=True) @@ -1484,16 +1462,16 @@ def previous_seq_slots_device(): num_generation_requests = len(scheduled_requests.generation_requests) # Cache indirection is only used for beam search on generation requests - if self.max_beam_width > 1 and num_generation_requests > 0 and cache_indirection_buffer is not None: - cache_indirection_attention = torch.zeros_like( - cache_indirection_buffer) - #Copy cache indirection to local buffer with offsets changing: seq_slots[i] -> i - cache_indirection_attention[:num_generation_requests].copy_( - cache_indirection_buffer[gen_request_seq_slots]) - attn_metadata.cache_indirection = cache_indirection_attention - attn_metadata.beam_width = self.max_beam_width + if self.use_beam_search and num_generation_requests > 0: + # CUDA Graph needs to set beam width during warmup (where the graph is captured), to ensure that cache indirection buffer is correctly picked up by the CUDA graph + is_cuda_graph_during_warmup = self.in_warmup and attn_metadata.is_cuda_graph + if cache_indirection_buffer is not None: + #Copy cache indirection to local buffer with offsets changing: seq_slots[i] -> i + self.cache_indirection_attention[:num_generation_requests].copy_( + cache_indirection_buffer[gen_request_seq_slots]) + if cache_indirection_buffer is not None or is_cuda_graph_during_warmup: + attn_metadata.beam_width = self.max_beam_width else: - attn_metadata.cache_indirection = None attn_metadata.beam_width = 1 attn_metadata.request_ids = request_ids @@ -1506,8 +1484,7 @@ def previous_seq_slots_device(): attn_metadata.kv_cache_params = KVCacheParams( use_cache=True, num_cached_tokens_per_seq=num_cached_tokens_per_seq, - num_extra_kv_tokens=0 if self.spec_config is None else - self.spec_config.num_extra_kv_tokens) + num_extra_kv_tokens=get_num_extra_kv_tokens(self.spec_config)) attn_metadata.kv_cache_manager = kv_cache_manager attn_metadata.prepare() @@ -2113,7 +2090,6 @@ def forward( with MoeLoadBalancerIterContext(moe_load_balancer): return self._forward_step(inputs, gather_ids, gather_context_logits) - with self._maybe_pad_batch(scheduled_requests, kv_cache_manager) as scheduled_requests: maybe_graph = self._maybe_get_cuda_graph( @@ -2160,18 +2136,6 @@ def capture_forward_fn(inputs: Dict[str, Any]): with MoeLoadBalancerIterContext(moe_load_balancer): outputs = maybe_graph.run(inputs) - # Note: To overlap the CPU and GPU computation as much as possible, - # guided_decoder.build should be called immediately after the launch of the single step; - # while guided_decoder.execute should be called right before the samplings. - # We can insert other CPU computation between them in the future. - if self.mapping.is_last_pp_rank( - ) and self.guided_decoder is not None: - seq_slot_manager = resource_manager.get_resource_manager( - ResourceManagerType.SEQ_SLOT_MANAGER) - self.guided_decoder.build(scheduled_requests, seq_slot_manager) - self.guided_decoder.execute(scheduled_requests, - outputs['logits'], seq_slot_manager) - self._execute_logit_post_processors(scheduled_requests, outputs) return outputs @@ -2182,6 +2146,14 @@ def model_forward(self, **kwargs): attrs["attention_metadata"] = weakref.ref(kwargs['attn_metadata']) attrs.update(self.model.model_config.extra_attrs) + if self._torch_compile_backend is not None: + # Register aux streams and events to model extra attrs. + # The streams and events are list which could be updated during compilation. + attrs["aux_streams"] = weakref.ref( + self._torch_compile_backend.aux_streams) + attrs["events"] = weakref.ref(self._torch_compile_backend.events) + attrs["global_stream"] = torch.cuda.current_stream() + if is_trace_enabled("TLLM_TRACE_MODEL_FORWARD"): return trace_func(self.model.forward)(**kwargs) else: @@ -2264,7 +2236,7 @@ def _execute_logit_post_processors(self, # Skip as we only need to apply logit processor on the last context request continue - logits_row = logits_tensor[request.py_batch_idx] + logits_row = logits_tensor[idx] # Reshape to align w/ the shape used in the TRT backend, # so the same logit processors can be used across both backends. logits_row = logits_row.view(1, 1, -1) @@ -2277,4 +2249,4 @@ def _execute_logit_post_processors(self, "defined in `tensorrtllm.sampling_params`.") lp(request.py_request_id, logits_row, token_ids, None, None) - logits_tensor[request.py_batch_idx] = logits_row.view(-1) + logits_tensor[idx] = logits_row.view(-1) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index dc7b79c265c8..715a70139856 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2,16 +2,13 @@ import datetime import functools import gc -import heapq import os -import queue import threading import time import traceback import weakref -from collections import deque, namedtuple from contextlib import contextmanager -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Optional, Union import torch @@ -23,7 +20,7 @@ FinishReason, InflightBatchingStats, IterationStats, KvCacheStats, RequestStage, RequestStats, - RequestType, SpecDecodingStats, + SpecDecodingStats, StaticBatchingStats) from tensorrt_llm.bindings.internal.batch_manager import (LlmRequestType, ReqIdsSet) @@ -31,11 +28,13 @@ from ..distributed import Distributed from ..speculative.drafter import Drafter +from .executor_request_queue import ExecutorRequestQueue, RequestQueueItem +from .guided_decoder import GuidedDecoder from .kv_cache_transceiver import KvCacheTransceiver from .llm_request import (ExecutorRequest, LlmRequest, LlmRequestState, - LlmResponse, executor_request_to_llm_request) + LlmResponse) from .model_engine import ModelEngine -from .sampler import Sampler, SampleState, SampleStateTensors, TorchSampler +from .sampler import Sampler, SampleState, SampleStateTensors from .scheduler import RequestScheduler, ScheduledRequests # Environment variable to specify iteration ranges for profiling start/stop. @@ -50,68 +49,6 @@ # Set to a path to save detailed tracing of PyTorch operations. PROFILE_TRACE_ENV_VAR_NAME = "TLLM_TORCH_PROFILE_TRACE" -SHUTDOWN_REQUEST_ID = -1 - - -@dataclasses.dataclass -class RequestQueueItem: - id: int - request: Optional[ExecutorRequest] = None - is_canceled_request: bool = False - query: Optional[list] = None # only used in `StarAttention` - - @property - def is_shutdown_request(self): - return self.id == SHUTDOWN_REQUEST_ID - - @property - def is_normal_request(self): - return not (self.is_shutdown_request or self.is_canceled_request) - - -def _get_from_request_queue( - request_queue, - timeout: Optional[datetime.timedelta]) -> List[RequestQueueItem]: - items = [] - timeout_secs = timeout.total_seconds() if timeout is not None else None - try: - if request_queue.empty() and (timeout_secs is None or timeout_secs > 0): - # if queue is empty and want to wait, wait - items.append(request_queue.get(timeout=timeout_secs)) - else: - # if not empty or don't want to wait, just return all items in queue - while True: - queue_item = request_queue.get_nowait() - items.append(queue_item) - except queue.Empty: - pass - return items - - -def _get_from_waiting_queue( - waiting_queue: deque[RequestQueueItem], - max_req_count: int, -) -> List[RequestQueueItem]: - """Safely extracts up to max_req_count items from a deque. - - Args: - waiting_queue: The queue to pop items from. - max_req_count: Maximum items to retrieve. Returns empty list if <=0. - - Returns: - List of retrieved items (may be shorter than max_req_count if queue empties first). - """ - # Edge case handling - if max_req_count <= 0: # Handles negative/zero counts - return [] - - items = [] - req_count = 0 - while req_count < max_req_count and waiting_queue: - items.append(waiting_queue.popleft()) - req_count += 1 - return items - @functools.cache def _load_iteration_indexes(env_var: str): @@ -204,13 +141,12 @@ def __init__(self, max_draft_len: int = 0, kv_cache_transceiver: Optional[KvCacheTransceiver] = None, draft_model_engine: Optional[ModelEngine] = None, + guided_decoder: Optional[GuidedDecoder] = None, garbage_collection_gen0_threshold: Optional[int] = None, start_worker: bool = True): super(PyExecutor, self).__init__() self.device_id = torch.cuda.current_device() self.global_rank = global_mpi_rank() - self.request_queue: queue.Queue[RequestQueueItem] = queue.Queue() - self.waiting_queue: deque[RequestQueueItem] = deque() # profile config self.profile_start_iters, self.profile_stop_iters = _load_iteration_indexes( @@ -225,6 +161,7 @@ def __init__(self, self.enable_attention_dp = model_engine.enable_attention_dp self.sampler = sampler self.drafter = drafter + self.guided_decoder = guided_decoder self.dist = dist self.disable_overlap_scheduler = disable_overlap_scheduler @@ -232,8 +169,6 @@ def __init__(self, self.draft_model_engine = draft_model_engine # enqueue and _fetch_new_requests used data - self.enqueue_lock = threading.Lock() - self.active = True self.next_req_id = max_batch_size # The first max_batch_size request IDs are reserved for dummy requests self.max_beam_width = max_beam_width self.max_draft_len = max_draft_len @@ -255,18 +190,11 @@ def __init__(self, ResourceManagerType.KV_CACHE_MANAGER) self.enable_kv_cache_events = self.kv_cache_manager is not None and self.kv_cache_manager.event_buffer_max_size > 0 - if self.draft_model_engine is not None and self.kv_cache_manager is not None: - if self.kv_cache_manager.enable_block_reuse: - raise NotImplementedError( - "Draft model engine + KV cache reuse is not supported yet. " - "This will be fixed in the near future!") - self.max_input_len = max_input_len # _executor_loop private data self.max_num_active_requests = model_engine.get_max_num_sequences() self.active_requests: List[LlmRequest] = [] self.expected_num_active_requests = 0 - self.has_context_request = False self.ctx_in_transmission_requests = [] self.previous_batch: Optional[BatchState] = None self.num_scheduled_requests: int = 0 @@ -280,7 +208,6 @@ def __init__(self, self.send_handles = [None] * self.num_micro_batches self.inflight_req_ids = ReqIdsSet() - self.canceled_req_ids = [] self.model_engine.warmup(self.resource_manager) if self.draft_model_engine is not None: @@ -288,10 +215,21 @@ def __init__(self, self.is_shutdown = False + # request fetcher initialization + self.executor_request_queue = ExecutorRequestQueue( + dist=self.dist, + enable_attention_dp=self.enable_attention_dp, + max_batch_size=max_batch_size, + max_beam_width=self.max_beam_width, + max_num_active_requests=self.max_num_active_requests, + enable_iter_perf_stats=self.enable_iter_perf_stats, + is_disaggregated=kv_cache_transceiver is not None, + ) + self.executor_request_queue.set_exclude_last_generation_logits( + self.disable_overlap_scheduler, self.sampler) + self.stats_lock = threading.Lock() self.stats = [] - self.start_times = {} - self.new_active_requests_queue_latency_ms = 0 self.gather_all_responses = False self.kv_cache_transceiver = kv_cache_transceiver @@ -299,13 +237,10 @@ def __init__(self, self.event_loop = self._executor_loop_pp else: self.event_loop = self._executor_loop if disable_overlap_scheduler else self._executor_loop_overlap - if not disable_overlap_scheduler and model_engine.max_beam_width > 1: - raise NotImplementedError( - "Overlap scheduler is not supported for beam search.") if is_trace_enabled("TLLM_TRACE_EXECUTOR_LOOP"): self.event_loop = trace_func(self.event_loop) - if self.draft_model_engine is not None: + if self.drafter is not None: if self.event_loop.__name__ != self._executor_loop.__name__: raise NotImplementedError( "Drafting is not supported for selected executor loop. " @@ -352,19 +287,7 @@ def enqueue_requests(self, requests: List[ExecutorRequest]): """ Enqueue new requests """ - req_ids = [] - try: - self.enqueue_lock.acquire() - assert self.active, "PyExecutor has already been shutdown." - start_time = time.time() - for request in requests: - self.start_times[self.next_req_id] = start_time - self.request_queue.put( - RequestQueueItem(self.next_req_id, request)) - req_ids.append(self.next_req_id) - self.next_req_id += 1 - finally: - self.enqueue_lock.release() + req_ids = self.executor_request_queue.enqueue_requests(requests) return req_ids def await_responses( @@ -397,23 +320,13 @@ def cancel_request(self, id: int): Args: id (int): The request id for which to cancel the response """ - try: - self.enqueue_lock.acquire() - self.request_queue.put( - RequestQueueItem(id, is_canceled_request=True)) - finally: - self.enqueue_lock.release() + self.executor_request_queue.enqueue_cancel_request(id) def shutdown(self): """ Signals the server to shutdown. """ - try: - self.enqueue_lock.acquire() - self.request_queue.put(RequestQueueItem(SHUTDOWN_REQUEST_ID)) - self.active = False - finally: - self.enqueue_lock.release() + self.executor_request_queue.enqueue_shutdown_request() self.shutdown_event.wait() self.worker_thread.join() self.worker_started = False @@ -428,10 +341,7 @@ def can_enqueue_requests(self) -> bool: """ Indicates if the current process is allowed to enqueue requests """ - self.enqueue_lock.acquire() - can_enqueue = self.active - self.enqueue_lock.release() - return can_enqueue and self.dist.rank == 0 + return self.executor_request_queue.can_enqueue_request() def get_latest_iteration_stats(self): """ @@ -469,20 +379,8 @@ def enqueue_request(self, """ Enqueue a new request, query is only used in `StarAttention`. """ - try: - self.enqueue_lock.acquire() - assert self.active, "PyExecutor has already been shutdown." - req_id = self.next_req_id - if self.enable_iter_perf_stats: - self.start_times[req_id] = time.time() - - if query is not None: - self.request_queue.put(RequestQueueItem(req_id, request, query)) - else: - self.request_queue.put(RequestQueueItem(req_id, request)) - self.next_req_id += 1 - finally: - self.enqueue_lock.release() + req_id = self.executor_request_queue.enqueue_request(request, query) + return req_id def set_gather_responses(self, gather_all_responses): @@ -490,8 +388,8 @@ def set_gather_responses(self, gather_all_responses): @property def should_stop_processing(self): - return self.is_shutdown and len(self.active_requests) == 0 and len( - self.waiting_queue) == 0 + return self.is_shutdown and len(self.active_requests) == 0 and \ + self.executor_request_queue.get_waiting_queue_size() == 0 @contextmanager def _profiler(self): @@ -630,7 +528,7 @@ def get_queued_req_stats(request_id: int) -> RequestStats: req_stat.stage = req.stage req_stats.append(req_stat) - for req in list(self.request_queue.queue): + for req in list(self.executor_request_queue.get_request_queue().queue): if isinstance(req, RequestQueueItem): req_stat = get_queued_req_stats(req.id) req_stat.stage = RequestStage.QUEUED @@ -647,7 +545,8 @@ def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests, scheduled_batch) -> IterationStats: stats.iter_latency_ms = iter_latency_ms - stats.num_queued_requests = self.request_queue.qsize() + stats.num_queued_requests = self.executor_request_queue.get_request_queue_size( + ) stats.num_completed_requests = num_completed_requests stats.max_num_active_requests = self.max_num_active_requests @@ -749,7 +648,7 @@ def _executor_loop_pp(self): with self._profiler() as profile_step: iter_start_time = time.time() iter_stats = None - while not self.should_stop_processing: + while True: profile_step() if self.enable_iter_perf_stats: iter_start_time = time.time() @@ -760,7 +659,8 @@ def _executor_loop_pp(self): if self.enable_iter_perf_stats: iter_stats = self._get_init_iter_stats( len(new_requests), - self.new_active_requests_queue_latency_ms) + self.executor_request_queue. + get_new_active_requests_queue_latency()) self._pad_attention_dp_dummy_request() @@ -801,6 +701,12 @@ def _executor_loop_pp(self): if self._need_return_logits(scheduled_batch): logits_host = batch_outputs["logits"].to( "cpu", non_blocking=True) + + if self.guided_decoder is not None: + self.guided_decoder.build(scheduled_batch) + self.guided_decoder.execute( + scheduled_batch, batch_outputs['logits']) + sample_state = self._sample_async( scheduled_batch, batch_outputs) sample_state.host.logits = logits_host @@ -900,7 +806,7 @@ def _executor_loop(self): sample_state = None iter_start_time = time.time() iter_stats = None - while not self.should_stop_processing: + while True: profile_step() if self.enable_iter_perf_stats: iter_start_time = time.time() @@ -914,12 +820,12 @@ def _executor_loop(self): if self.enable_iter_perf_stats: iter_stats = self._get_init_iter_stats( len(new_requests), - self.new_active_requests_queue_latency_ms) + self.executor_request_queue. + get_new_active_requests_queue_latency()) self._pad_attention_dp_dummy_request() - if self.draft_model_engine is not None or hasattr( - self, 'drafter') and self.drafter is not None: + if self.drafter is not None: self._prepare_draft_requests(self.active_requests) scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule( @@ -953,18 +859,6 @@ def _executor_loop(self): if scheduled_batch.batch_size > 0 or ( self.enable_attention_dp and self.dist.tp_size > 1): - if self.kv_cache_transceiver: - # For generation requests which have completed KV cache transfer - self._prepare_disagg_gen_transmission_complete( - scheduled_batch) - - self.resource_manager.prepare_resources(scheduled_batch) - if self.draft_model_engine is not None: - self._prepare_draft_tokens(scheduled_batch) - - if self.drafter is not None: - self.drafter.prepare_draft_tokens(scheduled_batch) - if self.kv_cache_transceiver: # For generation requests which have completed KV cache transfer self._prepare_disagg_gen_transmission_complete( @@ -973,8 +867,18 @@ def _executor_loop(self): # Return the first token to the client self._handle_first_token_response(scheduled_batch) + self.resource_manager.prepare_resources(scheduled_batch) + if self.drafter is not None: + self.drafter.prepare_draft_tokens( + scheduled_batch, self.resource_manager) + batch_outputs = self._forward_step(scheduled_batch) + if self.guided_decoder is not None: + self.guided_decoder.build(scheduled_batch) + self.guided_decoder.execute(scheduled_batch, + batch_outputs['logits']) + sample_state = self._sample_async(scheduled_batch, batch_outputs) @@ -1036,16 +940,17 @@ def _prepare_draft_requests(self, requests): def _executor_loop_overlap(self): torch.cuda.set_device(self.device_id) if self.dist.rank == 0 and not self.is_warmup and self.benchmark_req_queues_size > 0 and self.kv_cache_transceiver: - while self.request_queue.qsize() < self.benchmark_req_queues_size: + while self.executor_request_queue.get_request_queue_size( + ) < self.benchmark_req_queues_size: logger.info( - f"sleep 5 seconds, num_request_queue: {self.request_queue.qsize()}" + f"sleep 5 seconds, num_request_queue: {self.executor_request_queue.get_request_queue_size()}" ) time.sleep(5) with self._profiler() as profile_step: iter_start_time = time.time() iter_stats = None - while not self.should_stop_processing: + while True: profile_step() if self.enable_iter_perf_stats: iter_start_time = time.time() @@ -1059,7 +964,8 @@ def _executor_loop_overlap(self): if self.enable_iter_perf_stats: iter_stats = self._get_init_iter_stats( len(new_requests), - self.new_active_requests_queue_latency_ms) + self.executor_request_queue. + get_new_active_requests_queue_latency()) self._pad_attention_dp_dummy_request() @@ -1111,10 +1017,6 @@ def _executor_loop_overlap(self): ) if self.kv_cache_transceiver: - # For generation requests which have completed KV cache transfer - self._prepare_disagg_gen_transmission_complete( - scheduled_batch) - # Return the first token to the client self._handle_first_token_response(scheduled_batch) @@ -1123,6 +1025,14 @@ def _executor_loop_overlap(self): batch_outputs = self._forward_step(scheduled_batch, previous_tensors_device) + if self.previous_batch is not None: + self._update_requests(self.previous_batch.sample_state) + + if self.guided_decoder is not None: + self.guided_decoder.build(scheduled_batch) + self.guided_decoder.execute(scheduled_batch, + batch_outputs['logits']) + sample_state = self._sample_async(scheduled_batch, batch_outputs) assert sample_state is not None, "Sampling failed" @@ -1156,8 +1066,6 @@ def _executor_loop_overlap(self): self._terminate_ctx_finished_requests() def _process_previous_batch(self): - self._update_requests(self.previous_batch.sample_state) - if self.kv_cache_transceiver and self.previous_batch.ctx_transmission_reqs: for req in self.previous_batch.ctx_transmission_reqs: req.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS @@ -1185,181 +1093,17 @@ def _forward_step_inter_pp(self, scheduled_batch) -> SampleState: sampler_event=sampler_event, ) - def _update_new_active_requests_queue_latency( - self, new_requests: List[RequestQueueItem]): - if self.enable_iter_perf_stats and self.dist.rank == 0: - now = time.time() - for req_item in new_requests: - if req_item.id in self.start_times: - self.new_active_requests_queue_latency_ms += now - self.start_times.pop( - req_item.id) - - @nvtx_range("_broadcast_new_requests") - def _broadcast_new_requests( - self, - new_requests: List[RequestQueueItem], - py_request_objects: Optional[dict[str, tuple[str, dict]]] = None, - ) -> tuple[List[RequestQueueItem], Optional[dict[str, tuple[str, dict]]]]: - """Broadcasts new_requests and optional Python-only metadata (`py_request_objects`) across pipeline stages. - `py_request_objects` is a tuple of (attribute_name, {request_id: object}). - """ - payloads = (new_requests, py_request_objects) - - if not self.dist.has_pp: - return self.dist.broadcast(payloads, root=0) - - # broadcast within first tp group before send/recv chain to other tp groups - if self.dist.tp_size > 1 and self.dist.is_first_pp_rank: - payloads = self.dist.tp_broadcast(payloads, root=0) - - # tag = [0, num_micro_batches - 1] used for new_tokens send/recv - tag = self.num_micro_batches - - # send payloads - if not self.dist.is_first_pp_rank: - payloads = self.dist.recv_object(self.dist.prev_pp_rank, tag) - - if not self.dist.is_last_pp_rank: - self.dist.send_object(payloads, self.dist.next_pp_rank, tag) - - return payloads - @nvtx_range("_fetch_new_requests") def _fetch_new_requests(self) -> List[RequestQueueItem]: - if self.enable_attention_dp: - all_ranks_num_active_requests = [] - responses_list = self.dist.tp_allgather(len(self.active_requests)) - for num_active_requests in responses_list: - all_ranks_num_active_requests.append(num_active_requests) - total_num_active_requests = sum(all_ranks_num_active_requests) - total_max_num_active_requests = self.dist.tp_size * self.max_num_active_requests - else: - total_num_active_requests = len(self.active_requests) - total_max_num_active_requests = self.max_num_active_requests - - timeout = None if (total_num_active_requests == 0) and len( - self.waiting_queue) == 0 else datetime.timedelta(0) - new_requests = [] - if self.dist.rank == 0: - new_requests = _get_from_request_queue(self.request_queue, timeout) - - if self.dist.rank == 0: - py_logits_post_processors = self._collect_py_objects_from_requests( - new_requests, "py_logits_post_processors") - py_multimodal_data = self._collect_py_objects_from_requests( - new_requests, "py_multimodal_data") - py_request_objects = tuple( - filter(None, [py_logits_post_processors, py_multimodal_data])) - else: - py_request_objects = None - - if self.dist.rank == 0: - # Preserve original `new_requests` on rank 0 since it may contain - # Python-only objects (e.g., custom logits processors) not serializable by pybind. - _ = self._broadcast_new_requests(new_requests, py_request_objects) - else: - new_requests, py_request_objects = self._broadcast_new_requests( - new_requests, py_request_objects) - - # drop requests arriving after shutdown - valid_new_requests = [] - for req_item in new_requests: - if req_item.is_shutdown_request: - self.is_shutdown = True - break - elif req_item.is_canceled_request: - self.canceled_req_ids.append(req_item.id) - else: - valid_new_requests.append(req_item) - # Check if the beam width of the requests is equal to the max_beam_width - for req_item in valid_new_requests: - assert req_item.request.sampling_config.beam_width == self.max_beam_width, f"Request beam width {req_item.request.sampling_config.beam_width} is not equal to max_beam_width {self.max_beam_width}. This is not supported!" - - if py_request_objects and (self.dist.tp_size > 1 - or self.dist.has_pp) and self.dist.rank > 0: - for attr_name, req_obj_dict in py_request_objects: - self._attach_py_objects_to_requests(valid_new_requests, - attr_name, req_obj_dict) - - self.waiting_queue.extend(valid_new_requests) + new_requests = self.executor_request_queue.fetch_new_requests( + len(self.active_requests)) + self.active_requests.extend(new_requests) - new_requests = _get_from_waiting_queue( - self.waiting_queue, - total_max_num_active_requests - total_num_active_requests) - - if not self.enable_attention_dp: - self._update_new_active_requests_queue_latency(new_requests) - new_requests = self._merge_requests(new_requests) - self.active_requests.extend(new_requests) - return new_requests - - num_new_requests_all_ranks = len(new_requests) - self.expected_num_active_requests = max( - (total_num_active_requests + num_new_requests_all_ranks + - self.dist.tp_size - 1) // self.dist.tp_size, - max(all_ranks_num_active_requests), + self.is_shutdown = self.executor_request_queue.is_shutdown + self.expected_num_active_requests = self.executor_request_queue.get_expected_num_active_requests( ) - self.has_context_request = False - new_requests_cur_rank = [] - if new_requests != [] and self.expected_num_active_requests > all_ranks_num_active_requests[ - self.dist.tp_rank]: - # Balance context tokens across ranks - HeapVal = namedtuple( - 'HeapVal', - [ - 'num_tokens', # number of context tokens that have been added - 'num_requests', # number of requests to be added - 'rank', # rank - 'request_list', # new requests that have been added - ], - ) - all_ranks_new_requests_heap = [ - HeapVal(0, self.expected_num_active_requests - val, tp_rank, []) - for tp_rank, val in enumerate(all_ranks_num_active_requests) - ] - new_requests_cur_rank = all_ranks_new_requests_heap[ - self.dist.tp_rank].request_list - all_ranks_new_requests_heap = [ - val for val in all_ranks_new_requests_heap - if val.num_requests > 0 - ] - heapq.heapify(all_ranks_new_requests_heap) - new_requests = sorted(new_requests, - key=lambda x: len(x.request.input_token_ids), - reverse=True) - for req_item in new_requests: - val = heapq.heappop(all_ranks_new_requests_heap) - val = val._replace( - num_tokens=val.num_tokens + - len(req_item.request.input_token_ids), - num_requests=val.num_requests - 1, - ) - val.request_list.append(req_item) - if val.num_requests > 0: - heapq.heappush(all_ranks_new_requests_heap, val) - elif val.rank == self.dist.tp_rank: - break - - # In disaggregated serving, we might get either context request or - # generation request. In IFB, we only get context request from request queue - if self.kv_cache_transceiver: - for req_item in new_requests_cur_rank: - if req_item.request.request_type == RequestType.REQUEST_TYPE_CONTEXT_ONLY: - self.has_context_request = True - break - else: - self.has_context_request = len(new_requests_cur_rank) > 0 - self._update_new_active_requests_queue_latency( - new_requests_cur_rank) - - self.num_fetch_requests = self.num_fetch_requests + num_new_requests_all_ranks - self.num_fetch_requests_cur_rank = self.num_fetch_requests_cur_rank + len( - new_requests_cur_rank) - - new_requests_cur_rank = self._merge_requests(new_requests_cur_rank) - self.active_requests.extend(new_requests_cur_rank) - return new_requests_cur_rank + return new_requests def _add_kv_cache_events(self): kv_cache_manager = self.resource_manager.resource_managers.get( @@ -1370,149 +1114,6 @@ def _add_kv_cache_events(self): # to be transferred to main thread when user needs them. kv_cache_manager.flush_iteration_events() - def _collect_py_objects_from_requests( - self, requests: list[RequestQueueItem], - attribute_name: str) -> Optional[tuple[str, dict]]: - """WAR to gather dynamic Python-only attributes (e.g., custom logits processors) - that cannot be handled by pybind serialization during MP communication. - - Returns: - A tuple of (attribute_name, {request_id: object}) or None. - """ - req_id_to_obj = {} - for item in requests: - if not item.is_normal_request: - continue - obj = getattr(item.request, attribute_name, None) - if obj is not None: - req_id_to_obj[item.id] = obj - return None if not req_id_to_obj else (attribute_name, req_id_to_obj) - - def _attach_py_objects_to_requests(self, requests: list[RequestQueueItem], - attribute_name: str, - py_request_objects: dict): - """Attaches Python-only objects (e.g., dynamic attributes not handled by pybind) - to each request. - """ - for item in requests: - py_obj = py_request_objects.get(item.id) - if py_obj is not None: - setattr(item.request, attribute_name, py_obj) - - def _partition_context(self, ctx_ids_list): - ctx_ids = torch.tensor(ctx_ids_list).unsqueeze(0) - ctx_len = ctx_ids.shape[-1] - block_size = self.dist.cp_config['block_size'] - if block_size is None: - block_size = ctx_len // self.dist.cp_size - anchor_block_size = self.dist.cp_config['cp_anchor_size'] - if anchor_block_size is None: - anchor_block_size = block_size - - assert anchor_block_size <= block_size, f'cp_anchor_size {anchor_block_size} should be smaller than block_size {block_size}' - padding = 0 - if ctx_len % block_size != 0: - padding = block_size - (ctx_len % block_size) - assert padding <= ctx_len, f'block size is too large for context, please set it smaller' - ctx_ids = torch.cat( - (ctx_ids, torch.zeros_like(ctx_ids)[:, :padding]), dim=-1) - position_ids = torch.arange(0, ctx_ids.shape[-1]).unsqueeze(0) - - ctx_ids_blocks = torch.tensor_split( - torch.stack(ctx_ids.split(block_size, dim=-1)), self.dist.cp_size) - position_ids_blocks = torch.tensor_split( - torch.stack(position_ids.split(block_size, dim=-1)), - self.dist.cp_size) - if self.dist.cp_rank != 0: - ctx_blocks, position_blocks = [ - ctx_ids_blocks[0][0].tolist()[0][:anchor_block_size] - ], [position_ids_blocks[0][0].tolist()[0][:anchor_block_size]] - else: - ctx_blocks, position_blocks = [], [] - - for idx in range(len(ctx_ids_blocks[self.dist.cp_rank])): - ctx_block = ctx_ids_blocks[self.dist.cp_rank][idx] - position_block = position_ids_blocks[self.dist.cp_rank][idx] - ctx_blocks.append(ctx_block.tolist()[0]) - position_blocks.append(position_block.tolist()[0]) - return ctx_blocks, position_blocks, padding - - def _merge_star_attention_requests(self, - new_requests: list[RequestQueueItem]): - result = [] - for req_item in new_requests: - req_id, exe_req, query_token_ids = req_item.id, req_item.request, req_item.query - ctx_len0 = len(exe_req.input_token_ids) - ctx_blocks, position_blocks, last_block_padding_num = [ - exe_req.input_token_ids - ], [[i for i in range(ctx_len0)]], 0 - ctx_blocks, position_blocks, last_block_padding_num = self._partition_context( - exe_req.input_token_ids) - if self.dist.cp_rank == self.dist.cp_size - 1 and last_block_padding_num > 0: - ctx_blocks[-1] = ctx_blocks[-1][:-last_block_padding_num] - position_blocks[-1] = position_blocks[ - -1][:-last_block_padding_num] - #if has query - if query_token_ids: - ctx_blocks.append(query_token_ids) - position_blocks.append([ - i for i in range(ctx_len0, ctx_len0 + len(query_token_ids)) - ]) - - # insert the dummy block to align the number of ctx iterations of each rank - block_size = self.dist.cp_config['block_size'] - total_blocks = (ctx_len0 + block_size - 1) // block_size - num_blocks_per_rank = ( - total_blocks + self.dist.cp_size - - 1) // self.dist.cp_size + 1 # 1 for query block - if len(ctx_blocks) == num_blocks_per_rank: - ctx_blocks.insert(1, []) - position_blocks.insert(1, []) - elif len(ctx_blocks) == num_blocks_per_rank + 1: - # anchor + ctx_blocks + qry_block - pass - else: - print( - f'rank = {self.dist.cp_rank}, len(ctx_blocks) = {len(ctx_blocks) }, num_blocks_per_rank = {num_blocks_per_rank}' - ) - assert False, f'invalid context partition' - - # fake data for scheduler - ctx_blocks_list = [0] * (block_size + - self.dist.cp_config['cp_anchor_size']) - - req = executor_request_to_llm_request( - req_id, exe_req, self._should_exclude_last_generation_logits(), - ctx_blocks_list) - req.gen_iters = 0 - req.ctx_iters = 0 - req.ctx_blocks = ctx_blocks - req.ctx_position_blocks = position_blocks - req.query_id = query_token_ids - - result.append(req) - - return result - - @nvtx_range("_merge_requests") - def _merge_requests(self, new_requests: list[RequestQueueItem]): - cp_config = self.dist.cp_config - if 'cp_type' in cp_config: - cp_type = cp_config['cp_type'] - if cp_type == 'star_attention': - return self._merge_star_attention_requests(new_requests) - elif cp_type == 'ring_attention': - raise NotImplementedError("ring attention not implemented yet") - else: - raise NotImplementedError(f'unsupport cp type {cp_type}') - else: - return [ - executor_request_to_llm_request( - req_item.id, req_item.request, - self._should_exclude_last_generation_logits()) - for req_item in new_requests - ] - @nvtx_range("_schedule") def _schedule(self): scheduler_output = self.scheduler.schedule_request( @@ -1545,7 +1146,7 @@ def _check_disagg_gen_transfer_status(self): @nvtx_range("_pad_attention_dp_dummy_request") def _pad_attention_dp_dummy_request(self): """ - Pad with a dummy request, if required, to ensure every attention_dp rank has at least one active request. + Pad with a generation dummy request, if required, to ensure every attention_dp rank has at least one active request. """ if not self.enable_attention_dp: return @@ -1563,8 +1164,8 @@ def _pad_attention_dp_dummy_request(self): if self.expected_num_active_requests - num_active_request > 0 and num_active_request == 0: llm_request = self.kv_cache_manager.add_dummy_requests( request_ids=[0], - is_gen=not self.has_context_request, - prepare_resource=not self.has_context_request, + is_gen=True, + prepare_resource=True, max_num_draft_tokens=self.max_draft_len, )[0] llm_request.is_attention_dp_dummy = True @@ -1652,8 +1253,13 @@ def _send_disagg_ctx_cache(self, scheduled_ctx_requests): if req.is_context_only_request and (req.is_context_finished or req.is_finished_due_to_length): self.kv_cache_transceiver.respond_and_send_async(req) - self.resource_manager.resource_managers[ - ResourceManagerType.SEQ_SLOT_MANAGER].free_resources(req) + for resource_mgr_type in ( + ResourceManagerType.SEQ_SLOT_MANAGER, + ResourceManagerType.SPEC_RESOURCE_MANAGER): + if resource_mgr_type in self.resource_manager.resource_managers and self.resource_manager.resource_managers[ + resource_mgr_type] is not None: + self.resource_manager.resource_managers[ + resource_mgr_type].free_resources(req) self.kv_cache_transceiver.check_context_transfer_status(0) @@ -1710,6 +1316,10 @@ def _update_request_states_tp(self, scheduled_requests: ScheduledRequests): for request in scheduled_requests.context_requests: if request.state != LlmRequestState.GENERATION_COMPLETE: # skip failed requests + request.py_last_context_chunk = ( + request.context_current_position, + request.context_current_position + + request.context_chunk_size) request.move_to_next_context_chunk() if request.context_remaining_length == 0: request.state = LlmRequestState.GENERATION_IN_PROGRESS @@ -1768,188 +1378,6 @@ def _update_requests(self, sample_state: SampleState): logger.error(f"Encountered an error in sampling: {error_msg}") self._handle_errors(error_msg) - @nvtx_range("_prepare_draft_batch") - def _prepare_draft_batch( - self, scheduled_requests: ScheduledRequests - ) -> Tuple[ScheduledRequests, Dict[int, LlmRequest]]: - """ - Prepares a batch for the draft model engine. Draft tokens are only produced - for generation requests. - - The requests are prepared as follows: - 1. The first time the draft engine sees a request, it's a context request. - 2. Otherwise, if draft tokens were accepted on the last target model decoding - step, it's a chunked context request (we process all the accepted tokens together). - 3. Otherwise, it's a generation request. - """ - try: - draft_batch = ScheduledRequests() - - for request in scheduled_requests.generation_requests: - if request.py_draft_pages_allocated == 0: - # No space for draft tokens. - continue - - # Stop drafting when we hit the max seqlen. We still need dummy draft - # tokens attached to the requests to make sure everything works properly - # with CUDA graph. These dummy tokens are already added by - # _prepare_draft_requests to make the KV cache/scheduler aware of the fact - # that we want to do spec decoding, so no need to do anything else here. - # This makes the perf for this case suboptimal, but that's OK - this is - # a corner case for weird models like the llama 3.1 8b EAGLE3 implementation. - if request.max_beam_num_tokens - 1 >= self.draft_model_engine.max_seq_len: - continue - - num_draft_tokens = len( - request.py_last_draft_tokens - ) if request.py_last_draft_tokens is not None else 0 - request.py_draft_tokens = [] - - num_accepted_tokens = request.py_num_accepted_draft_tokens - num_rejected_tokens = num_draft_tokens - num_accepted_tokens - assert num_rejected_tokens >= 0 - - spec_config = self.model_engine.spec_config - beam_idx = 0 - input_tokens = spec_config.get_draft_model_prompt( - request.get_tokens()[beam_idx]) - - def create_new_request(input_tokens): - return LlmRequest( - request_id=request.py_request_id, - max_new_tokens=request.py_max_new_tokens, - input_tokens=input_tokens, - sampling_config=request.sampling_config, - return_perf_metrics=request.return_perf_metrics, - is_streaming=False, - is_draft=True) - - if request.max_beam_num_tokens - 1 == request.py_prompt_len: - # This is the first time the draft model is seeing this request. - # Prepare a context request. We discard the first token and take - # the newly decoded one - this is the convention for EAGLE 2 and 3. - new_request = create_new_request(input_tokens) - draft_batch.context_requests.append(new_request) - elif num_accepted_tokens == 0: - new_request = create_new_request(input_tokens[:-1]) - # Explicitly add the last token so get_last_tokens() returns - # the right value - new_request.add_new_token(input_tokens[-1], beam_idx) - new_request.state = LlmRequestState.GENERATION_IN_PROGRESS - draft_batch.generation_requests.append(new_request) - else: - new_request = create_new_request(input_tokens) - new_request.context_chunk_size = num_accepted_tokens + 1 - new_request.context_current_position = len( - input_tokens) - num_accepted_tokens - 1 - new_request.context_chunk_size = num_accepted_tokens + 1 - new_request.context_current_position = len( - input_tokens) - num_accepted_tokens - 1 - - draft_batch.context_requests.append(new_request) - - new_request.py_stop_words_list = request.py_stop_words_list - - return draft_batch - - except Exception as e: - traceback.print_exc() - error_msg = str(e) - logger.error(f"Encountered an error in decode: {error_msg}") - self._handle_errors(error_msg) - - @nvtx_range("_prepare_draft_tokens") - def _prepare_draft_tokens(self, scheduled_requests: ScheduledRequests): - if not self.draft_model_engine: - raise ValueError("Draft model engine is not set") - - try: - draft_batch = self._prepare_draft_batch(scheduled_requests) - - if draft_batch.batch_size == 0: - return - self.draft_seq_slot_manager.prepare_resources(draft_batch) - - req_id_to_old_request = { - req.py_request_id: req - for req in scheduled_requests.all_requests() - } - - # Disable cuda graph for the 1st draft model forward - if self.model_engine.spec_config.spec_dec_mode.needs_kv_cache_recompute( - ): - with self.draft_model_engine.no_cuda_graph(): - outputs = self.draft_model_engine.forward( - draft_batch, self.resource_manager) - else: - outputs = self.draft_model_engine.forward( - draft_batch, self.resource_manager) - if hasattr(self.draft_model_engine.model.model, 'd2t'): - outputs['d2t'] = self.draft_model_engine.model.model.d2t.data - - sample_state = self._sample_async(draft_batch, outputs) - previous_batch = sample_state - - self._update_request_states(draft_batch) - - def _process_decoded_tokens(draft_batch): - new_requests = [] - for req in draft_batch.all_requests(): - target_model_req = req_id_to_old_request[req.py_request_id] - target_model_req.py_draft_tokens.append( - req.get_last_tokens(0)) - if req.state != LlmRequestState.GENERATION_COMPLETE and len( - target_model_req.py_draft_tokens - ) < target_model_req.py_draft_pages_allocated: - new_requests.append(req) - else: - self.draft_seq_slot_manager.free_resources(req) - - return new_requests - - # The TRTLLM attention kernels cannot handle generation requests with - # different seqlens. No issues with flashinfer, should we look into removing - # this? Just needs proper kernel support. - def _pad_to_max_draft_tokens(): - for req in scheduled_requests.generation_requests: - max_draft_len = self.max_draft_len - num_draft_tokens = len(req.py_draft_tokens) - req.py_draft_tokens.extend( - 0 for _ in range(max_draft_len - num_draft_tokens)) - - draft_batch.generation_requests = draft_batch.context_requests + draft_batch.generation_requests - draft_batch.context_requests = [] - - for i in range(self.max_draft_len - 1): - if len(draft_batch.generation_requests) == 0: - break - - outputs = self.draft_model_engine.forward( - draft_batch, - self.resource_manager, - new_tensors_device=previous_batch.device) - - if hasattr(self.draft_model_engine.model.model, 'd2t'): - outputs[ - 'd2t'] = self.draft_model_engine.model.model.d2t.data - sample_state = self._sample_async(draft_batch, outputs) - self._update_request_states(draft_batch) - self._update_requests(previous_batch) - new_requests = _process_decoded_tokens( - previous_batch.scheduled_requests) - draft_batch.generation_requests = new_requests - previous_batch = sample_state - self._update_requests(previous_batch) - new_requests = _process_decoded_tokens( - previous_batch.scheduled_requests) - _pad_to_max_draft_tokens() - - except Exception as e: - traceback.print_exc() - error_msg = str(e) - logger.error(f"Encountered an error in decode: {error_msg}") - self._handle_errors(error_msg) - def _handle_errors(self, error_msg: Optional[str] = None): error_responses = {} error_msg = error_msg or "error" @@ -1969,16 +1397,15 @@ def _terminate_request(self, request: LlmRequest): @nvtx_range("_handle_canceled_requests") def _handle_canceled_requests(self): - if len(self.canceled_req_ids) == 0: + if self.executor_request_queue.get_canceled_req_ids_size() == 0: return - # cancel request in the waiting queue - self.waiting_queue = deque(req for req in self.waiting_queue - if req.id not in self.canceled_req_ids) + # Remove cancel request in the waiting queue + self.executor_request_queue.update_waiting_queue() for request in self.active_requests: req_id = request.py_request_id - if req_id in self.canceled_req_ids: + if req_id in self.executor_request_queue.get_canceled_req_ids(): # Mark requests as finished, then, we reuse all existing code # to clean up the KV cache resources. request.finish_by_reason(FinishReason.CANCELLED) @@ -1988,7 +1415,7 @@ def _handle_canceled_requests(self): # TODO: revisit the cancel logic of attention dp # When enable attention dp, each rank does not have full copy of requests # so we need to remove the cancel requests not in the local rank - self.canceled_req_ids.clear() + self.executor_request_queue.clear_canceled_req_ids() @nvtx_range("_enqueue_responses") def _enqueue_responses(self, responses: Dict[int, LlmResponse]): @@ -2080,7 +1507,8 @@ def _handle_responses(self): requests_to_terminate.append(request) else: new_active_requests.append(request) - self.active_requests = new_active_requests + self.active_requests.clear() + self.active_requests.extend(new_active_requests) self._enqueue_responses(new_responses) for request in requests_to_terminate: self._terminate_request(request) @@ -2140,19 +1568,3 @@ def _remove_inflight_ids(self, scheduled_requests): """Remove reqids of current requests from self.inflight_req_ids.""" for req in scheduled_requests.all_requests(): self.inflight_req_ids.erase(req.request_id) - - def _should_exclude_last_generation_logits(self) -> bool: - # When overlap scheduler is enabled then when starting to handle a new prompt, - # sample_async is called twice before the first call to update_requests: - # - 1st time as a context request that handles on the 1st generated token - # - 2nd time as a generation request that handles on the 2nd generated token. - # and only after these two calls the sampler's update_request method is called. - # So in a sampler that works by the expected flow of handling the logits in - # sample_async (TorchSampler is an anomaly that instead does that on - # update_requests), every update_request doesn't handle the newest token, but one - # before it. Since all these calls work on the same request object, then its - # logits storage contains the logits of both the token update_requests should work - # on, and also its next token. Thus, excluding the last generation logits from any - # getter is required, when not using TorchSampler. - return not self.disable_overlap_scheduler and not isinstance( - self.sampler, TorchSampler) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index b6893d69e268..674a85741be8 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -19,11 +19,13 @@ from ..attention_backend.interface import AttentionRuntimeFeatures from ..distributed import MPIDist -from ..speculative import get_spec_drafter, get_spec_resource_manager +from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter, + get_spec_resource_manager) from ._util import (KvCacheCreator, _adjust_torch_mem_fraction, create_py_executor_instance, instantiate_sampler, is_mla) from .config import PyTorchConfig from .config_utils import is_mla +from .guided_decoder import GuidedDecoder from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor @@ -161,21 +163,6 @@ def _mangle_executor_config(executor_config: ExecutorConfig): ) executor_config.kv_cache_config.enable_block_reuse = False - spec_config = executor_config.speculative_config - if spec_config is not None and spec_config.spec_dec_mode.has_draft_model(): - # The draft and target models have different KV cache managers to support - # different head sizes, dtypes, etc in the generic case. - # However, this line will set context_current_position > 0 if there are - # cached blocks: https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/pyexecutor/resource_manager.py#L310. - # It actually mutates the LLM request! As a result, when we try to allocate KV cache - # pages for the draft model, is_first_context_chunk returns False and - # no pages are allocated. - # We need to refactor LLMRequest to fix this. Disable block reuse for now. - logger.warning( - f"Disabling block reuse for speculation algorithm {spec_config.spec_dec_mode}" - ) - executor_config.kv_cache_config.enable_block_reuse = False - if pytorch_backend_config.attn_backend == "FLASHINFER_STAR_ATTENTION" and executor_config.enable_chunked_context: logger.warning( f"Disabling chunked context for {pytorch_backend_config.attn_backend} backend" @@ -237,8 +224,8 @@ def create_py_executor( attn_runtime_features=attn_runtime_features, dist=dist, spec_config=spec_config, - guided_decoding_config=executor_config.guided_decoding_config, lora_config=lora_config, + checkpoint_loader=executor_config.checkpoint_loader, ) if has_draft_model_engine: @@ -262,6 +249,7 @@ def create_py_executor( attn_runtime_features=attn_runtime_features, dist=dist, spec_config=draft_spec_config, + checkpoint_loader=executor_config.checkpoint_loader, is_draft_model=True, ) draft_model_engine.kv_cache_manager_key = ResourceManagerType.DRAFT_KV_CACHE_MANAGER @@ -279,7 +267,7 @@ def create_py_executor( max_seq_len += spec_config.max_draft_len if spec_config is not None: - max_seq_len += spec_config.num_extra_kv_tokens + max_seq_len += get_num_extra_kv_tokens(spec_config) max_seq_len += spec_config.max_draft_len executor_config.max_seq_len = max_seq_len @@ -342,6 +330,17 @@ def create_py_executor( sampler = instantiate_sampler(model_engine, executor_config, pytorch_backend_config, mapping) + guided_decoder: Optional[GuidedDecoder] = None + if executor_config.guided_decoding_config is not None: + if spec_config is not None: + raise ValueError( + "Guided decoding is not supported with speculative decoding.") + if mapping.is_last_pp_rank(): + guided_decoder = GuidedDecoder( + executor_config.guided_decoding_config, + executor_config.max_batch_size, + model_engine.model.vocab_size_padded) + resources = {} estimating_kv_cache = False kv_cache_creator = None @@ -358,18 +357,20 @@ def create_py_executor( if estimating_kv_cache else _ExecutorCreationStage.KV_CACHE): kv_cache_creator.build_managers(resources) - # Drafter for speculative decoding - with mem_monitor.observe_creation_stage(_ExecutorCreationStage.DRAFTER): - drafter = get_spec_drafter(model_engine) - # Resource managers for speculative decoding + # For user-specified drafters, use extra_resource_managers in PyTorchBackend config + # to provide a resource manager if required. spec_resource_manager = get_spec_resource_manager(model_engine, - draft_model_engine, - drafter) + draft_model_engine) if spec_resource_manager is not None: resources[ ResourceManagerType.SPEC_RESOURCE_MANAGER] = spec_resource_manager + # Drafter for speculative decoding + with mem_monitor.observe_creation_stage(_ExecutorCreationStage.DRAFTER): + drafter = get_spec_drafter(model_engine, draft_model_engine, sampler, + spec_resource_manager) + with mem_monitor.observe_creation_stage( _ExecutorCreationStage.INIT_EXTRA_RESOURCES if estimating_kv_cache else _ExecutorCreationStage.EXTRA_RESOURCES): @@ -385,6 +386,7 @@ def create_py_executor( start_worker=False, sampler=sampler, drafter=drafter, + guided_decoder=guided_decoder, lora_config=lora_config, garbage_collection_gen0_threshold=garbage_collection_gen0_threshold, ) @@ -427,6 +429,7 @@ def create_py_executor( start_worker=False, sampler=sampler, drafter=drafter, + guided_decoder=guided_decoder, lora_config=lora_config, garbage_collection_gen0_threshold= garbage_collection_gen0_threshold, diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index e52096727c6b..adcae974354e 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -1,8 +1,9 @@ +import copy import enum import math from abc import ABC, abstractmethod from collections import OrderedDict, defaultdict -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Optional, Set, Tuple, Union import torch @@ -11,7 +12,7 @@ from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE from tensorrt_llm.sampling_params import SamplingParams -from ..._utils import nvtx_range +from ..._utils import binding_dtype_size, nvtx_range from ...logger import logger from ...mapping import Mapping from .llm_request import LlmRequest, LlmRequestState, SamplingConfig @@ -175,7 +176,9 @@ def __init__( self.kv_factor = 1 if kv_cache_type == CacheTypeCpp.SELFKONLY else 2 # Some speculative decoding methods need to use different kv lengths for the # draft/target layers. Add extra tokens to handle this issue. - self.num_extra_kv_tokens = 0 if spec_config is None else spec_config.num_extra_kv_tokens + # Import here to avoid circular imports + from ..speculative import get_num_extra_kv_tokens + self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config) self.event_buffer_max_size = kv_cache_config.event_buffer_max_size self.max_num_tokens = max_num_tokens @@ -192,10 +195,10 @@ def __init__( else 0) # Determine if this is VSWA (Variable Sliding Window Attention) - is_vswa = len(self.max_attention_window_vec) > 1 + self.is_vswa = len(self.max_attention_window_vec) > 1 # Calculate blocks per window using appropriate method - if is_vswa: + if self.is_vswa: # VSWA case: use C++ implementation for variable window sizes # model config check if model_config is None: @@ -372,11 +375,15 @@ def add_dummy_requests( prepare_resource: bool = True, max_num_draft_tokens: int = 0, use_mrope: bool = False, + max_beam_width: int = 1, ): - beam_width = 1 # TODO: more than 1 beam? + beam_width = max_beam_width requests = [] for i, req_id in enumerate(request_ids): - sampling_params = SamplingParams() + # exact choice of n can be ignored for dummy requests + sampling_params = SamplingParams(n=beam_width, + best_of=beam_width, + use_beam_search=beam_width > 1) # Here 1+max_num_draft_tokens is used to extend the prompt length to # a non-zero number to skip illegal memory access issue in MLA kernel # during warmup. @@ -437,14 +444,10 @@ def calculate_max_num_blocks(self, cache_size_per_token = kv_factor * sum( self.num_kv_heads_per_layer) * head_dim - if dtype == DataType.FP8: - kv_cache_dtype_bytes = 1 - elif dtype in (DataType.HALF, DataType.BF16): - kv_cache_dtype_bytes = 2 - elif dtype == DataType.FLOAT: - kv_cache_dtype_bytes = 4 - else: + if dtype not in (DataType.FP8, DataType.HALF, DataType.BF16, + DataType.FLOAT): raise ValueError(f'Cannot support {dtype} KV cache.') + kv_cache_dtype_bytes = binding_dtype_size(dtype) cache_size_bytes_per_token = cache_size_per_token * kv_cache_dtype_bytes free_mem, total_mem = torch.cuda.mem_get_info() @@ -526,7 +529,14 @@ def get_batch_cache_indices( return result def get_num_free_blocks(self) -> int: - return self.impl.get_kv_cache_stats().free_num_blocks + if self.is_vswa: + logger.info( + f"For VSWA case, we return the minimum of the number of free blocks for each window size: {self.impl.get_kv_cache_stats().num_free_blocks_per_window_size}" + ) + return min(self.impl.get_kv_cache_stats(). + num_free_blocks_per_window_size.values()) + else: + return self.impl.get_kv_cache_stats().free_num_blocks def get_num_kv_blocks(self, num_tokens: int) -> int: return (num_tokens + self.tokens_per_block - 1) // self.tokens_per_block @@ -603,6 +613,102 @@ def _get_window_size_to_layers(self) -> dict[int, list[int]]: window_size_to_layers_map[window_size].append(local_layer_idx) return window_size_to_layers_map + @staticmethod + def adjust_window_sizes_for_vswa( + window_size_to_layers: Dict[int, List[int]], + kv_cache_config: KvCacheConfigCpp, + model_config: ModelConfig, + pool_memory_bytes: int, + kv_factor: int, + dtype: DataType, + is_cross_attention: bool = False, + ) -> Dict[int, List[int]]: + + assert is_cross_attention is False, 'Cross attention is not supported' + + max_tokens_from_config = kv_cache_config.max_tokens + + def calculate_cache_size_per_token(layers: Set[int]) -> int: + # Same as BaseKVCacheManager::calculateCacheSizePerTokenForSingleWindowSize + total_kv_heads = sum(model_config.num_kv_heads_per_layer[i] + for i in layers) + return total_kv_heads * kv_factor * model_config.head_size + + # Calculate the required memory bytes per sequence. + required_mem_bytes_per_seq = 0 + for window_size in sorted(window_size_to_layers): + layers = window_size_to_layers[window_size] + cache_size_per_token = calculate_cache_size_per_token(layers) + cache_size_bytes_per_token = cache_size_per_token * binding_dtype_size( + dtype) + required_mem_bytes_per_seq += window_size * cache_size_bytes_per_token + logger.debug( + f'Required memory per sequence: {required_mem_bytes_per_seq} bytes') + + if required_mem_bytes_per_seq < pool_memory_bytes: + # No need to adjust the window sizes. + return copy.deepcopy(window_size_to_layers) + + logger.debug( + f'Adjusting the window sizes {list(window_size_to_layers)} to fit ' + f'the memory {pool_memory_bytes} bytes.') + adjusted_window_size_to_layers = {} + + remaining_mem_bytes = pool_memory_bytes + remaining_layers = set(i for layers in window_size_to_layers.values() + for i in layers) + + accum_max_tokens = 0 + prev_window_size = 0 + + for window_size in sorted(window_size_to_layers): + layers = window_size_to_layers[window_size] + if remaining_mem_bytes > 0 and remaining_layers: + # Calculate cache size per token for remaining layers only + cache_size_per_token = calculate_cache_size_per_token( + remaining_layers) + cache_size_bytes_per_token = cache_size_per_token * binding_dtype_size( + dtype) + logger.debug( + f'Cache size per token for {len(remaining_layers)} layers: ' + f'{cache_size_bytes_per_token} bytes') + # Calculate max tokens that can fit in this window with remaining memory. + max_tokens_in_window = min( + remaining_mem_bytes // cache_size_bytes_per_token, + window_size - prev_window_size) + remaining_mem_bytes -= max_tokens_in_window * cache_size_bytes_per_token + accum_max_tokens += max_tokens_in_window + logger.debug(f'Remaining memory: {remaining_mem_bytes} bytes') + logger.debug( + f'Max token of window {window_size}: {accum_max_tokens}') + + if accum_max_tokens < window_size: + logger.debug( + f'Max tokens ({accum_max_tokens}) cannot fill the current window ({window_size}). ' + f'The larger windows will have the same max tokens.') + remaining_mem_bytes = 0 + + # Clamp the sequence length if provided explicitly. + if max_tokens_from_config is not None: + accum_max_tokens = min(max_tokens_from_config, + accum_max_tokens) + # If max tokens from config is reached, stop allocating + # more memory. Since the maximum number of tokens is + # already reached, for the remaining windows maxTokens + # will be set by the current value of accumMaxTokens. + if accum_max_tokens == max_tokens_from_config: + remaining_mem_bytes = 0 + + if accum_max_tokens not in adjusted_window_size_to_layers: + adjusted_window_size_to_layers[accum_max_tokens] = layers.copy() + else: + adjusted_window_size_to_layers[accum_max_tokens].extend(layers) + + remaining_layers -= set(layers) + prev_window_size = window_size + + return adjusted_window_size_to_layers + def calculate_max_num_blocks_from_cpp( self, kv_cache_config: KvCacheConfigCpp, @@ -622,6 +728,11 @@ def calculate_max_num_blocks_from_cpp( A dict of (max_attention_window, (blocks_in_primary_pool, blocks_in_secondary_pool)). """ + # VSWA on Torch backend has not supported the cross attention. + is_cross_attention = False + # check model config + assert model_config.layer_types is not None, "layer_types have to be set correctly for VSWA" + # Construct WorldConfig from self.mapping world_config_cpp = WorldConfig( tensor_parallelism=self.mapping.tp_size, @@ -636,12 +747,26 @@ def calculate_max_num_blocks_from_cpp( primary_pool_memory_bytes = free_mem secondary_pool_memory_bytes = 0 logger.debug( - f"primary_pool_memory_bytes is set to {primary_pool_memory_bytes/1024**3}GB, \nsecondary_pool_memory_bytes is set to {secondary_pool_memory_bytes/1024**3}GB" + f"primary_pool_memory_bytes is set to {primary_pool_memory_bytes/1024**3}GB, \n" + f"secondary_pool_memory_bytes is set to {secondary_pool_memory_bytes/1024**3}GB" + ) + + # Adjust the window sizes to fit the memory if even a single sequence + # cannot fit in the memory. + window_size_to_layers = self.adjust_window_sizes_for_vswa( + window_size_to_layers=window_size_to_layers, + model_config=model_config, + kv_cache_config=kv_cache_config, + pool_memory_bytes=primary_pool_memory_bytes, + kv_factor=self.kv_factor, + dtype=self.dtype, + is_cross_attention=is_cross_attention, ) blocks_per_window = KVCacheManagerCpp.calculate_max_num_blocks( config=kv_cache_config, - is_cross_attention=False, #TODO: support cross attention + # TODO: support cross attention + is_cross_attention=is_cross_attention, dtype=self.dtype, model_config=model_config, world_config=world_config_cpp, @@ -1099,7 +1224,7 @@ def update_resources(self, scheduled_batch: ScheduledRequests): pass def free_resources(self, request: LlmRequest): - pass + self.impl.mark_request_done(request) def shutdown(self): pass diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index b4dfdf25d45b..f6f4a7420dda 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -194,7 +194,7 @@ def add_token(request: LlmRequest, *, beam: int, step: int = 0) -> int: - seq_slot = request.seq_slot + seq_slot = request.py_seq_slot assert seq_slot is not None new_token = int(new_tokens[step, seq_slot, beam]) request.add_new_token(new_token, beam) @@ -285,14 +285,14 @@ def _handle_stop_criteria(self, request: LlmRequest, def handle_logits(self, request: LlmRequest, state: SampleState, *, beam: int, count: int): - current_slice = slice(0, count), request.seq_slot, beam + current_slice = slice(0, count), request.py_seq_slot, beam if request.py_return_generation_logits: assert state.host.logits is not None current_logits = state.host.logits[current_slice] request.py_result.append_generation_logits(current_logits) if request.py_return_log_probs: assert state.host.log_probs is not None - log_probs = state.host.log_probs[request.seq_slot][beam][:count] + log_probs = state.host.log_probs[request.py_seq_slot][beam][:count] current_tokens = state.host.new_tokens[current_slice] token_log_probs = [{ @@ -406,7 +406,7 @@ def _process_requests(self, no_draft_tokens = len(requests) == sum_steps fast_path = not self.enable_mixed_sampler and no_draft_tokens and gen_logits_host is None and log_probs_host is None - seq_slots = torch.as_tensor([r.seq_slot for r in requests]) + seq_slots = torch.as_tensor([r.py_seq_slot for r in requests]) seq_slots = seq_slots.to(device="cuda", non_blocking=True) if fast_path: @@ -473,10 +473,12 @@ class SampleStateTensorsHostTRTLLM(SampleStateTensors): finish_reasons: torch.Tensor sequence_lengths: torch.Tensor cum_log_probs: torch.Tensor | None = None + gathered_ids: torch.Tensor | None = None @dataclass(kw_only=True) class SampleStateTRTLLM(SampleState): + finalize_events: dict[str, CudaEvent] host: SampleStateTensorsHostTRTLLM @@ -536,8 +538,7 @@ def _initialize_store(self): "buffer_manager": buffer_manager, "decoder_input_buffers": [ - DecoderInputBuffers(self.max_num_sequences, - self.executor_config.max_batch_size, + DecoderInputBuffers(self.executor_config.max_batch_size, self.MAX_DECODING_TOKENS, buffer_manager) for _ in range(self.num_micro_batches) ], @@ -617,9 +618,9 @@ def _update_cache_indirection_buffer(self, # Copy cache indirection output to input for request in scheduled_requests.generation_requests: self.store["decoder_state"].cache_indirection_input[ - request.seq_slot].copy_( + request.py_seq_slot].copy_( self.store["decoder_state"].cache_indirection_output[ - request.seq_slot], + request.py_seq_slot], non_blocking=True) @torch.inference_mode() @@ -673,6 +674,24 @@ def sample_async(self, scheduled_requests: ScheduledRequests, self.store["decoder_state"], self.store["decoding_input"][self.micro_batch_idx]) + finalize_events = {} + gathered_ids = None + if beam_width > 1: + finished_sum_device = self.store["decoder_state"].finished_sum + + for request in scheduled_requests.all_requests(): + if request.is_context_init_state: + continue + if finished_sum_device[request.seq_slot] == beam_width: + finalize_events[ + request.request_id] = self._finalize_request( + request, False) + elif request.streaming: + finalize_events[ + request.request_id] = self._finalize_request( + request, True) + gathered_ids = self.store["decoder_state"].gathered_ids.to( + 'cpu', non_blocking=True) new_output_tokens = self.store["decoder_state"].all_new_tokens.to( 'cpu', non_blocking=True) finished_sum = self.store["decoder_state"].finished_sum.to( @@ -699,7 +718,8 @@ def sample_async(self, scheduled_requests: ScheduledRequests, finish_reasons=finish_reasons, sequence_lengths=sequence_lengths, log_probs=log_probs, - cum_log_probs=cum_log_probs) + cum_log_probs=cum_log_probs, + gathered_ids=gathered_ids) sampler_event = torch.cuda.Event() sampler_event.record() @@ -710,7 +730,8 @@ def sample_async(self, scheduled_requests: ScheduledRequests, return SampleStateTRTLLM(scheduled_requests=scheduled_requests, device=device, host=host, - sampler_event=sampler_event) + sampler_event=sampler_event, + finalize_events=finalize_events) @torch.inference_mode() def update_requests(self, state: SampleStateTRTLLM): @@ -751,8 +772,7 @@ def update_requests_single_beam_single_step(self, state: SampleStateTRTLLM): reqs_with_new_tokens = [ r for r in reqs - if (sequence_lengths_host_data[r.py_seq_slot] > r.get_num_tokens(0) - or self.is_trt_overlap) + if (sequence_lengths_host_data[r.py_seq_slot] > r.get_num_tokens(0)) ] # Add new tokens @@ -799,7 +819,7 @@ def update_requests_multiple_beams_or_drafting(self, ) if state.host.cum_log_probs is not None else None log_probs_host = state.host.log_probs.tolist( ) if state.host.log_probs is not None else None - finalize_events = {} + finalize_events = state.finalize_events reqs = [ r for r in state.scheduled_requests.context_requests @@ -821,7 +841,6 @@ def update_requests_multiple_beams_or_drafting(self, for beam in range(beam_width): seq_len = sequence_lengths_host_data[seq_slot * beam_width + beam] - seq_len = seq_len + 1 if self.is_trt_overlap else seq_len num_new_tokens[beam] = min( num_generated_tokens, seq_len - request.get_num_tokens(beam)) @@ -846,8 +865,7 @@ def update_requests_multiple_beams_or_drafting(self, }) if request.py_return_log_probs: - cum_log_probs.append( - cum_log_probs_host[seq_slot * beam_width + beam]) + cum_log_probs.append(cum_log_probs_host[seq_slot][beam]) finished_state = FinishedState( finish_reasons[seq_slot * beam_width + beam]) @@ -869,48 +887,37 @@ def update_requests_multiple_beams_or_drafting(self, if finished_sum_host[seq_slot] == beam_width: request.state = LlmRequestState.GENERATION_COMPLETE - if beam_width > 1: - finalize_events[ - request.request_id] = self._finalize_request( - request, False) - elif request.streaming and beam_width > 1: - finalize_events[request.request_id] = self._finalize_request( - request, True) - # post process all requests if necessary - if beam_width > 1: - for request in reqs: - if request.request_id in finalize_events: - self._post_process_request( - request, finalize_events[request.request_id]) + for request in reqs: + if request.request_id in finalize_events: + self._post_process_request(request, state) def _finalize_request(self, request: LlmRequest, streaming: bool): """ Finalizes the request. This is necessary for beam search. """ - seq_slot = request.seq_slot + seq_slot = request.py_seq_slot event = self.algs.decoder.finalize(self.store["decoder_state"], seq_slot, request.sampling_config, streaming) return event def _post_process_request(self, request: LlmRequest, - finalize_event: CudaEvent): + state: SampleStateTRTLLM): """ Post Process the request. Updates the sequence according to the beam search results. request: LlmRequest which shall be post processed finalize_event: CudaEvent to wait for the finalize step to finish """ - seq_slot = request.seq_slot + seq_slot = request.py_seq_slot beam_width = request.sampling_config.beam_width # synchronize on the finalize event before continuing the post processing. - finalize_event.synchronize() + # should be unnecessary, as already wait for the sampler event in update_requests + state.finalize_events[request.request_id].synchronize() # Get these values again, as they might have changed during the finalize step - output_ids_host = self.store["decoder_state"].gathered_ids.to('cpu') - sequence_lengths_host = self.store["decoder_state"].sequence_lengths.to( - 'cpu') + output_ids_host = state.host.gathered_ids + sequence_lengths_host = state.host.sequence_lengths if request.py_return_log_probs: - log_probs_host = self.store["decoder_state"].log_probs.to('cpu') - cum_log_probs_host = self.store["decoder_state"].cum_log_probs.to( - 'cpu') + log_probs_host = state.host.log_probs + cum_log_probs_host = state.host.cum_log_probs generated_tokens = [[0]] * beam_width log_probs = [[] for _ in range(beam_width)] diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler.py index 26df44874a09..d7a9249dd365 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler.py @@ -73,12 +73,14 @@ def __init__( self, max_num_requests: int, kv_cache_manager, + peft_cache_manager: tb_internal.batch_manager.PeftCacheManager | None, scheduler_policy: tb_executor.CapacitySchedulerPolicy = tb_executor. CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, two_step_lookahead: bool = False, ): super(BindCapacityScheduler, self).__init__() self.kv_cache_manager = kv_cache_manager + self.peft_cache_manager = peft_cache_manager self.impl = tb_internal.algorithms.CapacityScheduler( max_num_requests=max_num_requests, @@ -91,7 +93,8 @@ def __init__( def schedule_request( self, active_requests: RequestList ) -> tuple[list[LlmRequest], list[LlmRequest], list[LlmRequest]]: - return self.impl(active_requests, self.kv_cache_manager) + return self.impl(active_requests, self.kv_cache_manager, + self.peft_cache_manager) class GuaranteedNoEvictScheduler(CapacityScheduler): diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py index dd709cfbfe84..6918b5739059 100644 --- a/tensorrt_llm/_torch/speculative/__init__.py +++ b/tensorrt_llm/_torch/speculative/__init__.py @@ -2,9 +2,10 @@ from .interface import SpecMetadata from .mtp import MTPEagleWorker, MTPSpecMetadata, MTPWorker from .ngram import NGramDrafter, NGramPoolManager -from .utils import (get_num_spec_layers, get_spec_decoder, get_spec_drafter, - get_spec_metadata, get_spec_resource_manager, - get_spec_worker) +from .utils import (get_num_extra_kv_tokens, get_num_spec_layers, + get_spec_decoder, get_spec_drafter, get_spec_metadata, + get_spec_resource_manager, get_spec_worker, + update_spec_config_from_model_config) __all__ = [ "Eagle3SpecMetadata", @@ -14,10 +15,12 @@ "NGramDrafter", "NGramPoolManager", "SpecMetadata", + "get_num_extra_kv_tokens", "get_num_spec_layers", "get_spec_decoder", "get_spec_drafter", "get_spec_metadata", "get_spec_resource_manager", "get_spec_worker", + "update_spec_config_from_model_config", ] diff --git a/tensorrt_llm/_torch/speculative/drafter.py b/tensorrt_llm/_torch/speculative/drafter.py index d0f5a44d7781..e08044cbb4f6 100644 --- a/tensorrt_llm/_torch/speculative/drafter.py +++ b/tensorrt_llm/_torch/speculative/drafter.py @@ -1,24 +1,23 @@ from abc import ABC, abstractmethod from typing import Optional -from ..pyexecutor.resource_manager import BaseResourceManager +from ..pyexecutor.resource_manager import ResourceManager from ..pyexecutor.scheduler import ScheduledRequests class Drafter(ABC): - - def __init__( - self, - spec_resource_manager: Optional[BaseResourceManager] = None, - ): - self.spec_resource_manager = spec_resource_manager + """Abstract base class for all drafter implementations.""" @abstractmethod def prepare_draft_tokens( self, scheduled_requests: ScheduledRequests, + resource_manager: Optional[ResourceManager] = None, ) -> None: """ Prepare the drafter tokens for the forward computation this step. + + Args: + scheduled_requests: The scheduled requests for this iteration """ raise NotImplementedError diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 3006ccdb4ef3..46fe18e0584a 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -77,7 +77,8 @@ def has_spec_decoder(self): return self.is_mtp() or self.is_eagle3() or self.is_eagle3_one_model() def has_spec_drafter(self): - return self.is_ngram() or self.is_user_provided() + return self.is_eagle3() or self.is_draft_target() or self.is_ngram( + ) or self.is_user_provided() def extend_ctx(self, attention_backend: Type[AttentionBackend]): """ diff --git a/tensorrt_llm/_torch/speculative/model_drafter.py b/tensorrt_llm/_torch/speculative/model_drafter.py new file mode 100644 index 000000000000..318cce8c736f --- /dev/null +++ b/tensorrt_llm/_torch/speculative/model_drafter.py @@ -0,0 +1,401 @@ +from __future__ import annotations + +import traceback +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +import torch + +from tensorrt_llm._utils import nvtx_range +from tensorrt_llm.logger import logger + +from ..pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig +from ..pyexecutor.resource_manager import BaseResourceManager, ResourceManager +from ..pyexecutor.sampler import Sampler, SampleState +from ..pyexecutor.scheduler import ScheduledRequests +from ..pyexecutor.seq_slot_manager import SeqSlotManager +from .drafter import Drafter + +if TYPE_CHECKING: + from ..pyexecutor.model_engine import ModelEngine + from .interface import SpeculativeDecodingMode + + +# Place the tool function here to avoid circular import +def get_draft_model_prompt(spec_dec_mode: SpeculativeDecodingMode, + input_tokens: torch.Tensor) -> torch.Tensor: + """ + Can be used to modify prompts for speculative algorithms that need to update tokens + before drafting. + """ + if spec_dec_mode.is_eagle3(): + # EAGLE3 always throws away the first token when processing draft inputs + return input_tokens[1:] + return input_tokens + + +class ModelDrafter(Drafter): + """Model-based drafter that uses a draft model to generate draft tokens.""" + + def __init__( + self, + spec_config: "DecodingBaseConfig", + draft_model_engine: "ModelEngine", + max_draft_tokens: int, + draft_seq_slot_manager: SeqSlotManager, + sampler: Sampler, + spec_resource_manager: Optional[BaseResourceManager] = None, + ): + # Validate required parameters + if draft_model_engine is None: + raise ValueError("draft_model_engine cannot be None") + if max_draft_tokens < 0: + raise ValueError(f"max_draft_tokens must be >= 0") + + # Model and resource management + self.draft_model_engine = draft_model_engine + self.draft_seq_slot_manager = draft_seq_slot_manager + self.spec_resource_manager = spec_resource_manager + + # Configuration + self.spec_config = spec_config + self.max_draft_tokens = max_draft_tokens + + # Sampling + self.sampler = sampler + + def _create_draft_request(self, request_id: int, max_new_tokens: int, + input_tokens: Optional[List], + sampling_config: SamplingConfig, + return_perf_metrics: bool) -> LlmRequest: + """Create a draft request with common parameters.""" + return LlmRequest(request_id=request_id, + max_new_tokens=max_new_tokens, + input_tokens=input_tokens, + sampling_config=sampling_config, + return_perf_metrics=return_perf_metrics, + is_streaming=False, + is_draft=True) + + def _initialize_draft_tokens(self, request: LlmRequest) -> Tuple[int, int]: + """Initialize draft token tracking for a request.""" + num_draft_tokens = len( + request.py_last_draft_tokens + ) if request.py_last_draft_tokens is not None else 0 + request.py_draft_tokens = [] + + num_accepted_tokens = request.py_num_accepted_draft_tokens + num_rejected_tokens = num_draft_tokens - num_accepted_tokens + assert num_rejected_tokens >= 0 + + return num_draft_tokens, num_accepted_tokens + + def _create_context_request(self, request: LlmRequest, + input_tokens: Any) -> LlmRequest: + """Create a context request for first-time drafting.""" + new_request = self._create_draft_request(request.py_request_id, + request.py_max_new_tokens, + input_tokens, + request.sampling_config, + request.return_perf_metrics) + + begin_compute, end_compute = request.py_last_context_chunk + if begin_compute is not None: + new_request.context_current_position = begin_compute + new_request.context_chunk_size = end_compute - begin_compute + return new_request + + def _create_generation_request(self, request: LlmRequest, + input_tokens: Any) -> LlmRequest: + """Create a generation request when no tokens were accepted.""" + new_request = self._create_draft_request(request.py_request_id, + request.py_max_new_tokens, + input_tokens[:-1], + request.sampling_config, + request.return_perf_metrics) + # Explicitly add the last token so get_last_tokens() returns the right value + new_request.add_new_token(input_tokens[-1], 0) + new_request.state = LlmRequestState.GENERATION_IN_PROGRESS + return new_request + + def _create_accepted_tokens_request(self, request: LlmRequest, + input_tokens: Any, + num_accepted_tokens: int) -> LlmRequest: + """ + Create a chunked context request for accepted tokens. + Only applicable if the draft model needs to recompute KV cache for accepted tokens (e.g. eagle 3) + """ + new_request = self._create_draft_request(request.py_request_id, + request.py_max_new_tokens, + input_tokens, + request.sampling_config, + request.return_perf_metrics) + new_request.context_chunk_size = num_accepted_tokens + 1 + new_request.context_current_position = len( + input_tokens) - num_accepted_tokens - 1 + return new_request + + def _create_draft_request_for_request( + self, request: LlmRequest) -> Optional[LlmRequest]: + """Create a draft request based on the original request state.""" + num_draft_tokens, num_accepted_tokens = self._initialize_draft_tokens( + request) + input_tokens = get_draft_model_prompt(self.spec_config.spec_dec_mode, + request.get_tokens()[0]) + + # First time seeing this request - context request + if request.max_beam_num_tokens - 1 == request.py_prompt_len: + # This is the first time the draft model is seeing this request. + # Prepare a context request. We discard the first token and take + # the newly decoded one - this is the convention for EAGLE 2 and 3. + assert num_draft_tokens == 0 + return self._create_context_request(request, input_tokens) + + # No tokens accepted - generation request + elif num_accepted_tokens == 0: + return self._create_generation_request(request, input_tokens) + + # Tokens accepted - chunked context request + else: + return self._create_accepted_tokens_request(request, input_tokens, + num_accepted_tokens) + + def _add_to_draft_batch(self, draft_batch: ScheduledRequests, + draft_request: LlmRequest, + original_request: LlmRequest) -> None: + """Add the draft request to the appropriate batch list.""" + # Copy additional properties + draft_request.py_stop_words_list = original_request.py_stop_words_list + + # Add to appropriate batch based on request type + if draft_request.state == LlmRequestState.GENERATION_IN_PROGRESS: + draft_batch.generation_requests.append(draft_request) + else: + draft_batch.context_requests.append(draft_request) + + @nvtx_range("_prepare_draft_batch") + def _prepare_draft_batch( + self, scheduled_requests: ScheduledRequests) -> ScheduledRequests: + """ + Prepares a batch for the draft model engine. Draft tokens are only produced + for generation requests. + + The requests are prepared as follows: + 1. The first time the draft engine sees a request, it's a context request. + 2. Otherwise, if draft tokens were accepted on the last target model decoding + step, it's a chunked context request (we process all the accepted tokens together). + 3. Otherwise, it's a generation request. + + Args: + scheduled_requests: The scheduled requests to prepare draft batch for + + Returns: + ScheduledRequests: The prepared draft batch + """ + try: + draft_batch = ScheduledRequests() + + for request in scheduled_requests.context_requests: + if request.is_first_context_chunk: + # Ignore requests which still need to be processed by the target model. + continue + + # We hit this path if we're doing chunked prefill. The target model processed + # a prefill chunk on the last iteration. Now, we need to fill in the KV cache + # for the draft model too. + all_tokens = request.get_tokens()[0] + input_tokens = get_draft_model_prompt( + self.spec_config.spec_dec_mode, all_tokens) + + new_request = self._create_context_request( + request, input_tokens) + self._add_to_draft_batch(draft_batch, new_request, request) + + for request in scheduled_requests.generation_requests: + if request.py_draft_pages_allocated == 0: + # No space for draft tokens + continue + + # Stop drafting when we hit the max seqlen. We still need dummy draft + # tokens attached to the requests to make sure everything works properly + # with CUDA graph. These dummy tokens are already added by + # _prepare_draft_requests to make the KV cache/scheduler aware of the fact + # that we want to do spec decoding, so no need to do anything else here. + # This makes the perf for this case suboptimal, but that's OK - this is + # a corner case for weird models like the llama 3.1 8b EAGLE3 implementation. + if request.max_beam_num_tokens - 1 >= self.draft_model_engine.max_seq_len: + continue + + draft_request = self._create_draft_request_for_request(request) + if draft_request is not None: + self._add_to_draft_batch(draft_batch, draft_request, + request) + + return draft_batch + + except Exception as e: + logger.error(f"Error in _prepare_draft_batch: {str(e)}") + traceback.print_exc() + raise e + + def _should_disable_cuda_graph( + self, previous_batch: Optional[SampleState]) -> bool: + """Check if CUDA graph should be disabled for the current forward pass.""" + if previous_batch is not None: + return False + return self.spec_config.spec_dec_mode.needs_kv_cache_recompute() + + def _forward_draft_model( + self, + draft_batch: ScheduledRequests, + resource_manager: ResourceManager, + previous_batch: Optional[SampleState] = None) -> Dict[str, Any]: + """Forward pass through the draft model.""" + if self._should_disable_cuda_graph(previous_batch): + with self.draft_model_engine.no_cuda_graph(): + outputs = self.draft_model_engine.forward( + draft_batch, resource_manager) + else: + new_tensors_device = previous_batch.device if previous_batch else None + outputs = self.draft_model_engine.forward( + draft_batch, + resource_manager, + new_tensors_device=new_tensors_device) + + # Handle d2t data if available + if hasattr(self.draft_model_engine.model.model, 'd2t'): + outputs['d2t'] = self.draft_model_engine.model.model.d2t.data + + return outputs + + def _sample_async(self, draft_batch: ScheduledRequests, + outputs: Dict[str, Any]) -> Optional[SampleState]: + """Sample tokens from draft model outputs.""" + try: + if self.sampler is not None: + return self.sampler.sample_async(draft_batch, outputs) + return None + except Exception as e: + logger.error(f"Error in sampling: {str(e)}") + return None + + def _update_request_states(self, + scheduled_requests: ScheduledRequests) -> None: + """Update request states after processing.""" + for request in scheduled_requests.context_requests: + if request.state != LlmRequestState.GENERATION_COMPLETE: + request.move_to_next_context_chunk() + if request.context_remaining_length == 0: + request.state = LlmRequestState.GENERATION_IN_PROGRESS + + def _update_requests(self, sample_state: SampleState) -> None: + """Update requests with sample state.""" + if self.sampler is not None: + self.sampler.update_requests(sample_state) + + def _process_decoded_tokens( + self, draft_batch: ScheduledRequests, + req_id_to_old_request: Dict[int, LlmRequest]) -> List[LlmRequest]: + """Process decoded tokens and determine which requests to continue processing.""" + new_requests = [] + for req in draft_batch.all_requests(): + target_model_req = req_id_to_old_request[req.py_request_id] + if target_model_req.state != LlmRequestState.GENERATION_IN_PROGRESS: + # This is a chunked prefill request and we have more prefill chunks + # to process. Defer adding draft tokens until the whole prompt is processed. + self.draft_seq_slot_manager.free_resources(req) + continue + + target_model_req.py_draft_tokens.append(req.get_last_tokens(0)) + if req.state != LlmRequestState.GENERATION_COMPLETE and len( + target_model_req.py_draft_tokens + ) < target_model_req.py_draft_pages_allocated: + new_requests.append(req) + else: + self.draft_seq_slot_manager.free_resources(req) + + return new_requests + + def _pad_to_max_draft_tokens(self, + scheduled_requests: ScheduledRequests) -> None: + """Pad draft tokens to maximum length for all generation requests.""" + for req in scheduled_requests.generation_requests: + max_draft_tokens = self.max_draft_tokens + num_draft_tokens = len(req.py_draft_tokens) + req.py_draft_tokens.extend( + 0 for _ in range(max_draft_tokens - num_draft_tokens)) + + @nvtx_range("prepare_draft_tokens") + def prepare_draft_tokens( + self, + scheduled_requests: ScheduledRequests, + resource_manager: Optional[ResourceManager] = None, + ) -> None: + """ + Prepare draft tokens for the scheduled requests. + + Args: + scheduled_requests: The scheduled requests for this iteration + resource_manager: The resource manager for this iteration + """ + if not self.draft_model_engine: + raise ValueError("Draft model engine is not set") + + if resource_manager is None: + raise ValueError("Resource manager is required") + + try: + draft_batch = self._prepare_draft_batch(scheduled_requests) + + if draft_batch.batch_size == 0: + return + + self.draft_seq_slot_manager.prepare_resources(draft_batch) + + req_id_to_old_request = { + req.py_request_id: req + for req in scheduled_requests.all_requests() + } + + # Initial forward pass + outputs = self._forward_draft_model(draft_batch, resource_manager) + sample_state = self._sample_async(draft_batch, outputs) + previous_batch = sample_state + + self._update_request_states(draft_batch) + + # Convert context requests to generation requests + draft_batch.generation_requests = draft_batch.context_requests + draft_batch.generation_requests + draft_batch.context_requests = [] + + # Generate remaining draft tokens iteratively + for i in range(self.max_draft_tokens - 1): + if len(draft_batch.generation_requests) == 0: + break + + outputs = self._forward_draft_model(draft_batch, + resource_manager, + previous_batch) + sample_state = self._sample_async(draft_batch, outputs) + self._update_request_states(draft_batch) + if previous_batch is not None: + self._update_requests(previous_batch) + new_requests = self._process_decoded_tokens( + previous_batch.scheduled_requests, + req_id_to_old_request) + else: + new_requests = [] + draft_batch.generation_requests = new_requests + previous_batch = sample_state + + # Final cleanup + if previous_batch is not None: + self._update_requests(previous_batch) + self._process_decoded_tokens(previous_batch.scheduled_requests, + req_id_to_old_request) + self._pad_to_max_draft_tokens(scheduled_requests) + + except Exception as e: + traceback.print_exc() + error_msg = str(e) + logger.error(f"Encountered an error in decode: {error_msg}") + raise e diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 72316a2e474a..7d383257b5ec 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -232,7 +232,7 @@ def _request_common_handling(self, request: LlmRequest, assert not request.py_return_context_logits, "return_context_logits not implemented for MTPSampler" assert not request.py_return_generation_logits, "return_generation_logits not implemented for MTPSampler" assert not request.py_return_log_probs, "return_log_probs not implemented for MTPSampler" - request.py_draft_tokens = next_draft_tokens[request.seq_slot] + request.py_draft_tokens = next_draft_tokens[request.py_seq_slot] request.py_decoding_iter += 1 def update_requests(self, state: SampleStateMTP) -> None: @@ -253,7 +253,7 @@ def update_requests(self, state: SampleStateMTP) -> None: for req in state.scheduled_requests.generation_requests: if req.state == LlmRequestState.GENERATION_COMPLETE: continue - num_new_tokens = new_tokens_lens[req.seq_slot] + num_new_tokens = new_tokens_lens[req.py_seq_slot] for i in range(num_new_tokens): new_token = add_token(req, new_tokens, beam=beam_idx, step=i) if self._handle_stop_criteria(req, new_token): @@ -269,7 +269,7 @@ def sample_async(self, scheduled_requests: ScheduledRequests, # next_new_tokens_device: input tokens for the next iteration, device tensor, shape: batch_size, nextn + 1 requests = scheduled_requests.all_requests() - slots = torch.as_tensor([r.seq_slot for r in requests]) + slots = torch.as_tensor([r.py_seq_slot for r in requests]) slots = slots.to(device="cuda", non_blocking=True) o_new_tokens = outputs['new_tokens'][:len(requests)] diff --git a/tensorrt_llm/_torch/speculative/ngram.py b/tensorrt_llm/_torch/speculative/ngram.py index 1d015a58b9b9..9113900ef94c 100644 --- a/tensorrt_llm/_torch/speculative/ngram.py +++ b/tensorrt_llm/_torch/speculative/ngram.py @@ -5,7 +5,7 @@ from tensorrt_llm.logger import logger from ..pyexecutor.llm_request import * -from ..pyexecutor.resource_manager import BaseResourceManager +from ..pyexecutor.resource_manager import BaseResourceManager, ResourceManager from ..pyexecutor.scheduler import ScheduledRequests from .drafter import Drafter @@ -59,10 +59,10 @@ def __init__(self, spec_config: "NGramDecodingConfig", self.start_index = {} def get_max_resource_count(self) -> int: - raise self.max_num_requests + return self.max_num_requests def get_needed_resource_to_completion(self, request: LlmRequest) -> int: - raise 0 + return 0 def prepare_resources(self, scheduled_batch: ScheduledRequests): pass @@ -167,12 +167,13 @@ def __init__( ngram_pool_manager: NGramPoolManager = None, ): assert ngram_pool_manager is not None, "NGram needs a resource manager to maintain the pool." - super().__init__(spec_resource_manager=ngram_pool_manager) self.max_draft_len = spec_config.max_draft_len + self.spec_resource_manager = ngram_pool_manager def prepare_draft_tokens( self, scheduled_requests: ScheduledRequests, + resource_manager: Optional[ResourceManager] = None, ) -> None: # Sort by request_id when py_batch_idx is None as a fallback. # This happens in the disagg case: for a set of new requests, we draft diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 882dfdf924d7..bc866550470f 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -1,9 +1,11 @@ from tensorrt_llm._torch.pyexecutor.sampler import TorchSampler from tensorrt_llm._torch.speculative.interface import SpecMetadata +from ..pyexecutor.seq_slot_manager import SeqSlotManager from .eagle3 import (Eagle3OneModelSampler, Eagle3OneModelSpecMetadata, Eagle3OneModelWorker, Eagle3ResourceManager, Eagle3SpecMetadata) +from .model_drafter import ModelDrafter from .mtp import (MTPEagleWorker, MTPHiddenStatesManager, MTPSampler, MTPSpecMetadata, MTPWorker) from .ngram import NGramDrafter, NGramPoolManager @@ -55,9 +57,7 @@ def get_spec_metadata(spec_config, return None -def get_spec_resource_manager(model_engine, - draft_model_engine=None, - drafter=None): +def get_spec_resource_manager(model_engine, draft_model_engine=None): spec_config = model_engine.spec_config if spec_config is None: return None @@ -93,9 +93,10 @@ def get_spec_resource_manager(model_engine, max_seq_len, max_num_tokens, ) - if spec_dec_mode.is_ngram() or spec_dec_mode.is_user_provided(): - assert drafter is not None, "Drafter is required for ngram or user provided speculative decoding." - return drafter.spec_resource_manager + if spec_dec_mode.is_ngram(): + return NGramPoolManager(spec_config, max_num_requests) + if spec_dec_mode.is_user_provided(): + return spec_config.resource_manager return None @@ -113,16 +114,26 @@ def get_spec_decoder(sampler_args: TorchSampler.Args, f"Unsupported speculative decoding mode: {spec_config.spec_dec_mode}") -def get_spec_drafter(model_engine): +def get_spec_drafter(model_engine, draft_model_engine, sampler, + spec_resource_manager): spec_config = model_engine.spec_config - max_num_requests = model_engine.batch_size if spec_config is None: return None - if spec_config.spec_dec_mode.is_ngram(): - return NGramDrafter(spec_config, - NGramPoolManager(spec_config, max_num_requests)) + if spec_config.spec_dec_mode.is_user_provided(): return spec_config.drafter + + max_num_requests = model_engine.batch_size + if spec_config.spec_dec_mode.is_draft_target( + ) or spec_config.spec_dec_mode.is_eagle3(): + return ModelDrafter(spec_config, draft_model_engine, + spec_config.max_draft_len, + SeqSlotManager(max_num_requests), sampler, + spec_resource_manager) + + if spec_config.spec_dec_mode.is_ngram(): + return NGramDrafter(spec_config, spec_resource_manager) + return None @@ -142,3 +153,24 @@ def get_spec_worker(spec_config, mapping): if spec_config.spec_dec_mode.is_eagle3_one_model(): return Eagle3OneModelWorker(spec_config, mapping) return None + + +def get_num_extra_kv_tokens(spec_config): + """ + Implementation detail for one model implementations of speculative decoding. Extra + KV cache tokens are required. + """ + if spec_config is None: + return 0 + if spec_config.spec_dec_mode.is_eagle3_one_model( + ) or spec_config.spec_dec_mode.is_mtp_eagle(): + return spec_config.max_draft_len - 1 + return 0 + + +def update_spec_config_from_model_config(spec_config, model_config): + if spec_config.spec_dec_mode.is_mtp(): + # Use `max_draft_len` for several low-level APIs. TODO: Remove this after distinguishing them. + spec_config.max_draft_len = spec_config.num_nextn_predict_layers + # Use `num_nextn_predict_layers_from_model_config` to decide decoding mode MTP / MTP_EAGLE. + spec_config.num_nextn_predict_layers_from_model_config = model_config.num_nextn_predict_layers diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index f687e9d9f55b..5710dbdc6ae4 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -1,5 +1,4 @@ import contextlib -import os import threading from dataclasses import dataclass from enum import Enum @@ -100,13 +99,6 @@ def shape(self): return self.fp4_tensor.shape -_disable_fp4_allgather = os.getenv("TLLM_DISABLE_FP4_ALLGATHER", "0") == "1" - - -def disable_fp4_allgather(): - return _disable_fp4_allgather - - def compute_swizzled_sf_shape(row: int, col: int): padded_row = pad_up(row, 128) padded_col = pad_up(col, 4) @@ -204,7 +196,17 @@ def next_positive_power_of_2(x: int) -> int: if x < 1: return 1 - return 1 << (x - 1).bit_length() + # Following code is equivalent to 1 << (x - 1).bit_length() + # But this impl does not contain bit_length() so can be used by torch compile. + # It can correctly handle 64bit number which should be enough for now. + n = x - 1 + n |= n >> 1 + n |= n >> 2 + n |= n >> 4 + n |= n >> 8 + n |= n >> 16 + n |= n >> 32 + return n + 1 def last_positive_power_of_2(x: int) -> int: diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 9c3197e7c93c..b07430224afc 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -180,6 +180,22 @@ def str_dtype_to_torch(dtype): fp8=DataType.FP8, ) +_binding_dtype_size = { + DataType.INT64: 8, + DataType.FLOAT: 4, + DataType.INT32: 4, + DataType.BF16: 2, + DataType.HALF: 2, + DataType.BOOL: 1, + DataType.FP8: 1, + DataType.INT8: 1, + DataType.UINT8: 1, +} + + +def binding_dtype_size(dtype: DataType): + return _binding_dtype_size[dtype] + def str_dtype_to_binding(dtype): ret = _str_to_binding_dtype_dict.get(dtype) @@ -493,7 +509,7 @@ def mpi_barrier(): def mpi_broadcast(obj, root=0): - return mpi_comm().bcast(obj, root) if ENABLE_MULTI_DEVICE else obj + return mpi_comm().bcast(obj, root) if is_multi_device_enable() else obj def mpi_allgather(obj): @@ -1063,3 +1079,14 @@ def _unique_tokens_to_json(data): "token_id": data.token_id, "token_extra_id": data.token_extra_id } + + +def is_multi_device_enable(): + """ + This method evaluates if we are running on multiple GPUs and the flag ENABLE_MULTI_DEVICE is set. + So we can avoid broadcast calls on single GPU. + Issue: https://github.com/NVIDIA/TensorRT-LLM/issues/5927 + ENABLE_MULTI_DEVICE is true by default when building tensorrt-llm so we need to also check + the number of devices + """ + return local_mpi_size() > 1 diff --git a/tensorrt_llm/bench/benchmark/low_latency.py b/tensorrt_llm/bench/benchmark/low_latency.py index 490ac62f4f5b..fd701700a29a 100644 --- a/tensorrt_llm/bench/benchmark/low_latency.py +++ b/tensorrt_llm/bench/benchmark/low_latency.py @@ -9,11 +9,14 @@ import yaml from click_option_group import (MutuallyExclusiveOptionGroup, OptionGroup, optgroup) +from huggingface_hub import snapshot_download +from tensorrt_llm import LLM as PyTorchLLM from tensorrt_llm._tensorrt_engine import LLM from tensorrt_llm.bench.benchmark.utils.asynchronous import async_benchmark from tensorrt_llm.bench.benchmark.utils.general import generate_warmup_dataset from tensorrt_llm.bench.benchmark.utils.processes import IterationWriter +from tensorrt_llm.bench.build.build import get_model_config from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment from tensorrt_llm.bench.dataclasses.reporting import ReportUtility @@ -21,10 +24,11 @@ from tensorrt_llm.models.modeling_utils import SpeculativeDecodingMode # isort: off -from tensorrt_llm.bench.benchmark.utils.general import get_settings_from_engine +from tensorrt_llm.bench.benchmark.utils.general import get_settings_from_engine, get_settings, ALL_SUPPORTED_BACKENDS # isort: on from tensorrt_llm.bench.utils.data import (create_dataset_from_stream, - initialize_tokenizer) + initialize_tokenizer, + update_metadata_for_multimodal) from tensorrt_llm.logger import logger from tensorrt_llm.sampling_params import SamplingParams @@ -38,15 +42,25 @@ readable=True, path_type=Path, resolve_path=True), - required=True, + default=None, help="Path to a serialized TRT-LLM engine.", ) +@optgroup.option("--backend", + type=click.Choice(ALL_SUPPORTED_BACKENDS), + default="pytorch", + help="The backend to use when running benchmarking.") @optgroup.option( "--kv_cache_free_gpu_mem_fraction", type=float, default=.90, help="The percentage of memory to use for KV Cache after model load.", ) +@optgroup.option( + "--max_seq_len", + type=int, + default=None, + help="Maximum sequence length.", +) @optgroup.group( "Engine Input Configuration", help="Input configuration for driving the engine.", @@ -60,6 +74,20 @@ default=None, help="Pass in a dataset file for parsing instead of stdin.", ) +@optgroup.option( + "--modality", + type=click.Choice(["image", "video"]), + default=None, + help="Modality of the multimodal requests.", +) +@optgroup.option( + "--max_input_len", + type=int, + default=4096, + help= + "Maximum input sequence length to use for multimodal models. This is used only when --modality " + "is specified since the actual number of vision tokens is unknown before the model is run.", +) @optgroup.option( "--num_requests", type=int, @@ -73,6 +101,24 @@ default=2, help="Number of requests warm up benchmark.", ) +@optgroup.option( + "--tp", + type=int, + default=1, + help="tensor parallelism size", +) +@optgroup.option( + "--pp", + type=int, + default=1, + help="pipeline parallelism size", +) +@optgroup.option( + "--ep", + type=int, + default=None, + help="expert parallelism size", +) @optgroup.group("Request Load Control Options", cls=MutuallyExclusiveOptionGroup, help="Limits how requests are loaded.") @@ -134,29 +180,83 @@ def latency_command( logger.info("Preparing to run latency benchmark...") # Parameters from CLI # Model, experiment, and engine params - dataset_path: Path = params.pop("dataset") - num_requests: int = params.pop("num_requests") + dataset_path: Path = params.get("dataset") + num_requests: int = params.get("num_requests") model: str = bench_env.model checkpoint_path: Path = bench_env.checkpoint_path or bench_env.model - engine_dir: Path = params.pop("engine_dir") - concurrency: int = params.pop("concurrency") - beam_width: int = params.pop("beam_width") + engine_dir: Path = params.get("engine_dir") + concurrency: int = params.get("concurrency") + beam_width: int = params.get("beam_width") warmup: int = params.get("warmup") - # Engine configuration parsing - exec_settings, build_cfg = get_settings_from_engine(engine_dir) - exec_settings["model"] = model - engine_tokens = exec_settings["settings_config"]["max_num_tokens"] - engine_max_seq_len = build_cfg["max_seq_len"] + modality: str = params.get("modality") + max_input_len: int = params.get("max_input_len") + max_seq_len: int = params.get("max_seq_len") + backend: str = params.get("backend") + model_type = get_model_config(model, checkpoint_path).model_type # Runtime Options - kv_cache_percent = params.pop("kv_cache_free_gpu_mem_fraction") - medusa_choices = params.pop("medusa_choices") + kv_cache_percent = params.get("kv_cache_free_gpu_mem_fraction") + medusa_choices = params.get("medusa_choices") # Reporting Options report_json: Path = params.pop("report_json") iteration_log: Path = params.pop("iteration_log") iteration_writer = IterationWriter(iteration_log) + # Initialize the HF tokenizer for the specified model. + tokenizer = initialize_tokenizer(checkpoint_path) + + # Dataset Loading and Preparation + with open(dataset_path, "r") as dataset: + metadata, requests = create_dataset_from_stream( + tokenizer, + dataset, + num_requests=num_requests, + model_dir=checkpoint_path, + model_type=model_type, + modality=modality, + max_input_seq_len_for_multimodal=max_input_len) + + metadata.dataset_path = dataset_path + + if modality is None: + # Log dataset info + # NOTE: This table is only accurate for non-multimodal models. + # The accurate table for multimodal models will be logged after the benchmark is done. + logger.info(metadata.get_summary_for_print()) + + # Engine configuration parsing for PyTorch backend + kwargs = {} + if backend and backend.lower() in ALL_SUPPORTED_BACKENDS and backend.lower( + ) != "tensorrt": + if bench_env.checkpoint_path is None: + snapshot_download(model) + + exec_settings = get_settings(params, metadata, bench_env.model, + bench_env.checkpoint_path) + kwargs_max_sql = max_seq_len or metadata.max_sequence_length + logger.info(f"Setting PyTorch max sequence length to {kwargs_max_sql}") + kwargs["max_seq_len"] = kwargs_max_sql + elif backend.lower() == "tensorrt": + assert max_seq_len is None, ( + "max_seq_len is not a runtime parameter for C++ backend") + exec_settings, build_cfg = get_settings_from_engine(engine_dir) + engine_max_seq_len = build_cfg["max_seq_len"] + + if metadata.max_sequence_length > engine_max_seq_len: + raise RuntimeError( + f"Engine supports a max sequence of {engine_max_seq_len}. Provided " + "dataset contains a maximum sequence of " + f"{metadata.max_sequence_length}. Please rebuild a new engine to" + "support this dataset.") + else: + raise RuntimeError( + f"Invalid backend: {backend}, please use one of the following: " + f"{ALL_SUPPORTED_BACKENDS}") + + exec_settings["model"] = model + engine_tokens = exec_settings["settings_config"]["max_num_tokens"] + # Update configuration with runtime options exec_settings["settings_config"]["kv_cache_percent"] = kv_cache_percent exec_settings["settings_config"]["max_batch_size"] = 1 @@ -187,32 +287,25 @@ def latency_command( # Construct the runtime configuration dataclass. runtime_config = RuntimeConfig(**exec_settings) - # Initialize the HF tokenizer for the specified model. - ignore_eos = True if runtime_config.decoding_config.decoding_mode == SpeculativeDecodingMode.NONE else False - tokenizer = initialize_tokenizer(checkpoint_path) - eos_id = tokenizer.eos_token_id if not ignore_eos else -1 - pad_id = tokenizer.pad_token_id if not ignore_eos else -1 + llm = None + kwargs = kwargs | runtime_config.get_llm_args() + kwargs['backend'] = backend - # Dataset Loading and Preparation - with open(dataset_path, "r") as dataset: - metadata, requests = create_dataset_from_stream( - tokenizer, dataset, num_requests=num_requests) - metadata.dataset_path = dataset_path + try: + logger.info("Setting up latency benchmark.") - if metadata.max_sequence_length > engine_max_seq_len: - raise RuntimeError( - f"Engine supports a max sequence of {engine_max_seq_len}. Provided " - "dataset contains a maximum sequence of " - f"{metadata.max_sequence_length}. Please rebuild a new engine to" - "support this dataset.") + if "pytorch_backend_config" in kwargs and iteration_log is not None: + kwargs["pytorch_backend_config"].enable_iter_perf_stats = True - logger.info(metadata.get_summary_for_print()) - logger.info("Running experimental latency benchmark.") + if runtime_config.backend == 'pytorch': + llm = PyTorchLLM(**kwargs) + else: + llm = LLM(**kwargs) - llm = None - kwargs = runtime_config.get_llm_args() + ignore_eos = True if runtime_config.decoding_config.decoding_mode == SpeculativeDecodingMode.NONE else False + eos_id = tokenizer.eos_token_id if not ignore_eos else -1 + pad_id = tokenizer.pad_token_id if not ignore_eos else -1 - try: sampling_params = SamplingParams( end_id=eos_id, pad_id=pad_id, @@ -220,7 +313,6 @@ def latency_command( use_beam_search=beam_width > 1, ) post_proc_params = None # No detokenization - llm = LLM(**kwargs) # Perform warmup if requested. if warmup > 0: @@ -228,8 +320,13 @@ def latency_command( warmup_dataset = generate_warmup_dataset(requests, warmup) logger.info("Running warmup.") asyncio.run( - async_benchmark(llm, sampling_params, post_proc_params, - warmup_dataset, False, concurrency)) + async_benchmark(llm, + sampling_params, + post_proc_params, + warmup_dataset, + False, + concurrency, + modality=modality)) # WAR: IterationResult is a singleton tied to the executor. # Since the benchmark calls asyncio.run() multiple times (e.g., during warmup), # we must reset it to ensure it attaches to the correct event loop. @@ -238,11 +335,21 @@ def latency_command( with iteration_writer.capture(): statistics = asyncio.run( - async_benchmark(llm, sampling_params, post_proc_params, - requests, True, concurrency, - iteration_writer.full_address)) + async_benchmark(llm, + sampling_params, + post_proc_params, + requests, + True, + concurrency, + iteration_writer.full_address, + modality=modality)) logger.info(f"Benchmark done. Reporting results...") + + if modality is not None: + # For multimodal models, we need to update the metadata with the correct input lengths + metadata = update_metadata_for_multimodal(metadata, statistics) + report_utility = ReportUtility(statistics, metadata, runtime_config, logger, kwargs, True) if report_json: diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index 6c5279e2b361..9dbee903ec2c 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -15,7 +15,7 @@ # isort: off from tensorrt_llm.bench.benchmark.utils.general import ( - get_settings_from_engine, get_settings) + get_settings_from_engine, get_settings, ALL_SUPPORTED_BACKENDS) # isort: on from tensorrt_llm import LLM as PyTorchLLM from tensorrt_llm._tensorrt_engine import LLM @@ -45,7 +45,7 @@ help="Path to a serialized TRT-LLM engine.", ) @optgroup.option("--backend", - type=click.Choice(["pytorch", "tensorrt", "_autodeploy"]), + type=click.Choice(ALL_SUPPORTED_BACKENDS), default="pytorch", help="The backend to use when running benchmarking.") @optgroup.option( @@ -305,10 +305,11 @@ def throughput_command( logger.info(metadata.get_summary_for_print()) # Engine configuration parsing - if backend and backend.lower() in ["pytorch", "_autodeploy"]: + if backend and backend.lower() in ALL_SUPPORTED_BACKENDS and backend.lower( + ) != "tensorrt": # If we're dealing with a model name, perform a snapshot download to # make sure we have a local copy of the model. - if checkpoint_path is None: + if bench_env.checkpoint_path is None: snapshot_download(model) exec_settings = get_settings(params, metadata, bench_env.model, @@ -387,6 +388,9 @@ def throughput_command( logger.warning( "Ignore extended_runtime_perf_knob_config for _autodeploy backend." ) + kwargs["world_size"] = kwargs.pop("tensor_parallel_size", None) + kwargs.pop("pipeline_parallel_size", None) + llm = AutoDeployLLM(**kwargs) else: llm = LLM(**kwargs) diff --git a/tensorrt_llm/bench/benchmark/utils/asynchronous.py b/tensorrt_llm/bench/benchmark/utils/asynchronous.py index ae20343f45bd..ed8338d9243b 100644 --- a/tensorrt_llm/bench/benchmark/utils/asynchronous.py +++ b/tensorrt_llm/bench/benchmark/utils/asynchronous.py @@ -47,7 +47,9 @@ def __init__(self, def _task_done_callback(self, task: asyncio.Task) -> None: self._tasks.discard(task) if task.exception() is not None and not self._stop.is_set(): - logger.error("Exception raised during inference - stopping") + logger.error( + f"Stopping benchmarking due to following exception raised during inference: {task.exception()}" + ) self.stop() async def process_request(self, request: InferenceRequest, diff --git a/tensorrt_llm/bench/benchmark/utils/general.py b/tensorrt_llm/bench/benchmark/utils/general.py index 153e262276fe..bc72b5e14670 100755 --- a/tensorrt_llm/bench/benchmark/utils/general.py +++ b/tensorrt_llm/bench/benchmark/utils/general.py @@ -22,6 +22,8 @@ QuantAlgo.NVFP4.value: "fp8", } +ALL_SUPPORTED_BACKENDS = ["pytorch", "_autodeploy", "tensorrt"] + def get_settings_from_engine( engine_path: Path @@ -86,12 +88,14 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, enable_chunked_prefill = params.get("enable_chunked_prefill", False) kv_cache_dtype = "auto" + kv_cache_config = {} if extra_llm_api_options: with open(extra_llm_api_options, 'r') as f: llm_args_dict = yaml.safe_load(f) - - if "kv_cache_dtype" in llm_args_dict: - kv_cache_dtype = llm_args_dict["kv_cache_dtype"] + kv_cache_config = llm_args_dict.get("kv_cache_config", { + "dtype": "auto", + }) + kv_cache_dtype = kv_cache_config.get("dtype", "auto") enable_chunked_prefill = llm_args_dict.get("enable_chunked_prefill", enable_chunked_prefill) @@ -156,9 +160,11 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, "max_batch_size": max_batch_size } + kv_cache_config["dtype"] = kv_cache_dtype + pyt_options = { "cuda_graph_config": cuda_graph_config, - "kv_cache_dtype": kv_cache_dtype, + "kv_cache_config": kv_cache_config, } backend = params.get("backend", "pytorch") diff --git a/tensorrt_llm/bench/dataclasses/configuration.py b/tensorrt_llm/bench/dataclasses/configuration.py index 0d352fa068c2..a693333230c7 100755 --- a/tensorrt_llm/bench/dataclasses/configuration.py +++ b/tensorrt_llm/bench/dataclasses/configuration.py @@ -58,8 +58,6 @@ def get_llm_args(self) -> Dict: self.world_config.cluster_size, "trust_remote_code": True, - "kv_cache_config": - self.settings_config.get_kvcache_config(), "enable_chunked_prefill": self.settings_config.chunking, "extended_runtime_perf_knob_config": @@ -82,6 +80,10 @@ def get_llm_args(self) -> Dict: if self.backend in backend_config_map: llm_args.update(backend_config_map[self.backend]()) + kv_cache_config = self.settings_config.get_kvcache_config().__dict__ + backend_cache_config = llm_args.pop("kv_cache_config", {}) + llm_args["kv_cache_config"] = backend_cache_config | kv_cache_config + return update_llm_args_with_extra_options(llm_args, self.extra_llm_api_options) @@ -112,7 +114,6 @@ def get_pytorch_perf_config(self) -> PyTorchConfig: def get_autodeploy_perf_config(self) -> Dict: AutoDeployPerfConfig = dict ad_config = AutoDeployPerfConfig() - ad_config["kv_cache_dtype"] = "auto" ad_config["attn_backend"] = "flashinfer" return ad_config diff --git a/tensorrt_llm/bench/dataclasses/reporting.py b/tensorrt_llm/bench/dataclasses/reporting.py index d994000d6d08..476ff50ec2d4 100755 --- a/tensorrt_llm/bench/dataclasses/reporting.py +++ b/tensorrt_llm/bench/dataclasses/reporting.py @@ -11,6 +11,7 @@ from tensorrt_llm.bench.dataclasses.statistics import (BenchmarkStatistics, PercentileStats, RequestRecord) +from tensorrt_llm.llmapi import KvCacheConfig from tensorrt_llm.logger import Logger from tensorrt_llm.models.modeling_utils import SpeculativeDecodingMode @@ -275,8 +276,17 @@ def get_statistics_dict(self) -> Dict[str, Any]: model = self.rt_cfg.model_path or self.rt_cfg.model model_config = ModelConfig.from_pretrained(model, trust_remote_code=True) - validate_and_set_kv_cache_quant(model_config, - self.kwargs["kv_cache_dtype"]) + kv_cache_config = self.kwargs.get("kv_cache_config", + KvCacheConfig()) + if isinstance(kv_cache_config, KvCacheConfig): + kv_cache_dtype = kv_cache_config.dtype + elif isinstance(kv_cache_config, dict): + kv_cache_dtype = kv_cache_config.get("dtype", "auto") + else: + raise ValueError( + f"Invalid kv_cache_config type: {type(kv_cache_config)}.") + + validate_and_set_kv_cache_quant(model_config, kv_cache_dtype) stats_dict["engine"] |= { "backend": diff --git a/tensorrt_llm/builder.py b/tensorrt_llm/builder.py index 5b41abf7e8eb..11d528a853dc 100644 --- a/tensorrt_llm/builder.py +++ b/tensorrt_llm/builder.py @@ -13,12 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. import copy +import dataclasses import json import math import os import shutil import time from dataclasses import dataclass, field +from functools import cache from pathlib import Path from typing import Dict, Optional, Union @@ -557,53 +559,89 @@ def override_attri(attr_name, value): else: override_attri('paged_state', False) + @classmethod + @cache + def get_build_config_defaults(cls): + return { + field.name: field.default + for field in dataclasses.fields(cls) + if field.default is not dataclasses.MISSING + } + @classmethod def from_dict(cls, config, plugin_config=None): config = copy.deepcopy( config ) # it just does not make sense to change the input arg `config` - max_input_len = config.pop('max_input_len') - max_seq_len = config.pop('max_seq_len') - max_batch_size = config.pop('max_batch_size') - max_beam_width = config.pop('max_beam_width') - max_num_tokens = config.pop('max_num_tokens') - opt_num_tokens = config.pop('opt_num_tokens') - opt_batch_size = config.pop('opt_batch_size', 8) + + defaults = cls.get_build_config_defaults() + max_input_len = config.pop('max_input_len', + defaults.get('max_input_len')) + max_seq_len = config.pop('max_seq_len', defaults.get('max_seq_len')) + max_batch_size = config.pop('max_batch_size', + defaults.get('max_batch_size')) + max_beam_width = config.pop('max_beam_width', + defaults.get('max_beam_width')) + max_num_tokens = config.pop('max_num_tokens', + defaults.get('max_num_tokens')) + opt_num_tokens = config.pop('opt_num_tokens', + defaults.get('opt_num_tokens')) + opt_batch_size = config.pop('opt_batch_size', + defaults.get('opt_batch_size')) max_prompt_embedding_table_size = config.pop( - 'max_prompt_embedding_table_size', 0) - - kv_cache_type = KVCacheType( - config.pop('kv_cache_type')) if 'plugin_config' in config else None - gather_context_logits = config.pop('gather_context_logits', False) - gather_generation_logits = config.pop('gather_generation_logits', False) - strongly_typed = config.pop('strongly_typed', True) - force_num_profiles = config.pop('force_num_profiles', None) - weight_sparsity = config.pop('weight_sparsity', False) + 'max_prompt_embedding_table_size', + defaults.get('max_prompt_embedding_table_size')) + + if "kv_cache_type" in config and config["kv_cache_type"] is not None: + kv_cache_type = KVCacheType.from_string(config.pop('kv_cache_type')) + else: + kv_cache_type = None + gather_context_logits = config.pop( + 'gather_context_logits', defaults.get('gather_context_logits')) + gather_generation_logits = config.pop( + 'gather_generation_logits', + defaults.get('gather_generation_logits')) + strongly_typed = config.pop('strongly_typed', + defaults.get('strongly_typed')) + force_num_profiles = config.pop('force_num_profiles', + defaults.get('force_num_profiles')) + weight_sparsity = config.pop('weight_sparsity', + defaults.get('weight_sparsity')) profiling_verbosity = config.pop('profiling_verbosity', - 'layer_names_only') - enable_debug_output = config.pop('enable_debug_output', False) - max_draft_len = config.pop('max_draft_len', 0) - speculative_decoding_mode = config.pop('speculative_decoding_mode', - SpeculativeDecodingMode.NONE) - use_refit = config.pop('use_refit', False) - input_timing_cache = config.pop('input_timing_cache', None) - output_timing_cache = config.pop('output_timing_cache', None) + defaults.get('profiling_verbosity')) + enable_debug_output = config.pop('enable_debug_output', + defaults.get('enable_debug_output')) + max_draft_len = config.pop('max_draft_len', + defaults.get('max_draft_len')) + speculative_decoding_mode = config.pop( + 'speculative_decoding_mode', + defaults.get('speculative_decoding_mode')) + use_refit = config.pop('use_refit', defaults.get('use_refit')) + input_timing_cache = config.pop('input_timing_cache', + defaults.get('input_timing_cache')) + output_timing_cache = config.pop('output_timing_cache', + defaults.get('output_timing_cache')) lora_config = LoraConfig.from_dict(config.get('lora_config', {})) auto_parallel_config = AutoParallelConfig.from_dict( config.get('auto_parallel_config', {})) - max_encoder_input_len = config.pop('max_encoder_input_len', 1024) - weight_streaming = config.pop('weight_streaming', False) - use_strip_plan = config.pop('use_strip_plan', False) + max_encoder_input_len = config.pop( + 'max_encoder_input_len', defaults.get('max_encoder_input_len')) + weight_streaming = config.pop('weight_streaming', + defaults.get('weight_streaming')) + use_strip_plan = config.pop('use_strip_plan', + defaults.get('use_strip_plan')) if plugin_config is None: plugin_config = PluginConfig() if "plugin_config" in config.keys(): plugin_config.update_from_dict(config["plugin_config"]) - dry_run = config.pop('dry_run', False) - visualize_network = config.pop('visualize_network', None) - monitor_memory = config.pop('monitor_memory', False) - use_mrope = config.pop('use_mrope', False) + dry_run = config.pop('dry_run', defaults.get('dry_run')) + visualize_network = config.pop('visualize_network', + defaults.get('visualize_network')) + monitor_memory = config.pop('monitor_memory', + defaults.get('monitor_memory')) + use_mrope = config.pop('use_mrope', defaults.get('use_mrope')) return cls( max_input_len=max_input_len, diff --git a/tensorrt_llm/commands/build.py b/tensorrt_llm/commands/build.py index a47e1485b711..e6b55f6e040b 100644 --- a/tensorrt_llm/commands/build.py +++ b/tensorrt_llm/commands/build.py @@ -38,6 +38,23 @@ from tensorrt_llm.quantization.mode import QuantAlgo +def enum_type(enum_class): + + def parse_enum(value): + if isinstance(value, enum_class): + return value + + if isinstance(value, str): + return enum_class.from_string(value) + + valid_values = [e.name for e in enum_class] + raise argparse.ArgumentTypeError( + f"Invalid value '{value}' of type {type(value).__name__}. Expected one of {valid_values}" + ) + + return parse_enum + + def parse_arguments(): parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter) @@ -131,7 +148,7 @@ def parse_arguments(): parser.add_argument( '--kv_cache_type', default=argparse.SUPPRESS, - type=KVCacheType, + type=enum_type(KVCacheType), help= "Set KV cache type (continuous, paged, or disabled). For disabled case, KV cache is disabled and only context phase is allowed." ) diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index ddbcba2a115e..7de263ea89f4 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -71,7 +71,7 @@ def _signal_handler_cleanup_child(signum, frame): def get_llm_args(model: str, tokenizer: Optional[str] = None, - backend: Optional[str] = None, + backend: str = "pytorch", max_beam_width: int = BuildConfig.max_beam_width, max_batch_size: int = BuildConfig.max_batch_size, max_num_tokens: int = BuildConfig.max_num_tokens, @@ -165,8 +165,8 @@ def launch_server(host: str, help="Hostname of the server.") @click.option("--port", type=int, default=8000, help="Port of the server.") @click.option("--backend", - type=click.Choice(["pytorch"]), - default=None, + type=click.Choice(["pytorch", "trt"]), + default="pytorch", help="Set to 'pytorch' for pytorch path. Default is cpp path.") @click.option('--log_level', type=click.Choice(severity_map.keys()), @@ -362,6 +362,7 @@ def disaggregated(config_file: Optional[str], gen_servers=gen_server_urls, req_timeout_secs=request_timeout, server_start_timeout_secs=server_start_timeout, + max_retries=disagg_cfg.max_retries, ctx_router_config=disagg_cfg.ctx_router_config, gen_router_config=disagg_cfg.gen_router_config, conditional_disagg_config=disagg_cfg.conditional_disagg_config, @@ -429,7 +430,6 @@ def disaggregated_mpi_worker(config_file: Optional[str], log_level: str): disagg_cfg.server_configs) logger.set_level(log_level) - os.environ['TRTLLM_USE_MPI_KVCACHE'] = "1" set_mpi_comm(sub_comm) logger.info( f"mpi_session is provided for LLM instance. Global MPI rank: {global_mpi_rank()}, sub-comm MPI rank: {mpi_rank()}" diff --git a/tensorrt_llm/disaggregated_params.py b/tensorrt_llm/disaggregated_params.py index 6c476b78359b..16cfb7d38441 100644 --- a/tensorrt_llm/disaggregated_params.py +++ b/tensorrt_llm/disaggregated_params.py @@ -6,10 +6,10 @@ @dataclass(slots=True, kw_only=True) class DisaggregatedParams: - """Disaggregated seving parameters. + """Disaggregated serving parameters. Args: - request_type (str): The type of request ("context_only" or "generation_only") + request_type (str): The type of request ("context_only" | "generation_only" | "context_and_generation") first_gen_tokens (List[int]): The first tokens of the generation request ctx_request_id (int): The context request id opaque_state(bytes): Any additional state needing to be exchanged between context and gen instances diff --git a/tensorrt_llm/executor/request.py b/tensorrt_llm/executor/request.py index 886831d0723a..52e3d8773e1e 100644 --- a/tensorrt_llm/executor/request.py +++ b/tensorrt_llm/executor/request.py @@ -25,10 +25,15 @@ class LoRARequest: lora_name: str lora_int_id: int lora_path: str = "" + lora_ckpt_source: str = "hf" def __post_init__(self): if self.lora_path is not None and not os.path.exists(self.lora_path): raise ValueError(f"lora_path ({self.lora_path}) does not exist.") + if self.lora_ckpt_source not in ["hf", "nemo"]: + raise ValueError( + f"lora_ckpt_source must be 'hf' or 'nemo', got '{self.lora_ckpt_source}'" + ) @property def adapter_id(self): @@ -42,6 +47,10 @@ def name(self): def path(self): return self.lora_path + @property + def ckpt_source(self): + return self.lora_ckpt_source + @dataclass(slots=True) class PromptAdapterRequest: diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 9cd539f33b34..679c5793fe43 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -228,6 +228,10 @@ def _handle_sequence(self, output.logprobs = response_tensors.log_probs[src_idx] # overcome some WAR in the cpp executor if finish_reasons[src_idx] != tllm.FinishReason.CANCELLED: + if len(output.logprobs) > output.length: + # LlmResult holds a reference to LogProbStorage, which may be updated by the worker before the result is serialized. + # Therefore, we treat extra logprobs/logits as expected and only consume what's needed. + output.logprobs = output.logprobs[:output.length] assert len(output.logprobs) == output.length if response_tensors.generation_logits is not None: output.generation_logits = response_tensors.generation_logits[ diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index da90fc8fe931..6ebd7adc03de 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -150,13 +150,23 @@ def _create_engine(): self._runtime_model_config = _engine_config_to_model_config( engine_config) if engine_config.build_config.plugin_config.lora_plugin: - self._lora_manager = LoraManager() + # 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(cpp_peft_cache_manager=None) if engine_config.build_config.max_prompt_embedding_table_size > 0: self._prompt_adapter_manager = PromptAdapterManager() if getattr(executor_config, "backend", "") == "pytorch" and lora_config is not None: - self._lora_manager = LoraManager() + from tensorrt_llm._torch.pyexecutor.resource_manager import \ + ResourceManagerType + peft_cache_manager = self.engine.resource_manager.resource_managers.get( + ResourceManagerType.PEFT_CACHE_MANAGER) + self._lora_manager = LoraManager( + cpp_peft_cache_manager=peft_cache_manager.impl) lora_model_config = self.engine.model_engine.lora_model_config assert lora_model_config is not None self._lora_model_config = lora_model_config @@ -349,7 +359,8 @@ def _load_lora_adapter(self, lora_request: LoRARequest) -> bool: model_config=self._runtime_model_config if self._runtime_model_config is not None else self._lora_model_config, runtime_mapping=None, - uids=[adapter_id]) + uids=[adapter_id], + ckpt_source=lora_request.ckpt_source) return adapter_id in newly_loaded_uids def _load_prompt_adapter(self, @@ -362,15 +373,16 @@ def _load_prompt_adapter(self, def _enqueue_request(self, request: GenerationRequest) -> int: assert request.id is not None if self._lora_manager is not None and request.lora_request is not None: - loaded_new_lora_adapter = self._load_lora_adapter( - request.lora_request) + adapter_in_cache = self._lora_manager.is_adapter_in_cpu_cache( + request.lora_request.adapter_id) + self._load_lora_adapter(request.lora_request) uid = str(request.lora_request.adapter_id) lora_config = tllm.LoraConfig( task_id=request.lora_request.adapter_id, weights=self._lora_manager.cpp_lora_weights[uid] - if loaded_new_lora_adapter else None, + if not adapter_in_cache else None, config=self._lora_manager.cpp_lora_config[uid] - if loaded_new_lora_adapter else None) + if not adapter_in_cache else None) else: lora_config = None @@ -406,6 +418,10 @@ def _enqueue_request(self, request: GenerationRequest) -> int: context_phase_params = None request_type = tllm.RequestType.REQUEST_TYPE_CONTEXT_AND_GENERATION if request.disaggregated_params is not None: + assert ( + not self._is_pytorch_backend + or self.engine.kv_cache_transceiver is not None + ), "kv_cache_transceiver is disabled, please set 'cache_transceiver_config: backend:` in config file for disaggregated serving" request_type = request.disaggregated_params.get_request_type() if request_type == tllm.RequestType.REQUEST_TYPE_GENERATION_ONLY: context_phase_params = request.disaggregated_params.get_context_phase_params( @@ -567,6 +583,12 @@ def shutdown(self): self.engine.shutdown() self.engine = None + if hasattr( + self._executor_config, "checkpoint_loader" + ) and self._executor_config.checkpoint_loader is not None: + self._executor_config.checkpoint_loader.cleanup() + self._executor_config.checkpoint_loader = None + # Check if there are any errors from the threads before shutdown. self._handle_background_error() diff --git a/tensorrt_llm/functional.py b/tensorrt_llm/functional.py index 52c96d40f59f..0532b995c595 100755 --- a/tensorrt_llm/functional.py +++ b/tensorrt_llm/functional.py @@ -4834,6 +4834,48 @@ def _compute_sinusoidal_positions(scale_factors, is_short, True), _compute_sinusoidal_positions( scaling_long_factors, False, True), short_mscale + @staticmethod + def create_sinusoidal_positions_long_rope_for_attention_plugin( + num_pos: int, + dim: int, + theta: float, + original_max_pos: int, + short_factor: List[float], + long_factor: List[float], + dtype=np.float32): + 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)) + + # Short part + inv_freq_short = inv_freq / short_factor + t_short = np.arange(np.min([num_pos, original_max_pos]), + dtype=np.float32) + freqs_short = np.einsum("i,j->ij", t_short, inv_freq_short) + + # Long part + inv_freq_long = inv_freq / long_factor + t_long = np.arange(np.max([0, num_pos - original_max_pos]), + dtype=np.float32) + original_max_pos + freqs_long = np.einsum("i,j->ij", t_long, inv_freq_long) + + freqs = np.concatenate([freqs_short, freqs_long], axis=0) + sinusoid_inp = freqs.astype(np.float32)[..., np.newaxis] + + # Apply scaling + scale = num_pos / original_max_pos + scaling_factor = np.sqrt(1.0 + np.log(scale) / np.log(original_max_pos)) + + # 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) diff --git a/tensorrt_llm/inputs/__init__.py b/tensorrt_llm/inputs/__init__.py index 34b719885d50..a20978cab430 100644 --- a/tensorrt_llm/inputs/__init__.py +++ b/tensorrt_llm/inputs/__init__.py @@ -5,8 +5,9 @@ register_input_processor) from .utils import (ALL_SUPPORTED_MULTIMODAL_MODELS, ConversationMessage, MultimodalData, MultimodalDataTracker, - add_multimodal_placeholders, async_load_image, - async_load_video, default_multimodal_input_loader, + add_multimodal_placeholders, async_load_audio, + async_load_image, async_load_video, + default_multimodal_input_loader, encode_base64_content_from_url, load_image, load_video) __all__ = [ @@ -24,6 +25,7 @@ "MultimodalDataTracker", "MultimodalData", "MultimodalInput", + "async_load_audio", "async_load_image", "async_load_video", "add_multimodal_placeholders", diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index a6b29a9f0183..19d55ae77448 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -82,6 +82,72 @@ def to_tensor(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: torch.tensor(self.multimodal_lengths, dtype=torch.int32)) +@dataclass +class MultimodalRuntimeData: + """Runtime data for tracking multimodal token caching and reuse per request sequence. + + This class tracks which multimodal tokens are cached vs. need to be processed + for each request sequence during KV cache reuse scenarios. + + Attributes: + num_cached_tokens: Total number of cached tokens for this sequence + mm_token_lengths: Length of each multimodal token chunk + mm_token_positions: Starting positions of each multimodal token chunk + prompt_tokens: Current iteration of prompt tokens for this sequence (optional). Need it for chunk prefill if enabled (#TODO) + num_cached_mm_tokens: Number of multimodal tokens that are cached in this iteration (computed) + total_mm_tokens: Total number of multimodal tokens in this sequence (computed) + """ + num_cached_tokens: int + mm_token_lengths: List[int] + mm_token_positions: List[int] + + # TODO: support chunk prefill for multimodal + # When chunk prefill is enabled, we need to pass the prompt tokens for current chunk and mask to find the included mm tokens + prompt_tokens: Optional[List[int]] = None + + num_cached_mm_tokens: Optional[int] = None + total_mm_tokens: Optional[int] = None + + def __post_init__(self): + # Validate input data + if len(self.mm_token_positions) != len(self.mm_token_lengths): + raise ValueError( + f"mm_token_positions ({len(self.mm_token_positions)}) and mm_token_lengths ({len(self.mm_token_lengths)}) must have the same length" + ) + + if self.num_cached_tokens < 0: + raise ValueError( + f"num_cached_tokens must be non-negative, got {self.num_cached_tokens}" + ) + + if any(length <= 0 for length in self.mm_token_lengths): + raise ValueError( + f"All mm_token_lengths must be positive, got {self.mm_token_lengths}" + ) + + if any(pos < 0 for pos in self.mm_token_positions): + raise ValueError( + f"All mm_token_positions must be non-negative, got {self.mm_token_positions}" + ) + + if self.num_cached_mm_tokens is None: + # Compute cached multimodal tokens based on positions and cached tokens + self.num_cached_mm_tokens = 0 + for pos, length in zip(self.mm_token_positions, + self.mm_token_lengths): + if pos + length <= self.num_cached_tokens: + self.num_cached_mm_tokens += length + elif pos < self.num_cached_tokens: + # Partial overlap - only count the cached portion + self.num_cached_mm_tokens += self.num_cached_tokens - pos + + if self.num_cached_mm_tokens > self.num_cached_tokens: + raise ValueError( + f"num_cached_mm_tokens ({self.num_cached_mm_tokens}) must be less than or equal to " + f"num_cached_tokens ({self.num_cached_tokens})") + self.total_mm_tokens = sum(self.mm_token_lengths) + + @dataclass class MultimodalParams: """Unified container for multimodal parameters. @@ -117,6 +183,7 @@ class MultimodalParams: multimodal_input: Optional[MultimodalInput] = None multimodal_data: Optional[Dict[str, Any]] = field(default_factory=dict) + multimodal_runtime: Optional[MultimodalRuntimeData] = None def __post_init__(self): """Ensure default values are properly set.""" diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index e6a939864edc..a4bf8570d0ae 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -5,12 +5,13 @@ from collections import defaultdict from io import BytesIO from pathlib import Path -from typing import Any, Coroutine, Dict, List, Optional, TypedDict, Union +from typing import Any, Coroutine, Dict, List, Optional, Tuple, TypedDict, Union from urllib.parse import urlparse import aiohttp import numpy as np import requests +import soundfile import torch from PIL import Image from torchvision.transforms import ToTensor @@ -44,7 +45,7 @@ def load_base64_image(parsed_url: str) -> Image.Image: def load_image(image: str, format: str = "pt", - device: str = "cuda") -> Union[Image.Image, torch.Tensor]: + device: str = "cpu") -> Union[Image.Image, torch.Tensor]: assert format in ["pt", "pil"], "format must be either Pytorch or PIL" parsed_url = urlparse(image) @@ -66,7 +67,7 @@ def load_image(image: str, async def async_load_image( image: str, format: str = "pt", - device: str = "cuda") -> Union[Image.Image, torch.Tensor]: + device: str = "cpu") -> Union[Image.Image, torch.Tensor]: assert format in ["pt", "pil"], "format must be either Pytorch or PIL" parsed_url = urlparse(image) @@ -91,7 +92,7 @@ def load_video( video: str, num_frames: int = 10, format: str = "pt", - device: str = "cuda") -> Union[List[Image.Image], List[torch.Tensor]]: + device: str = "cpu") -> Union[List[Image.Image], List[torch.Tensor]]: # Keep this import local to avoid importing cv2 if not needed import cv2 @@ -140,7 +141,7 @@ async def async_load_video( video: str, num_frames: int = 10, format: str = "pt", - device: str = "cuda") -> Union[List[Image.Image], List[torch.Tensor]]: + device: str = "cpu") -> Union[List[Image.Image], List[torch.Tensor]]: assert format in ["pt", "pil"], "format must be either Pytorch or PIL" parsed_url = urlparse(video) @@ -159,6 +160,35 @@ async def async_load_video( return load_video(video_path, num_frames, format, device) +def load_audio( + audio: str, + format: str = "pt", + device: str = "cuda", +) -> Tuple[np.ndarray, int]: + parsed_url = urlparse(audio) + if parsed_url.scheme in ["http", "https"]: + audio = requests.get(audio, stream=True, timeout=10) + audio = BytesIO(audio.content) + + audio = soundfile.read(audio) + return audio + + +async def async_load_audio( + audio: str, + format: str = "pt", + device: str = "cuda", +) -> Tuple[np.ndarray, int]: + parsed_url = urlparse(audio) + if parsed_url.scheme in ["http", "https"]: + async with aiohttp.ClientSession() as session: + async with session.get(audio) as response: + audio = BytesIO(await response.content.read()) + + audio = soundfile.read(audio) + return audio + + # Copied from https://github.com/vllm-project/vllm/blob/main/examples/online_serving/openai_chat_completion_client_for_multimodal.py#L38 def encode_base64_content_from_url(content_url: str) -> str: """Encode a content retrieved from a remote url to base64 format.""" @@ -186,19 +216,24 @@ def encode_base64_content_from_url(content_url: str) -> str: SUPPORTED_LLAVA_VIDEO_MODEL_GROUP = ["llava_llama"] SUPPORTED_MISTRAL_IMAGE_MODEL_GROUP = ["mistral3"] SUPPORTED_HYPERCLOVAX_MODEL_GROUP = ["hyperclovax_vlm"] +SUPPORTED_PHI_MODEL_GROUP = ["phi4mm"] ALL_SUPPORTED_IMAGE_MODELS = SUPPORTED_QWEN_MODEL_GROUP \ + SUPPORTED_LLAMA_MODEL_GROUP \ + SUPPORTED_LLAVA_IMAGE_MODEL_GROUP \ + SUPPORTED_HYPERCLOVAX_MODEL_GROUP \ + SUPPORTED_GEMMA_MODEL_GROUP \ - + SUPPORTED_MISTRAL_IMAGE_MODEL_GROUP + + SUPPORTED_MISTRAL_IMAGE_MODEL_GROUP \ + + SUPPORTED_PHI_MODEL_GROUP ALL_SUPPORTED_VIDEO_MODELS = SUPPORTED_QWEN_MODEL_GROUP \ + SUPPORTED_LLAVA_VIDEO_MODEL_GROUP +ALL_SUPPORTED_AUDIO_MODELS = SUPPORTED_PHI_MODEL_GROUP + ALL_SUPPORTED_MULTIMODAL_MODELS = list(set(ALL_SUPPORTED_IMAGE_MODELS) \ - | set(ALL_SUPPORTED_VIDEO_MODELS)) + | set(ALL_SUPPORTED_VIDEO_MODELS) \ + | set(ALL_SUPPORTED_AUDIO_MODELS)) HF_CHAT_TEMPLATE_EXCEPTIONS = ["llava_llama"] PLACEHOLDER_EXCEPTIONS = ["llava_next"] @@ -223,6 +258,7 @@ class MultimodalPlaceholderPlacement(enum.Enum): # Ref: https://github.com/mistralai/mistral-common/blob/039465db2bdc0486df36365c9bdb428188482a18/ # src/mistral_common/tokens/tokenizers/base.py#L326 "mistral3": MultimodalPlaceholderPlacement.AFTER_TEXT, + "phi4mm": MultimodalPlaceholderPlacement.BEFORE_TEXT, } assert len(PLACEHOLDER_PLACEMENT_MAP) == len(ALL_SUPPORTED_MULTIMODAL_MODELS) @@ -235,7 +271,7 @@ def retrieve_multimodal_placeholder(model_type: str, modality: str, Args: model_type: The type of the multimodal model. modality: The modality of the data. - current_count: The number of multimodal data already added. Currently not used. + current_count: The number of multimodal data already added. """ @@ -257,6 +293,8 @@ def retrieve_multimodal_placeholder(model_type: str, modality: str, # Ref: https://github.com/mistralai/mistral-common/blob/26a6bb3a07ee0b78a3808f2797f23e1d28514b93/ # src/mistral_common/tokens/tokenizers/base.py#L60 return "[IMG]" + elif model_type in SUPPORTED_PHI_MODEL_GROUP: + return f"<|image_{current_count}|>" raise TypeError( f"For image modality, only {ALL_SUPPORTED_IMAGE_MODELS} are supported but got {model_type}" ) @@ -268,6 +306,9 @@ def retrieve_multimodal_placeholder(model_type: str, modality: str, raise TypeError( f"For video modality, only {ALL_SUPPORTED_VIDEO_MODELS} are supported but got {model_type}" ) + elif modality == "audio": + if model_type in SUPPORTED_PHI_MODEL_GROUP: + return f"<|audio_{current_count}|>" raise TypeError(f"Unknown modality: {modality}") @@ -343,7 +384,10 @@ def add_multimodal_placeholders(model_type: str, text_prompt: str, case MultimodalPlaceholderPlacement.AFTER_TEXT: parts.append(text_prompt) parts.extend(placeholders) - return "\n".join(parts) + if model_type == "phi4mm": + return "".join(parts) + else: + return "\n".join(parts) def resolve_hf_chat_template( @@ -436,7 +480,7 @@ def default_multimodal_input_loader( media: Union[List[str], List[List[str]]], image_data_format: str = "pt", num_frames: int = 8, - device: str = "cuda") -> List[dict[str, Union[str, torch.Tensor]]]: + device: str = "cpu") -> List[dict[str, Union[str, torch.Tensor]]]: def convert_to_conversation_message(prompt: str, media: Union[str, List[str]], @@ -458,6 +502,34 @@ def convert_to_conversation_message(prompt: str, media: Union[str, format=image_data_format, device=device)) for i in media ] + elif modality == "audio": + mm_data = [ + MultimodalData(modality=modality, + data=load_audio(i, device=device)) for i in media + ] + elif modality == "image_audio": + # Use different load_xxx functions to match the modality. + mm_data = [] + for m in media: + data = None + _modal = None + if _modal is None: + try: + data = load_image(m, + format=image_data_format, + device=device) + _modal = "image" + except Exception: + pass + if _modal is None: + try: + data = load_audio(m, device=device) + _modal = "audio" + except Exception: + pass + if _modal is None: + raise ValueError(f"Unknown matching modality: {modality}") + mm_data.append(MultimodalData(modality=_modal, data=data)) else: raise ValueError(f"Unknown modality: {modality}") return ConversationMessage(role="user", content=prompt, media=mm_data) diff --git a/tensorrt_llm/llmapi/disagg_utils.py b/tensorrt_llm/llmapi/disagg_utils.py index 42cff0b06018..f929c701fe4c 100644 --- a/tensorrt_llm/llmapi/disagg_utils.py +++ b/tensorrt_llm/llmapi/disagg_utils.py @@ -50,6 +50,7 @@ class DisaggServerConfig(): ctx_router_config: Optional[RouterConfig] = None gen_router_config: Optional[RouterConfig] = None conditional_disagg_config: Optional[ConditionalDisaggConfig] = None + max_retries: int = 3 @dataclass @@ -74,6 +75,7 @@ def parse_disagg_config_file(yaml_config_file: str): def extract_disagg_cfg(hostname: str = 'localhost', port: int = 8000, + max_retries: int = 3, context_servers: Optional[dict] = None, generation_servers: Optional[dict] = None, conditional_disagg_config: Optional[dict] = None, @@ -112,7 +114,7 @@ def extract_disagg_cfg(hostname: str = 'localhost', config = DisaggServerConfig(server_configs, hostname, port, ctx_router_config, gen_router_config, - conditional_disagg_config) + conditional_disagg_config, max_retries) return config diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index d5d0c9350022..934813aa4c4c 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -334,9 +334,9 @@ 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 - if (disaggregated_params is not None - and disaggregated_params.request_type == "context_only" - and not self._on_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: sampling_params.max_tokens = 1 inputs = prompt_inputs(inputs) @@ -401,7 +401,8 @@ def generate_async( self._check_arguments( len(prompt_token_ids), len(query_token_ids) if query_token_ids is not None else 0, - sampling_params) + sampling_params, + is_gen_only=is_gen_only) if _postproc_params: _postproc_params.postproc_args.num_prompt_tokens = len( prompt_token_ids) @@ -529,7 +530,8 @@ def _prepare_sampling_params( return sampling_params def _check_arguments(self, prompt_len: int, query_len: int, - sampling_params: SamplingParams) -> None: + sampling_params: SamplingParams, + is_gen_only: bool) -> None: if self.args.backend in ["pytorch", "_autodeploy"]: # TODO: remove these checks after PyTorch backend @@ -542,13 +544,6 @@ def _check_arguments(self, prompt_len: int, query_len: int, raise ValueError( f"PyTorch backend currently only supports `logprobs=1`. Received `logprobs={sampling_params.logprobs}` (Top{sampling_params.logprobs} logprobs). Please set `logprobs=1` in `sampling_params` instead." ) - # Check prompt length and query length against max_num_tokens to filter illegal requests. - if self.args.backend == "pytorch" and not self.args.enable_chunked_prefill: - max_num_tokens = self.args.max_num_tokens - if max_num_tokens and prompt_len / self.args.parallel_config.cp_size + query_len > max_num_tokens: - raise ValueError( - f"The sum of prompt length ({prompt_len/self.args.parallel_config.cp_size}), query length ({query_len}) and max_tokens ({sampling_params.max_tokens}) should not exceed " - f"max_num_tokens ({max_num_tokens})") return build_config = self.args.build_config @@ -565,7 +560,7 @@ def _check_arguments(self, prompt_len: int, query_len: int, (sampling_params.max_tokens or 0) > max_seq_len): raise ValueError( f"The sum of prompt length ({prompt_len/self.args.parallel_config.cp_size}) and query length ({query_len}) max_tokens ({sampling_params.max_tokens}) should not exceed " - f"max_seq_len ({max_seq_len})") + f"max_seq_len ({build_config.max_seq_len})") if sampling_params.use_beam_search and sampling_params.best_of > build_config.max_beam_width: if sampling_params.n == sampling_params.best_of: @@ -970,7 +965,11 @@ def _build_model(self): speculative_config=self.args.speculative_config, hf_model_dir=self._hf_model_dir, max_input_len=self.args.max_input_len, - max_seq_len=max_seq_len) + max_seq_len=max_seq_len, + checkpoint_format=None if self.args.backend == "_autodeploy" else + self.args.checkpoint_format, + checkpoint_loader=None if self.args.backend == "_autodeploy" else + self.args.checkpoint_loader) # TODO: revisit gather_context_logits return_logits = self.args.gather_generation_logits diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1b385b6e8fc6..6614391b4520 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -30,6 +30,7 @@ # isort: off from ..bindings.executor import ( BatchingType as _BatchingType, + CacheTransceiverBackendType as _CacheTransceiverBackendType, CacheTransceiverConfig as _CacheTransceiverConfig, CapacitySchedulerPolicy as _CapacitySchedulerPolicy, ContextChunkingPolicy as _ContextChunkingPolicy, @@ -248,7 +249,6 @@ class _ModelFormatKind(Enum): class DecodingBaseConfig(BaseModel): max_draft_len: Optional[int] = None speculative_model_dir: Optional[Union[str, Path]] = None - num_extra_kv_tokens: int = 0 @classmethod def from_dict(cls, data: dict): @@ -295,13 +295,6 @@ def spec_dec_mode(self): return TorchSpeculativeDecodingMode.from_string( self.decoding_type.upper()) - def update_from_model_config(self, model_config): - pass - - def get_draft_model_prompt(self, - input_tokens: torch.Tensor) -> torch.Tensor: - return input_tokens - class MedusaDecodingConfig(DecodingBaseConfig): medusa_choices: Optional[List[List[int]]] = None @@ -345,17 +338,11 @@ def spec_dec_mode(self): return TorchSpeculativeDecodingMode.EAGLE3_ONE_MODEL return TorchSpeculativeDecodingMode.EAGLE3 - def get_draft_model_prompt(self, - input_tokens: torch.Tensor) -> torch.Tensor: - """ - Eagle3 always throws away the first token when processing draft inputs - """ - return input_tokens[1:] - class UserProvidedDecodingConfig(DecodingBaseConfig): - # Type should be Drafter, but it leads to circular import - drafter: object + # Cannot use real type annotations due to circular imports + drafter: object # Type is Drafter + resource_manager: object = None # Type is Optional[ResourceManager] @classmethod def from_dict(cls, data: dict): @@ -447,11 +434,6 @@ def spec_dec_mode(self): return TorchSpeculativeDecodingMode.MTP_EAGLE return TorchSpeculativeDecodingMode.MTP - def update_from_model_config(self, model_config): - assert self.num_nextn_predict_layers > 0 - if model_config.num_nextn_predict_layers == 1 and not self.use_mtp_vanilla: - self.num_extra_kv_tokens = self.num_nextn_predict_layers - 1 - class PybindMirror(ABC): ''' A class containing the utilities for mirroring Python classes to @@ -821,6 +803,10 @@ class KvCacheConfig(BaseModel, PybindMirror): use_uvm: bool = Field(default=False, description="Whether to use UVM for the KV cache.") + # This is a pure python field, not a pybind field. It is only for the Pytorch backend. + dtype: str = Field(default="auto", + description="The data type to use for the KV cache.") + def _to_pybind(self): return _KvCacheConfig( enable_block_reuse=self.enable_block_reuse, @@ -874,12 +860,20 @@ class CacheTransceiverConfig(BaseModel, PybindMirror): """ Configuration for the cache transceiver. """ - max_num_tokens: Optional[int] = Field( + + backend: Optional[Literal["default", "ucx", "nixl", "mpi"]] = Field( + default=None, + description= + "The communication backend type to use for the cache transceiver.") + + max_tokens_in_buffer: Optional[int] = Field( default=None, description="The max number of tokens the transfer buffer can fit.") def _to_pybind(self): - return _CacheTransceiverConfig(max_num_tokens=self.max_num_tokens) + return _CacheTransceiverConfig( + backend=_CacheTransceiverBackendType.from_string(self.backend), + max_tokens_in_buffer=self.max_tokens_in_buffer) @dataclass @@ -1024,10 +1018,6 @@ class BaseLlmArgs(BaseModel): lora_config: Optional[LoraConfig] = Field( default=None, description="LoRA configuration for the model.") - # Quantization and calibration configurations - quant_config: Optional[QuantConfig] = Field( - default=None, description="Quantization config.", validate_default=True) - # Several options from ExecutorConfig, expanded here for less hierarchy kv_cache_config: KvCacheConfig = Field(default_factory=KvCacheConfig, description="KV cache config.") @@ -1208,13 +1198,6 @@ def validate_dtype(cls, v, info): raise RuntimeError("Pre SM 80 GPUs do not support bfloat16") return v - @field_validator("quant_config", mode='before') - @classmethod - def validate_quant_config(cls, v, info): - if v is None: - v = QuantConfig() - return v - @field_validator("gpus_per_node", mode='before') @classmethod def validate_gpus_per_node(cls, v, info): @@ -1286,7 +1269,8 @@ def validate_model_format_misc(self): 'pytorch', '_autodeploy' ]: # Load parallel_config from the engine. - model_format = get_model_format(self.model) + 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: @@ -1354,6 +1338,15 @@ def set_runtime_knobs_from_build_config(self): return self + @model_validator(mode="after") + def validate_runtime_args(self): + if self.max_batch_size is not None and self.max_num_tokens is not None: + if self.max_batch_size > self.max_num_tokens: + logger.warning( + f"max_batch_size [{self.max_batch_size}] should be less than or equal to max_num_tokens [{self.max_num_tokens}]" + ) + return self + @model_validator(mode="after") def validate_build_config_with_runtime_params(self): # Note: max_batch_size and max_num_tokens in LlmArgs are for runtime, @@ -1456,8 +1449,6 @@ def validate_speculative_config(self): assert self.speculative_config.speculative_model_dir is not None, "Path to EAGLE3 weights must be specified." self.build_config.max_draft_len = self.speculative_config.max_draft_len self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.EAGLE - if self.speculative_config.eagle3_one_model: - self.speculative_config.num_extra_kv_tokens = self.speculative_config.max_draft_len - 1 if self.backend not in ['pytorch', '_autodeploy']: eagle_config = _EagleConfig( self.speculative_config.eagle_choices, @@ -1478,6 +1469,7 @@ def validate_speculative_config(self): elif isinstance(self.speculative_config, DraftTargetDecodingConfig): assert self.backend in ['pytorch'] assert self.speculative_config.max_draft_len > 0 + assert self.speculative_config.speculative_model_dir is not None, "Path to draft model must be specified." self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.DRAFT_TOKENS_EXTERNAL self.build_config.max_draft_len = self.speculative_config.max_draft_len @@ -1656,6 +1648,10 @@ class TrtLlmArgs(BaseLlmArgs): calib_config: Optional[CalibConfig] = Field( default=None, description="Calibration config.", validate_default=True) + # Quantization and calibration configurations + quant_config: Optional[QuantConfig] = Field( + default=None, description="Quantization config.", validate_default=True) + embedding_parallel_mode: str = Field( default='SHARDING_ALONG_VOCAB', description="The embedding parallel mode.") @@ -1693,6 +1689,13 @@ def init_calib_config(cls, v): return CalibConfig() return v + @field_validator("quant_config", mode='before') + @classmethod + def validate_quant_config(cls, v, info): + if v is None: + v = QuantConfig() + return v + @model_validator(mode="after") def setup_embedding_parallel_mode(self): if self.embedding_parallel_mode == 'NONE': @@ -1737,6 +1740,11 @@ def validate_enable_build_cache(self): f"Invalid build_cache_config: {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 @@ -1764,6 +1772,20 @@ class TorchCompileConfig(BaseModel): description= "When torch compile is enabled, userbuffers is enabled by default.") + max_num_streams: int = Field( + default=1, + description= + "The maximum number of CUDA streams to use for torch.compile.") + + @field_validator('max_num_streams') + @classmethod + def validate_torch_compile_max_num_streams(cls, v): + """Validate torch_compile_config.max_num_streams >= 1.""" + if v < 1: + raise ValueError( + "torch_compile_config.max_num_streams must be >= 1") + return v + class TorchLlmArgs(BaseLlmArgs): # Just a dummy BuildConfig to allow code reuse with the TrtLlmArgs @@ -1810,9 +1832,6 @@ class TorchLlmArgs(BaseLlmArgs): "If true, will use the TRTLLM sampler instead of the PyTorch sampler. The TRTLLM sampler has a wide coverage of sampling strategies." ) - kv_cache_dtype: str = Field(default="auto", - description="Data type for KV cache.") - enable_iter_perf_stats: bool = Field( default=False, description="Enable iteration performance statistics.") @@ -1865,6 +1884,31 @@ class TorchLlmArgs(BaseLlmArgs): 'LOWPRECISION', 'MNNVL']] = Field(default='AUTO', description="Allreduce strategy to use.") + checkpoint_loader: Optional[object] = Field( + default=None, + description="The checkpoint loader to use for this LLM instance.", + json_schema_extra={ + "type": "Optional[tensorrt_llm._torch.BaseCheckpointLoader]" + }, + ) + + checkpoint_format: Optional[str] = Field( + default=None, + description="The format of the provided checkpoint.", + ) + + # PrivateVars + _quant_config: Optional[QuantConfig] = PrivateAttr(default=None) + + @property + def quant_config(self) -> QuantConfig: + if self._quant_config is None: + self._quant_config = QuantConfig() + return self._quant_config + + @quant_config.setter + def quant_config(self, value: QuantConfig): + self._quant_config = value # TODO: remove backend later @field_validator('backend', mode='before') @@ -1906,6 +1950,22 @@ def validate_stream_interval(self): f"stream_interval must be positive, got {self.stream_interval}") return self + @model_validator(mode="after") + def validate_checkpoint_format(self): + if self.checkpoint_format is not None and self.checkpoint_loader is not None: + logger.warning( + "checkpoint_format and checkpoint_loader are both provided, " + "checkpoint_loader will be ignored.") + self.checkpoint_loader = None + + if self.checkpoint_format is None and self.checkpoint_loader is None: + logger.info( + "neither checkpoint_format nor checkpoint_loader were provided, " + "checkpoint_format will be set to HF.") + self.checkpoint_format = "HF" + + return self + @staticmethod def _generate_cuda_graph_batch_sizes(max_batch_size: int, enable_padding: bool) -> List[int]: @@ -1993,6 +2053,22 @@ def validate_cuda_graph_config(self) -> 'TorchLlmArgs': return self + @model_validator(mode='after') + def sync_quant_config_with_kv_cache_config_dtype(self) -> 'TorchLlmArgs': + if self.kv_cache_config is None: + return self + + assert self.quant_config is not None + if self.kv_cache_config.dtype == "auto": + return self + elif self.kv_cache_config.dtype == 'fp8': + self.quant_config.kv_cache_quant_algo = QuantAlgo.FP8 + else: + logger.warning( + f"Cannot sync quant_config.kv_cache_quant_algo with kv_cache_config.dtype of {self.kv_cache_config.dtype}, " + "please update the validator") + return self + # TODO: Remove this after the PyTorch backend is fully migrated to TorchLlmArgs from ExecutorConfig def get_pytorch_backend_config(self) -> "PyTorchConfig": from tensorrt_llm._torch.pyexecutor.config import PyTorchConfig @@ -2016,7 +2092,7 @@ def get_pytorch_backend_config(self) -> "PyTorchConfig": moe_backend=self.moe_config.backend, enable_mixed_sampler=self.enable_mixed_sampler, enable_trtllm_sampler=self.enable_trtllm_sampler, - kv_cache_dtype=self.kv_cache_dtype, + kv_cache_dtype=self.kv_cache_config.dtype, enable_iter_perf_stats=self.enable_iter_perf_stats, enable_iter_req_stats=self.enable_iter_req_stats, print_iter_log=self.print_iter_log, @@ -2034,6 +2110,9 @@ def get_pytorch_backend_config(self) -> "PyTorchConfig": torch_compile_enable_userbuffers=self.torch_compile_config. enable_userbuffers if self.torch_compile_config is not None else TorchCompileConfig.model_fields['enable_userbuffers'].default, + torch_compile_max_num_streams=self.torch_compile_config. + max_num_streams if self.torch_compile_config is not None else + TorchCompileConfig.model_fields['max_num_streams'].default, enable_autotuner=self.enable_autotuner, enable_layerwise_nvtx_marker=self.enable_layerwise_nvtx_marker, load_format=self.load_format, @@ -2060,7 +2139,8 @@ def update_llm_args_with_extra_dict( } for field_name, field_type in field_mapping.items(): if field_name in llm_args_dict: - if field_name == "speculative_config": + # Some fields need to be converted manually. + if field_name in ["speculative_config", "build_config"]: llm_args_dict[field_name] = field_type.from_dict( llm_args_dict[field_name]) else: @@ -2083,7 +2163,8 @@ def update_llm_args_with_extra_options(llm_args: Dict, return llm_args -def get_model_format(model_dir: str) -> _ModelFormatKind: +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( @@ -2102,7 +2183,8 @@ def get_model_format(model_dir: str) -> _ModelFormatKind: PretrainedConfig.from_checkpoint(model_dir) else: model_format = _ModelFormatKind.HF - AutoConfig.from_hugging_face(model_dir) + 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}" diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index cf2bdb26c143..a62568a54e86 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -5,7 +5,6 @@ import tempfile import time import weakref -from argparse import Namespace from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Callable, List, Optional, Tuple, Union @@ -35,7 +34,7 @@ LookaheadDecodingConfig, MedusaDecodingConfig, MTPDecodingConfig, NGramDecodingConfig, UserProvidedDecodingConfig, _ModelFormatKind, - _ModelWrapper, _ParallelConfig, get_model_format, + _ModelWrapper, _ParallelConfig, update_llm_args_with_extra_dict, update_llm_args_with_extra_options) from .mpi_session import MPINodeState, MpiSession @@ -315,11 +314,6 @@ def save( if tokenizer is not None: tokenizer.save_pretrained(engine_dir) - @staticmethod - def get_model_format(model_dir: str) -> _ModelFormatKind: - ''' Get the format of the model. ''' - return get_model_format(model_dir) - def _download_hf_model(self): ''' Download HF model from third-party model hub like www.modelscope.cn or huggingface. ''' model_dir = None @@ -368,7 +362,11 @@ def _update_from_hf_quant_config(self) -> bool: hf_quant_algo = hf_quant_config.pop("quant_algo", None) if hf_quant_algo is not None: - hf_quant_algo = QuantAlgo(hf_quant_algo) + # fp8_pb_wo from modelopt is the same as fp8_block_scales + if hf_quant_algo == "fp8_pb_wo": + hf_quant_algo = QuantAlgo.FP8_BLOCK_SCALES + else: + hf_quant_algo = QuantAlgo(hf_quant_algo) if quant_config.quant_algo is None: logger.info( f"Setting quant_algo={hf_quant_algo} form HF quant config." @@ -407,6 +405,9 @@ def _update_from_hf_quant_config(self) -> bool: logger.info(f"Setting {key}={value} from HF quant config.") setattr(quant_config, key, value) + # Update the quant_config in llm_args for pytorch + self.llm_args.quant_config = quant_config + return True hf_config_path = f"{self._model_dir}/config.json" @@ -566,21 +567,6 @@ def _load_engine_buffer(self): # Load engine buffer from disk self._engine = Engine.from_dir(self._model_dir) - @staticmethod - def load_extra_build_configs_from_engine( - model_dir: str) -> Optional[Namespace]: - ''' Load the extra build configs from the engine directory, return None if model isn't an engine. ''' - if ModelLoader.get_model_format( - model_dir) is not _ModelFormatKind.TLLM_ENGINE: - return None - - with open(Path(model_dir) / "config.json", "r") as f: - engine_config = json.load(f) - - build_config = engine_config['build_config'] - build_config.pop("plugin_config") - return Namespace(**build_config) - @staticmethod def load_hf_tokenizer( model_dir, @@ -740,7 +726,8 @@ def get_pretrained_config(self) -> PretrainedConfig: self._hf_model_dir, mapping=self.llm_args.parallel_config.to_mapping(), quant_config=self.llm_args.quant_config, - dtype=self.llm_args.dtype) + 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 diff --git a/tensorrt_llm/lora_manager.py b/tensorrt_llm/lora_manager.py index 3c40917a194a..9f42fdad20db 100644 --- a/tensorrt_llm/lora_manager.py +++ b/tensorrt_llm/lora_manager.py @@ -4,13 +4,16 @@ import tarfile from collections import defaultdict from dataclasses import dataclass, field +from functools import lru_cache from pathlib import Path -from typing import TYPE_CHECKING, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union import numpy as np import torch import yaml +from tensorrt_llm.bindings import internal as tb_internal + from ._utils import DictConversion, pad_vocab_size, release_gc, str_dtype_to_torch, torch_to_numpy from .layers.linear import ColumnLinear from .mapping import Mapping @@ -20,8 +23,21 @@ from .runtime import ModelConfig -def get_all_nemo_lora_weights(lora_weights): - layer_weights = defaultdict(dict) +def get_all_nemo_lora_weights( + lora_weights: Dict[str, torch.Tensor], +) -> Dict[int, Dict[str, torch.Tensor]]: + """Extract and organize NeMo LoRA weights by layer and direction. + + Args: + lora_weights: Dictionary mapping weight keys to tensors from NeMo checkpoint + + Returns: + Dictionary mapping layer_idx -> {direction -> tensor} where direction is 'in' or 'out' + + Raises: + KeyError: If unsupported keys are found or layer extraction fails + """ + layer_weights: Dict[int, Dict[str, torch.Tensor]] = defaultdict(dict) adapter_key = "self_attention.adapter_layer.lora_kqv_adapter" layer_pattern = re.compile(r".*\.layers\.(\d+)\..*") for key, weights in lora_weights.items(): @@ -50,7 +66,28 @@ def get_all_nemo_lora_weights(lora_weights): ) -def iterate_hf_lora(iter_fn, lora_weights, hf_modules, component=None): +def iterate_hf_lora( + iter_fn, + lora_weights: Dict[str, torch.Tensor], + hf_modules: Set[str], + component: Optional[str] = None, +): + """Iterate over HuggingFace LoRA weights and call iterator function for each weight. + + Args: + iter_fn: Function to call for each weight with signature + (layer_idx, hf_module, expert_idx, inout_or_mag, weights) + lora_weights: Dictionary mapping weight keys to tensors from HF checkpoint + hf_modules: Set of supported HF module names + component: Optional component name to filter by (e.g., 'decoder') + + Returns: + Nested dictionary structure organizing the weights + + Raises: + KeyError: If unsupported keys are found + AssertionError: If HF module is not in supported list + """ all_weights = defaultdict(lambda: defaultdict(dict)) pattern = HF_LORA_PATTERN for key, weights in lora_weights.items(): @@ -94,7 +131,20 @@ def iterate_hf_lora(iter_fn, lora_weights, hf_modules, component=None): return all_weights -def get_all_hf_lora_weights(lora_weights, hf_modules, component=None): +def get_all_hf_lora_weights( + lora_weights: Dict[str, torch.Tensor], hf_modules: Set[str], component: Optional[str] = None +): + """Extract and organize all HuggingFace LoRA weights by layer and module. + + Args: + lora_weights: Dictionary mapping weight keys to tensors from HF checkpoint + hf_modules: Set of supported HF module names + component: Optional component name to filter by (e.g., 'decoder') + + Returns: + Nested dictionary organizing weights by layer, module, and potentially expert + """ + def iter_fn(layer_idx, hf_module, expert_idx, inout, weights): if expert_idx is None: all_weights[layer_idx][hf_module][inout] = weights @@ -116,8 +166,19 @@ def iter_fn(layer_idx, hf_module, expert_idx, inout, weights): return hf_target_modules -def invert_module_mapping(trtllm_modules_to_hf_modules): - hf_modules_to_trtllm_modules = {} +def invert_module_mapping( + trtllm_modules_to_hf_modules: Dict[str, Union[str, List[str]]], +) -> Dict[str, str]: + """Invert module mapping from TensorRT-LLM -> HF to HF -> TensorRT-LLM. + + Args: + trtllm_modules_to_hf_modules: Mapping from TensorRT-LLM module names to HF module names + (values can be strings or lists of strings) + + Returns: + Dictionary mapping HF module names to TensorRT-LLM module names + """ + hf_modules_to_trtllm_modules: Dict[str, str] = {} for k, hf_modules in trtllm_modules_to_hf_modules.items(): if isinstance(hf_modules, list): for hf_module in hf_modules: @@ -216,8 +277,88 @@ def get_target_modules(self, trtllm_modules_to_hf_modules): return list(lora_target_modules) +@lru_cache(maxsize=128) +def _find_nemo_files_single_path(lora_path: str) -> List[str]: + """Find .nemo files from a single path (file or directory). + + This function is cached per individual path to maximize cache efficiency + when the same paths appear in different collections. + + Args: + lora_path: A single path that can be either: + - Direct path to a .nemo file + - Directory containing .nemo files (will auto-detect *.nemo) + + Returns: + List[str]: List of paths to .nemo files found in this single path + + Raises: + ValueError: If path doesn't exist, no .nemo files found, or invalid file type + """ + path = Path(lora_path) + if not path.exists(): + raise ValueError(f"{path} does not exist") + + if path.is_file(): + if path.suffix == ".nemo": + return [str(path)] + else: + raise ValueError(f"{path} is not a .nemo file") + elif path.is_dir(): + nemo_files_in_dir = list(path.glob("*.nemo")) + if not nemo_files_in_dir: + raise ValueError(f"No .nemo files found in directory {path}") + return [str(f) for f in nemo_files_in_dir] + else: + raise ValueError(f"{path} is neither a file nor a directory") + + +def find_nemo_files(lora_dirs: List[str]) -> List[str]: + """Find all .nemo files from a list of directories or file paths. + + This function is optimized for repeated calls at generation time by using an internal LRU cache + on individual paths, which maximizes cache efficiency when the same paths + appear in different collections. + + Args: + lora_dirs: List of paths that can be either: + - Direct paths to .nemo files + - Directories containing .nemo files (will auto-detect *.nemo) + + Returns: + List[str]: List of paths to .nemo files + + Raises: + ValueError: If a path doesn't exist, no .nemo files are found in a directory + path, or a file path is of invalid file type + """ + if len(lora_dirs) == 0: + return [] + + all_nemo_files: List[str] = [] + for lora_path in lora_dirs: + nemo_files_for_path = _find_nemo_files_single_path(lora_path) + all_nemo_files.extend(nemo_files_for_path) + + if not all_nemo_files: + raise ValueError("No .nemo files found in the provided paths") + + return all_nemo_files + + class NemoLoraLoader: def __init__(self, lora_dirs: List[str]): + """Initialize NemoLoraLoader with paths to .nemo files or directories. + + Args: + lora_dirs: List of paths that can be either: + - Direct paths to .nemo files + - Directories containing .nemo files (will auto-detect *.nemo) + + Note: The parameter name 'lora_dirs' is misleading - it can accept both + directories and files. This is a design flaw that should be fixed + in a future version (e.g., rename to 'lora_paths'). + """ self.lora_target_modules = [] self.is_valid = False @@ -228,15 +369,28 @@ def __init__(self, lora_dirs: List[str]): path = Path(lora_file) if not path.exists(): raise ValueError(f"{path} does not exist") - if not path.is_file(): - raise ValueError(f"{path} is not a file") self.is_valid = True # Hardcoded since LoraManager only supports this case now self.lora_target_modules = ["attn_qkv"] + def get_target_modules(self): + """Get target modules for NeMo LoRA. + + Unlike the HF loader, this method does not accept trtllm_modules_to_hf_modules + as an argument since the module mapping is hardcoded for NeMo LoRA support. + + Returns: + List[str]: List of target module names supported by NeMo LoRA + """ + 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 @@ -285,6 +439,73 @@ def load_torch_hf_lora(lora_config: LoraConfig): lora_config.lora_target_modules.extend(missing_qkv_modules) +def load_torch_nemo_lora(lora_config: LoraConfig): + """Load NeMo LoRA checkpoint for PyTorch workflow. + + This is a PyTorch-specific loader for NeMo LoRA checkpoints, similar to + load_torch_hf_lora but handling NeMo checkpoint format. NeMo uses a combined + "attn_qkv" module rather than separate Q, K, V modules, so no missing QKV + module handling is needed. + + Note: This function only sets up the configuration. For PyTorch workflow, + the actual weight loading happens later via LoraManager when requests are + made with LoRA UIDs. + + Args: + lora_config: LoRA configuration with lora_ckpt_source="nemo" + + Raises: + ValueError: If NeMo LoRA directory is invalid or unsupported modules are specified + """ + lora_config.trtllm_modules_to_hf_modules = {"attn_qkv": "attn_qkv"} + + assert len(lora_config.lora_dir) == 1, "Expecting only a single lora dir" + 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.get_target_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." + ) + + supported_modules = {"attn_qkv"} + unsupported_modules = set(lora_config.lora_target_modules) - supported_modules + if unsupported_modules: + raise ValueError( + f"NeMo LoRA only supports {supported_modules} modules, " + f"but got unsupported modules: {unsupported_modules}. " + f"NeMo LoRA does not support embedding, lm_head, or MLP adapters." + ) + + +def load_torch_lora(lora_config: LoraConfig): + """Load LoRA checkpoint for PyTorch workflow. + + This function routes to the appropriate loader based on lora_ckpt_source. + + Args: + lora_config: LoRA configuration with lora_ckpt_source set to "hf" or "nemo" + + Raises: + ValueError: If lora_ckpt_source is not supported + """ + if lora_config.lora_ckpt_source == "nemo": + load_torch_nemo_lora(lora_config) + elif lora_config.lora_ckpt_source == "hf": + load_torch_hf_lora(lora_config) + else: + raise ValueError( + f"Unsupported lora_ckpt_source: {lora_config.lora_ckpt_source}. " + f"Supported sources: 'hf', 'nemo'" + ) + + def load_hf_lora( model, lora_config: LoraConfig, @@ -386,7 +607,18 @@ def use_lora( raise ValueError(f"Unsupported lora_ckpt_source: {lora_config.lora_ckpt_source}") -def unpack_nemo_weights(nemo_archive_path): +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. + + Args: + nemo_archive_path: Path to the .nemo archive file + + Returns: + Tuple of (model_config_dict, model_weights_dict) + + Raises: + Exception: If required files cannot be extracted from the archive + """ with tarfile.open(nemo_archive_path) as tar: try: model_weights_file = tar.extractfile("model_weights.ckpt") @@ -436,8 +668,16 @@ class LoraManager(object): "mlp_gate_up": 18, } - def __init__(self): - """Constructor.""" + def __init__( + self, cpp_peft_cache_manager: tb_internal.batch_manager.PeftCacheManager | None = None + ): + """Constructor. + + Args: + cpp_peft_cache_manager (PeftCacheManager, optional): used by is_adapter_in_cpu_cache method, that's used for + a performance optimization with LoRA of not sending the LoRA adapter weights with every LLM request when + the adapter is already loaded in the LoRA CPU cache. + """ # _lora_uid_to_low_ranks: dict[str -> dict[int -> dict[str -> int]]] # { # uid: { @@ -473,6 +713,19 @@ def __init__(self): self._cpp_lora_weights: Dict[str, torch.Tensor] = {} # on cpu self._cpp_lora_config: Dict[str, torch.Tensor] = {} # on cpu self.lora_target_modules: List[str] = [] + self._cpp_peft_cache_manager = cpp_peft_cache_manager + + def is_adapter_in_cpu_cache(self, adapter_uid: int) -> bool: + """Best effort to check if a LoRA adapter is in the LoRA CPU cache. + + If no cpp_peft_cache_manager instance was given at the construction of this LoraManager instance, then False is + returned. + """ + return ( + self._cpp_peft_cache_manager.is_task_cached(adapter_uid) + if self._cpp_peft_cache_manager + else False + ) @staticmethod def get_missing_qkv_modules(lora_target_modules): @@ -516,8 +769,12 @@ def load_from_ckpt( uids=uids, ) elif ckpt_source == "nemo": + # Find all .nemo files from directories or files + nemo_files = find_nemo_files(model_dirs_or_files) + + # Pass the actual .nemo files to the loader return self.load_from_nemo( - model_files=model_dirs_or_files, + model_files=nemo_files, model_config=model_config, runtime_mapping=runtime_mapping, uids=uids, diff --git a/tensorrt_llm/models/automodel.py b/tensorrt_llm/models/automodel.py index c210e58e445e..463ae334cad6 100644 --- a/tensorrt_llm/models/automodel.py +++ b/tensorrt_llm/models/automodel.py @@ -16,9 +16,10 @@ def from_hugging_face(hf_model_or_dir, quant_config: Optional[QuantConfig] = None, **kwargs): import transformers + trust_remote_code = kwargs.get('trust_remote_code', False) hf_config = transformers.AutoConfig.from_pretrained( - hf_model_or_dir, trust_remote_code=True) + hf_model_or_dir, trust_remote_code=trust_remote_code) if hasattr(hf_config, 'architectures') and hf_config.architectures is not None: diff --git a/tensorrt_llm/models/phi3/convert.py b/tensorrt_llm/models/phi3/convert.py index ddd654dadd47..1f11408b6940 100644 --- a/tensorrt_llm/models/phi3/convert.py +++ b/tensorrt_llm/models/phi3/convert.py @@ -126,8 +126,9 @@ def get_moe_weight(key, suffix): if "qkv." in key: weights[key] = shuffle_qkv_weights(weights[key], config) - if config.architecture in ['Phi3SmallForCausalLM', "PhiMoEForCausalLM" - ] and config.mapping.has_tp(): + 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/split_weights.py b/tensorrt_llm/models/phi3/split_weights.py index 62a8891230c1..bca33ba55116 100644 --- a/tensorrt_llm/models/phi3/split_weights.py +++ b/tensorrt_llm/models/phi3/split_weights.py @@ -145,19 +145,23 @@ def get_quant_weight(weight, prefix, bias): -1, hidden_size) split_weight = torch.cat( [split(x, tp_size, rank) for x in [q, k, v]], dim=0) - - 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) + 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) - split_bias = split_qkv_bias_tp(qkv_bias, 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' @@ -171,7 +175,10 @@ def get_quant_weight(weight, prefix, bias): 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) - bias = split_matrix_tp(mlp_fc_bias, 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) diff --git a/tensorrt_llm/quantization/functional.py b/tensorrt_llm/quantization/functional.py index c467499372ec..84dc1b74a534 100644 --- a/tensorrt_llm/quantization/functional.py +++ b/tensorrt_llm/quantization/functional.py @@ -959,7 +959,7 @@ def preprocess_weights_for_mixed_gemm(tensor: torch.Tensor, tensor = tensor.unsqueeze(0) elif sm_ >= 90: sm_ = 80 - if sm_ >= 120: + if sm_ > 90: sm_ = 80 permutation_map = { diff --git a/tensorrt_llm/runtime/model_runner.py b/tensorrt_llm/runtime/model_runner.py index 486c58f6d151..a9f0fe8de409 100644 --- a/tensorrt_llm/runtime/model_runner.py +++ b/tensorrt_llm/runtime/model_runner.py @@ -86,7 +86,7 @@ def _builder_to_model_config(config: dict) -> Tuple[ModelConfig, dict]: dtype = builder_config['precision'] tp_size = builder_config['tensor_parallel'] pp_size = builder_config.get('pipeline_parallel', 1) - kv_cache_type = KVCacheType(builder_config.get('kv_cache_type')) + kv_cache_type = KVCacheType.from_string(builder_config.get('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()})' diff --git a/tensorrt_llm/runtime/multimodal_model_runner.py b/tensorrt_llm/runtime/multimodal_model_runner.py index 9d6be2bddad8..bb3a5480fcb6 100644 --- a/tensorrt_llm/runtime/multimodal_model_runner.py +++ b/tensorrt_llm/runtime/multimodal_model_runner.py @@ -2647,7 +2647,7 @@ def setup_inputs(self, input_text, raw_image, raw_audio=None): ) image = None elif self.model_type in ['llava_onevision']: - pre_prompt = "<|im_start|>user " + pre_prompt = "<|im_start|>user " + "