Skip to content

Commit 458f6d4

Browse files
committed
Fix nanobind error.
Signed-off-by: Shiyu Li <[email protected]>
1 parent 7ae720c commit 458f6d4

File tree

1 file changed

+5
-6
lines changed

1 file changed

+5
-6
lines changed

cpp/tensorrt_llm/nanobind/runtime/bindings.cpp

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -301,9 +301,8 @@ void initBindings(nb::module_& m)
301301
.def("finalize", &tr::GptDecoderBatched::finalize, nb::arg("decoder_state"), nb::arg("batch_idx"),
302302
nb::arg("sampling_config"), nb::arg("streaming"), nb::call_guard<nb::gil_scoped_release>())
303303
.def_prop_ro(
304-
"decoder_stream",
305-
[](tr::GptDecoderBatched& self) -> tr::CudaStream const& { return *self.getDecoderStream(); },
306-
nb::rv_policy::reference);
304+
"decoder_stream", [](tr::GptDecoderBatched& self) -> tr::CudaStream const&
305+
{ return *self.getDecoderStream(); }, nb::rv_policy::reference);
307306

308307
m.def(
309308
"lamport_initialize_all",
@@ -314,8 +313,7 @@ void initBindings(nb::module_& m)
314313
},
315314
"Lamport initialize all buffers", nb::call_guard<nb::gil_scoped_release>());
316315
m.def(
317-
"lamport_initialize",
318-
[](intptr_t buffer, size_t size)
316+
"lamport_initialize", [](intptr_t buffer, size_t size)
319317
{ tensorrt_llm::kernels::ar_fusion::lamport_initialize(reinterpret_cast<void*>(buffer), size, 0); },
320318
"Lmaport initialize buffer", nb::call_guard<nb::gil_scoped_release>());
321319
m.def(
@@ -361,7 +359,8 @@ void initBindings(nb::module_& m)
361359
nb::call_guard<nb::gil_scoped_release>());
362360

363361
nb::class_<tensorrt_llm::runtime::McastGPUBuffer>(m, "McastGPUBuffer")
364-
.def(nb::init<size_t, uint32_t, uint32_t, at::Device, bool>(), nb::call_guard<nb::gil_scoped_release>())
362+
.def(nb::init<size_t, uint32_t, uint32_t, uint32_t, at::Device, bool>(),
363+
nb::call_guard<nb::gil_scoped_release>())
365364
.def("get_uc_buffer", &tensorrt_llm::runtime::McastGPUBuffer::getUCBuffer,
366365
nb::call_guard<nb::gil_scoped_release>())
367366
.def("get_mc_buffer", &tensorrt_llm::runtime::McastGPUBuffer::getMCBuffer,

0 commit comments

Comments
 (0)