We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent 9f2dc30 commit 1b9c4ccCopy full SHA for 1b9c4cc
cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp
@@ -363,6 +363,18 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m)
363
}
364
return block_pool_pointers;
365
})
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
+ })
378
.def("get_layer_to_pool_mapping",
379
[](tbk::BaseKVCacheManager& self)
380
{
0 commit comments