Skip to content

Commit 1b9c4cc

Browse files
authored
[None][fix] Fix nanobind failure (#7425)
Signed-off-by: Tian Zheng <[email protected]>
1 parent 9f2dc30 commit 1b9c4cc

File tree

1 file changed

+12
-0
lines changed

1 file changed

+12
-0
lines changed

cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -363,6 +363,18 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m)
363363
}
364364
return block_pool_pointers;
365365
})
366+
.def("get_block_scale_pool_pointers",
367+
[](tbk::BaseKVCacheManager& self)
368+
{
369+
std::optional<at::Tensor> block_scale_pool_pointers{std::nullopt};
370+
auto tensor = self.getBlockScalePoolPointers();
371+
if (tensor)
372+
{
373+
std::shared_ptr<tensorrt_llm::runtime::ITensor> _tensor = std::move(tensor);
374+
block_scale_pool_pointers = tr::Torch::tensor(_tensor);
375+
}
376+
return block_scale_pool_pointers;
377+
})
366378
.def("get_layer_to_pool_mapping",
367379
[](tbk::BaseKVCacheManager& self)
368380
{

0 commit comments

Comments
 (0)