From 3897e38656d05bfc479d647560ec070eeac9ca17 Mon Sep 17 00:00:00 2001 From: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> Date: Mon, 4 Aug 2025 10:35:00 -0700 Subject: [PATCH 1/7] [None][feat] Enable nanobind as the default binding library Signed-off-by: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> --- cpp/tensorrt_llm/nanobind/runtime/bindings.cpp | 2 +- jenkins/L0_Test.groovy | 10 +++++----- scripts/build_wheel.py | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp b/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp index a22a62bf8086..47be92e13f5d 100644 --- a/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp @@ -279,7 +279,7 @@ void initBindings(nb::module_& m) .def(nb::init(), nb::arg("stream")) .def("setup", &tr::GptDecoderBatched::setup, nb::arg("mode"), nb::arg("max_num_sequences"), 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("forward_async", &tr::GptDecoderBatched::forwardAsync, nb::arg("decoder_state"), 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")) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 61ed4978d96e..c4740cd210b6 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -65,7 +65,7 @@ def LLVM_CONFIG = "LLVM" LINUX_AARCH64_CONFIG = "linux_aarch64" @Field -def NANOBIND_CONFIG = "Nanobind" +def PYBIND_CONFIG = "Pybind" @Field def BUILD_CONFIGS = [ @@ -74,7 +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"], + (PYBIND_CONFIG) : [(TARNAME) : "pybind-TensorRT-LLM.tar.gz"], ] // TODO: Move common variables to an unified location @@ -1775,7 +1775,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], + "A10-Pybind": ["a10", "l0_a10_pybind", 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], @@ -1852,8 +1852,8 @@ def launchTestJobs(pipeline, testFilter, dockerNode=null) if (key.contains("llvm")) { config = LLVM_CONFIG } - if (key.contains("Nanobind")) { - config = NANOBIND_CONFIG + if (key.contains("Pybind")) { + config = PYBIND_CONFIG } runLLMTestlistOnPlatform(pipeline, values[0], values[1], config, key.contains("Perf"), key, values[2], values[3]) }]]} diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index e40543c78f35..dbc57c0b4565 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -437,7 +437,7 @@ def main(*, install: bool = False, skip_building_wheel: bool = False, linking_install_binary: bool = False, - binding_type: str = "pybind", + binding_type: str = "nanobind", benchmarks: bool = False, micro_benchmarks: bool = False, nvtx: bool = False, @@ -984,8 +984,8 @@ def add_arguments(parser: ArgumentParser): ) parser.add_argument("--binding_type", choices=["pybind", "nanobind"], - default="pybind", - help="Which binding type to build: pybind or nanobind") + default="nanobind", + help="Which binding library to use: pybind or nanobind") parser.add_argument("--benchmarks", action="store_true", help="Build the benchmarks for the C++ runtime") From a81f869c8bc2f6cadb71cb17f551c0c70b403d0b Mon Sep 17 00:00:00 2001 From: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> Date: Wed, 6 Aug 2025 02:44:53 -0700 Subject: [PATCH 2/7] test fix and adding e2e tests Signed-off-by: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> --- cpp/tensorrt_llm/nanobind/executor/request.cpp | 1 + tests/integration/test_lists/test-db/l0_a10.yml | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/nanobind/executor/request.cpp b/cpp/tensorrt_llm/nanobind/executor/request.cpp index 1949474a10bf..e56341b53e22 100644 --- a/cpp/tensorrt_llm/nanobind/executor/request.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/request.cpp @@ -27,6 +27,7 @@ #include #include +#include #include #include #include diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index ce285faa7994..2175704a0c6c 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -211,6 +211,7 @@ l0_a10_nanobind: linux_distribution_name: ubuntu* terms: stage: pre_merge - backend: tensorrt tests: - unittest/bindings + - test_e2e.py::test_openai_chat_example[trt] + - test_e2e.py::test_openai_chat_example[pytorch] TIMEOUT (90) From a7a8b905971bc79a9c81b9d8c00a38d32b7e2f78 Mon Sep 17 00:00:00 2001 From: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> Date: Fri, 8 Aug 2025 09:50:50 -0700 Subject: [PATCH 3/7] Pybind build stage Signed-off-by: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> --- jenkins/Build.groovy | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/jenkins/Build.groovy b/jenkins/Build.groovy index d689de393bc6..3c4d8e4543c4 100644 --- a/jenkins/Build.groovy +++ b/jenkins/Build.groovy @@ -48,10 +48,10 @@ CONFIG_LINUX_AARCH64 = "linux_aarch64" def CONFIG_LINUX_AARCH64_LLVM = "linux_aarch64_LLVM" @Field -def CONFIG_LINUX_X86_64_NANOBIND = "linux_x86_64_Nanobind" +def CONFIG_LINUX_X86_64_PYBIND = "linux_x86_64_Pybind" @Field -def CONFIG_LINUX_AARCH64_NANOBIND = "linux_aarch64_Nanobind" +def CONFIG_LINUX_AARCH64_PYBIND = "linux_aarch64_Pybind" @Field def BUILD_CONFIGS = [ @@ -62,9 +62,9 @@ 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", + (CONFIG_LINUX_X86_64_PYBIND) : [ + (WHEEL_EXTRA_ARGS) : "--binding_type pybind --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) : "pybind-TensorRT-LLM.tar.gz", (WHEEL_ARCHS): "80-real;86-real;89-real;90-real;100-real;120-real", ], (CONFIG_LINUX_X86_64_SINGLE_DEVICE) : [ @@ -82,9 +82,9 @@ 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", + (CONFIG_LINUX_AARCH64_PYBIND): [ + (WHEEL_EXTRA_ARGS) : "--binding_type pybind --extra-cmake-vars WARNING_IS_ERROR=ON", + (TARNAME) : "pybind-TensorRT-LLM-GH200.tar.gz", (WHEEL_ARCHS): "90-real;100-real;120-real", ], (CONFIG_LINUX_AARCH64_LLVM) : [ @@ -542,8 +542,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), + "Build TRT-LLM Pybind": [LLM_DOCKER_IMAGE] + prepareLLMBuild( + pipeline, cpu_arch == AARCH64_TRIPLE ? CONFIG_LINUX_AARCH64_PYBIND : CONFIG_LINUX_X86_64_PYBIND), ] if (cpu_arch == X86_64_TRIPLE) { From 78aa0090dff42f145df0880db77821e7a4be337f Mon Sep 17 00:00:00 2001 From: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> Date: Tue, 12 Aug 2025 01:06:58 -0700 Subject: [PATCH 4/7] Fix KVCacheEvent for nanobind Signed-off-by: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> --- cpp/tensorrt_llm/nanobind/executor/bindings.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/nanobind/executor/bindings.cpp b/cpp/tensorrt_llm/nanobind/executor/bindings.cpp index d3f482df8997..ae4936a4df8e 100644 --- a/cpp/tensorrt_llm/nanobind/executor/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/bindings.cpp @@ -240,7 +240,8 @@ void initBindings(nb::module_& m) 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); + .def_ro("window_size", &tle::KVCacheEvent::windowSize) + .def_ro("attention_dp_rank", &tle::KVCacheEvent::attentionDpRank); nb::class_(executor_kv_cache, "KVCacheEventManager") .def( From df84fa2b846c14ebacfb4c2435cf4129b332dfdf Mon Sep 17 00:00:00 2001 From: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> Date: Tue, 19 Aug 2025 08:23:54 -0700 Subject: [PATCH 5/7] custom caster for scalarType and pybind test fix Signed-off-by: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> --- .../nanobind/common/customCasters.h | 30 +++++++++++++++++++ .../integration/test_lists/test-db/l0_a10.yml | 2 +- 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/nanobind/common/customCasters.h b/cpp/tensorrt_llm/nanobind/common/customCasters.h index a77a3bcb5a7b..432ce5c26b73 100644 --- a/cpp/tensorrt_llm/nanobind/common/customCasters.h +++ b/cpp/tensorrt_llm/nanobind/common/customCasters.h @@ -285,5 +285,35 @@ struct type_caster>> return make_caster>::from_cpp(result, policy, cleanup); } }; + +template <> +struct type_caster +{ + NB_TYPE_CASTER(torch::ScalarType, const_name("torch.dtype")); + + bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) noexcept + { + std::string dtype_name = nb::cast(nb::str(src)); + if (dtype_name.substr(0, 6) == "torch.") + { + dtype_name = dtype_name.substr(6); + } + + auto const& dtype_map = c10::getStringToDtypeMap(); + auto it = dtype_map.find(dtype_name); + if (it != dtype_map.end()) + { + value = it->second; + return true; + } + + return false; + } + + static handle from_cpp(torch::ScalarType src, rv_policy policy, cleanup_list* cleanup) + { + throw std::runtime_error("from_cpp for torch::ScalarType is not implemented"); + } +}; } // namespace detail } // namespace NB_NAMESPACE diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 2175704a0c6c..30fc6c05b58c 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -199,7 +199,7 @@ l0_a10: tests: - stress_test/stress_test.py::test_run_stress_test[llama-v3-8b-instruct-hf_tp1-stress_time_300s_timeout_450s-MAX_UTILIZATION-pytorch-stress-test] - stress_test/stress_test.py::test_run_stress_test[llama-v3-8b-instruct-hf_tp1-stress_time_300s_timeout_450s-GUARANTEED_NO_EVICT-pytorch-stress-test] -l0_a10_nanobind: +l0_a10_pybind: - condition: ranges: system_gpu_count: From 514e76d4c145c1d710a1974f23f668e782a50dea Mon Sep 17 00:00:00 2001 From: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> Date: Tue, 19 Aug 2025 11:35:48 -0700 Subject: [PATCH 6/7] Change cuda driver to variable Signed-off-by: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> --- cpp/tensorrt_llm/nanobind/CMakeLists.txt | 2 +- cpp/tensorrt_llm/pybind/CMakeLists.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/nanobind/CMakeLists.txt b/cpp/tensorrt_llm/nanobind/CMakeLists.txt index 1ccb50a02b7e..8c62584108ea 100755 --- a/cpp/tensorrt_llm/nanobind/CMakeLists.txt +++ b/cpp/tensorrt_llm/nanobind/CMakeLists.txt @@ -43,7 +43,7 @@ target_link_libraries( ${Python3_LIBRARIES} ${TORCH_LIBRARIES} torch_python - CUDA::cuda_driver + ${CUDA_DRV_LIB} ${CUDA_NVML_LIB} th_common) target_compile_definitions( diff --git a/cpp/tensorrt_llm/pybind/CMakeLists.txt b/cpp/tensorrt_llm/pybind/CMakeLists.txt index 91b5ebf54820..9d758b427cc6 100755 --- a/cpp/tensorrt_llm/pybind/CMakeLists.txt +++ b/cpp/tensorrt_llm/pybind/CMakeLists.txt @@ -44,7 +44,7 @@ target_link_libraries( ${Python3_LIBRARIES} ${TORCH_LIBRARIES} torch_python - CUDA::cuda_driver + ${CUDA_DRV_LIB} ${CUDA_NVML_LIB} th_common) target_compile_definitions( From 7f2fcdf97a71a9a68714fe57025c2112823bd032 Mon Sep 17 00:00:00 2001 From: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> Date: Wed, 20 Aug 2025 02:16:46 -0700 Subject: [PATCH 7/7] nanobind default in cmake Signed-off-by: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> --- cpp/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 4a8c8e9267ff..79c533fb5bfc 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -69,7 +69,7 @@ add_compile_definitions("TLLM_GEN_EXPORT_INTERFACE") add_compile_definitions("TLLM_ENABLE_CUDA") set(BINDING_TYPE - "pybind" + "nanobind" CACHE STRING "Binding type of Python bindings for C++ runtime and batch manager")