Kv cache manager refactor 1

.gitignore

@@ -5,6 +5,8 @@ __pycache__/
 *.cache
 *.nsys-rep
 *.npy
+*.so
+*.whl
 .VSCodeCounter
 cpp/build*
 cpp/Release

cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp

@@ -1,12 +1,16 @@
 #include "tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h"
 #include "tensorrt_llm/common/logger.h"
+#include "tensorrt_llm/common/memoryUtils.h"
 #include <cassert>
 #include <cstdio>
 #include <cuda.h>
 #include <fcntl.h>
 #include <unistd.h>
 #include <vector>
 
+namespace tc = tensorrt_llm::common;
+using namespace tensorrt_llm::runtime;
+
 namespace tensorrt_llm::batch_manager::kv_cache_manager_v2
 {
 
@@ -142,4 +146,30 @@ CUresult copyHostToHost(
     return cuLaunchHostFunc(stream, hostFnHostToHostCopy, data);
 }
 
+void copyBatchBlockOffsets(ITensor& output, SizeType32 batchSize, std::vector<BlockIndices> const& batchBlockIndices,
+    SizeType32 numPools, SizeType32 offset) noexcept
+{
+    auto* dstPtr = bufferCast<tk::KVCacheIndex::UnderlyingType>(output);
+    auto const& dstShape = output.getShape();
+
+    for (SizeType32 poolIdx = 0; poolIdx < numPools; poolIdx++)
+    {
+        for (SizeType32 batchIdx = 0; batchIdx < batchSize; batchIdx++)
+        {
+            auto const& blockIndices = batchBlockIndices[poolIdx * batchSize + batchIdx];
+            auto const addr = reinterpret_cast<tk::KVCacheIndex::UnderlyingType*>(blockIndices.addr);
+            auto const length = blockIndices.length;
+            auto const dstIndexK = tc::flat_index(dstShape.d, poolIdx, batchIdx + offset, 0, 0);
+            auto const dstIndexV = tc::flat_index(dstShape.d, poolIdx, batchIdx + offset, 1, 0);
+            auto addrK = dstPtr + dstIndexK;
+            auto addrV = dstPtr + dstIndexV;
+            for (SizeType32 i = 0; i < length; i++)
+            {
+                addrK[i] = addr[i];
+                addrV[i] = addr[i] + 1;
+            }
+        }
+    }
+}
+
 } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2

cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h

@@ -17,10 +17,17 @@
 
 #pragma once
 
+#include "tensorrt_llm/kernels/kvCacheIndex.h"
+#include "tensorrt_llm/runtime/iBuffer.h"
+#include "tensorrt_llm/runtime/iTensor.h"
 #include <cstdint>
 #include <cuda.h>
 #include <vector>
 
+namespace tk = tensorrt_llm::kernels;
+using SizeType32 = tensorrt_llm::runtime::SizeType32;
+using ITensor = tensorrt_llm::runtime::ITensor;
+
 namespace tensorrt_llm::batch_manager::kv_cache_manager_v2
 {
 struct DiskAddress
@@ -38,6 +45,12 @@ struct Task
     SrcAddr src;
 };
 
+struct BlockIndices
+{
+    MemAddress addr;
+    SizeType32 length;
+};
+
 CUresult copyDiskToDisk(
     std::vector<Task<DiskAddress, DiskAddress>> const& tasks, ssize_t numBytes, CUstream stream) noexcept;
 CUresult copyDiskToHost(
@@ -52,4 +65,8 @@ CUresult copyDeviceToHost(
     std::vector<Task<MemAddress, MemAddress>> const& tasks, ssize_t numBytes, CUstream stream) noexcept;
 CUresult copyDeviceToDevice(
     std::vector<Task<MemAddress, MemAddress>> const& tasks, ssize_t numBytes, CUstream stream) noexcept;
+
+void copyBatchBlockOffsets(ITensor& output, SizeType32 batchSize, std::vector<BlockIndices> const& batchBlockIndices,
+    SizeType32 numPools, SizeType32 offset) noexcept;
+
 } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2

cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp

@@ -17,14 +17,31 @@
 
 #include "kvCacheManagerV2Utils.h"
 #include "tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h"
+#include "tensorrt_llm/runtime/iTensor.h"
+#include "tensorrt_llm/runtime/torchView.h"
+#include <ATen/ATen.h>
 #include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
 #include <nanobind/stl/vector.h>
+#include <torch/extension.h>
 
+namespace tr = tensorrt_llm::runtime;
 namespace nb = nanobind;
 
+using SizeType32 = tensorrt_llm::runtime::SizeType32;
+
 namespace tensorrt_llm::batch_manager::kv_cache_manager_v2
 {
 
+std::optional<tensorrt_llm::runtime::ITensor::UniquePtr> from_torch(std::optional<at::Tensor> torchPtr)
+{
+    if (torchPtr)
+    {
+        return tr::TorchView::of(torchPtr.value());
+    }
+    return std::nullopt;
+}
+
 void KVCacheManagerV2UtilsBindings::initBindings(nb::module_& module)
 {
     // Bind DiskAddress struct
@@ -54,6 +71,11 @@ void KVCacheManagerV2UtilsBindings::initBindings(nb::module_& module)
         .def_rw("dst", &Task<MemAddress, MemAddress>::dst)
         .def_rw("src", &Task<MemAddress, MemAddress>::src);
 
+    nb::class_<BlockIndices>(module, "BlockIndices")
+        .def(nb::init<MemAddress, SizeType32>(), nb::arg("addr"), nb::arg("length"))
+        .def_rw("addr", &BlockIndices::addr)
+        .def_rw("length", &BlockIndices::length);
+
     // Bind copy functions
     module.def(
         "copy_disk_to_disk",
@@ -103,6 +125,18 @@ void KVCacheManagerV2UtilsBindings::initBindings(nb::module_& module)
         { return copyDeviceToDevice(tasks, numBytes, reinterpret_cast<CUstream>(stream)); },
         nb::arg("tasks"), nb::arg("num_bytes"), nb::arg("stream"), nb::call_guard<nb::gil_scoped_release>(),
         "Copy data from device to device using CUDA kernels");
+
+    module.def(
+        "copy_batch_block_offsets",
+        [](at::Tensor output, SizeType32 batchSize, std::vector<BlockIndices> const& batchBlockIndices,
+            SizeType32 numPools, SizeType32 offset)
+        {
+            auto _output = from_torch(output);
+            TLLM_CHECK_WITH_INFO(_output.has_value(), "Invalid output tensor.");
+            copyBatchBlockOffsets(*(_output.value()), batchSize, batchBlockIndices, numPools, offset);
+        },
+        nb::arg("output"), nb::arg("batch_size"), nb::arg("batch_block_indices"), nb::arg("num_pools"),
+        nb::arg("offset"), nb::call_guard<nb::gil_scoped_release>(), "Copy batch block indices to output tensor");
 }
 
 } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2

cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.h

@@ -17,6 +17,7 @@
 
 #pragma once
 
+#include "tensorrt_llm/nanobind/common/customCasters.h"
 #include <nanobind/nanobind.h>
 
 namespace nb = nanobind;

cpp/tensorrt_llm/nanobind/executor/bindings.cpp

@@ -63,15 +63,15 @@ void initBindings(nb::module_& m)
         new (&self) tle::DecodingMode(nb::cast<tle::DecodingMode::UnderlyingType>(state[0]));
     };
     nb::class_<tle::DecodingMode>(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_static("Auto", &tle::DecodingMode::Auto)
+        .def_static("TopK", &tle::DecodingMode::TopK)
+        .def_static("TopP", &tle::DecodingMode::TopP)
+        .def_static("TopKTopP", &tle::DecodingMode::TopKTopP)
+        .def_static("BeamSearch", &tle::DecodingMode::BeamSearch)
+        .def_static("Medusa", &tle::DecodingMode::Medusa)
+        .def_static("Lookahead", &tle::DecodingMode::Lookahead)
+        .def_static("ExplicitDraftTokens", &tle::DecodingMode::ExplicitDraftTokens)
+        .def_static("Eagle", &tle::DecodingMode::Eagle)
         .def("isAuto", &tle::DecodingMode::isAuto)
         .def("isTopK", &tle::DecodingMode::isTopK)
         .def("isTopP", &tle::DecodingMode::isTopP)

cpp/tensorrt_llm/nanobind/thop/bindings.cpp

@@ -37,38 +37,32 @@ void initBindings(nb::module_& m)
 
     m.def("attention", &torch_ext::attention,
         // Parameters with default values using std::nullopt for optional arguments
-        nb::arg("q"), nb::arg("k") = std::nullopt, nb::arg("v") = std::nullopt, nb::arg("output"),
-        nb::arg("output_sf") = std::nullopt, nb::arg("out_dtype") = std::nullopt, nb::arg("workspace_") = std::nullopt,
-        nb::arg("sequence_length"), nb::arg("host_past_key_value_lengths"), nb::arg("host_total_kv_lens"),
-        nb::arg("context_lengths"), nb::arg("host_context_lengths"), nb::arg("host_request_types"),
-        nb::arg("kv_cache_block_offsets") = std::nullopt, nb::arg("host_kv_cache_block_offsets") = std::nullopt,
-        nb::arg("host_kv_cache_pool_pointers") = std::nullopt, nb::arg("host_kv_cache_pool_mapping") = std::nullopt,
-        nb::arg("cache_indirection") = std::nullopt, nb::arg("kv_scale_orig_quant") = std::nullopt,
-        nb::arg("kv_scale_quant_orig") = std::nullopt, nb::arg("out_scale") = std::nullopt,
-        nb::arg("rotary_inv_freq") = std::nullopt, nb::arg("rotary_cos_sin") = std::nullopt,
-        nb::arg("latent_cache") = std::nullopt, nb::arg("q_pe") = std::nullopt,
-        nb::arg("block_ids_per_seq") = std::nullopt, nb::arg("attention_sinks") = std::nullopt, nb::arg("is_fused_qkv"),
-        nb::arg("update_kv_cache"), nb::arg("predicted_tokens_per_seq"), nb::arg("layer_idx"), nb::arg("num_heads"),
-        nb::arg("num_kv_heads"), nb::arg("head_size"), nb::arg("tokens_per_block") = std::nullopt,
+        nb::arg("q"), nb::arg("k").none(), nb::arg("v").none(), nb::arg("output"), nb::arg("output_sf").none(),
+        nb::arg("out_dtype").none(), nb::arg("workspace_").none(), nb::arg("sequence_length"),
+        nb::arg("host_past_key_value_lengths"), nb::arg("host_total_kv_lens"), nb::arg("context_lengths"),
+        nb::arg("host_context_lengths"), nb::arg("host_request_types"), nb::arg("kv_cache_block_offsets").none(),
+        nb::arg("host_kv_cache_block_offsets").none(), nb::arg("host_kv_cache_pool_pointers").none(),
+        nb::arg("host_kv_cache_pool_mapping").none(), nb::arg("cache_indirection").none(),
+        nb::arg("kv_scale_orig_quant").none(), nb::arg("kv_scale_quant_orig").none(), nb::arg("out_scale").none(),
+        nb::arg("rotary_inv_freq").none(), nb::arg("rotary_cos_sin").none(), nb::arg("latent_cache").none(),
+        nb::arg("q_pe").none(), nb::arg("block_ids_per_seq").none(), nb::arg("attention_sinks").none(),
+        nb::arg("is_fused_qkv"), nb::arg("update_kv_cache"), nb::arg("predicted_tokens_per_seq"), nb::arg("layer_idx"),
+        nb::arg("num_heads"), nb::arg("num_kv_heads"), nb::arg("head_size"), nb::arg("tokens_per_block").none(),
         nb::arg("max_num_requests"), nb::arg("max_context_length"), nb::arg("attention_window_size"),
         nb::arg("sink_token_length"), nb::arg("beam_width"), nb::arg("mask_type"), nb::arg("quant_mode"),
         nb::arg("q_scaling"), nb::arg("position_embedding_type"), nb::arg("rotary_embedding_dim"),
         nb::arg("rotary_embedding_base"), nb::arg("rotary_embedding_scale_type"), nb::arg("rotary_embedding_scales"),
         nb::arg("rotary_embedding_max_position_info"), nb::arg("use_paged_context_fmha"),
-        nb::arg("attention_input_type") = std::nullopt, nb::arg("is_mla_enable"),
-        nb::arg("chunked_prefill_buffer_batch_size") = std::nullopt, nb::arg("q_lora_rank") = std::nullopt,
-        nb::arg("kv_lora_rank") = std::nullopt, nb::arg("qk_nope_head_dim") = std::nullopt,
-        nb::arg("qk_rope_head_dim") = std::nullopt, nb::arg("v_head_dim") = std::nullopt,
-        nb::arg("mrope_rotary_cos_sin") = std::nullopt, nb::arg("mrope_position_deltas") = std::nullopt,
-        nb::arg("mla_tensor_params"), nb::arg("attention_chunk_size") = std::nullopt,
-        nb::arg("softmax_stats_tensor") = std::nullopt, nb::arg("spec_decoding_bool_params"),
-        nb::arg("spec_decoding_tensor_params"), nb::arg("sparse_kv_indices") = std::nullopt,
-        nb::arg("sparse_kv_offsets") = std::nullopt, nb::arg("sparse_attn_indices") = std::nullopt,
-        nb::arg("sparse_attn_offsets") = std::nullopt, nb::arg("sparse_attn_indices_block_size"),
-        nb::arg("sparse_mla_topk") = std::nullopt, nb::arg("cu_q_seqlens") = std::nullopt,
-        nb::arg("cu_kv_seqlens") = std::nullopt, nb::arg("fmha_scheduler_counter") = std::nullopt,
-        nb::arg("mla_bmm1_scale") = std::nullopt, nb::arg("mla_bmm2_scale") = std::nullopt,
-        nb::arg("quant_q_buffer") = std::nullopt, "Multi-head attention operation",
-        nb::call_guard<nb::gil_scoped_release>());
+        nb::arg("attention_input_type").none(), nb::arg("is_mla_enable"),
+        nb::arg("chunked_prefill_buffer_batch_size").none(), nb::arg("q_lora_rank").none(),
+        nb::arg("kv_lora_rank").none(), nb::arg("qk_nope_head_dim").none(), nb::arg("qk_rope_head_dim").none(),
+        nb::arg("v_head_dim").none(), nb::arg("mrope_rotary_cos_sin").none(), nb::arg("mrope_position_deltas").none(),
+        nb::arg("mla_tensor_params"), nb::arg("attention_chunk_size").none(), nb::arg("softmax_stats_tensor").none(),
+        nb::arg("spec_decoding_bool_params"), nb::arg("spec_decoding_tensor_params"),
+        nb::arg("sparse_kv_indices").none(), nb::arg("sparse_kv_offsets").none(), nb::arg("sparse_attn_indices").none(),
+        nb::arg("sparse_attn_offsets").none(), nb::arg("sparse_attn_indices_block_size"),
+        nb::arg("sparse_mla_topk").none(), nb::arg("cu_q_seqlens").none(), nb::arg("cu_kv_seqlens").none(),
+        nb::arg("fmha_scheduler_counter").none(), nb::arg("mla_bmm1_scale").none(), nb::arg("mla_bmm2_scale").none(),
+        nb::arg("quant_q_buffer").none(), "Multi-head attention operation", nb::call_guard<nb::gil_scoped_release>());
 }
 } // namespace tensorrt_llm::nanobind::thop

cpp/tensorrt_llm/pybind/executor/bindings.cpp

@@ -64,15 +64,15 @@ void initBindings(pybind11::module_& m)
         return tle::DecodingMode(state[0].cast<tle::DecodingMode::UnderlyingType>());
     };
     py::class_<tle::DecodingMode>(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_static("Auto", &tle::DecodingMode::Auto)
+        .def_static("TopK", &tle::DecodingMode::TopK)
+        .def_static("TopP", &tle::DecodingMode::TopP)
+        .def_static("TopKTopP", &tle::DecodingMode::TopKTopP)
+        .def_static("BeamSearch", &tle::DecodingMode::BeamSearch)
+        .def_static("Medusa", &tle::DecodingMode::Medusa)
+        .def_static("Lookahead", &tle::DecodingMode::Lookahead)
+        .def_static("ExplicitDraftTokens", &tle::DecodingMode::ExplicitDraftTokens)
+        .def_static("Eagle", &tle::DecodingMode::Eagle)
         .def("isAuto", &tle::DecodingMode::isAuto)
         .def("isTopK", &tle::DecodingMode::isTopK)
         .def("isTopP", &tle::DecodingMode::isTopP)

cpp/tensorrt_llm/pybind/thop/bindings.cpp

@@ -37,38 +37,36 @@ void initBindings(pybind11::module_& m)
 
     m.def("attention", &torch_ext::attention,
         // Parameters with default values using std::nullopt for optional arguments
-        py::arg("q"), py::arg("k") = std::nullopt, py::arg("v") = std::nullopt, py::arg("output"),
-        py::arg("output_sf") = std::nullopt, py::arg("out_dtype") = std::nullopt, py::arg("workspace_") = std::nullopt,
+        py::arg("q"), py::arg("k").none(true), py::arg("v").none(true), py::arg("output"),
+        py::arg("output_sf").none(true), py::arg("out_dtype").none(true), py::arg("workspace_").none(true),
         py::arg("sequence_length"), py::arg("host_past_key_value_lengths"), py::arg("host_total_kv_lens"),
         py::arg("context_lengths"), py::arg("host_context_lengths"), py::arg("host_request_types"),
-        py::arg("kv_cache_block_offsets") = std::nullopt, py::arg("host_kv_cache_block_offsets") = std::nullopt,
-        py::arg("host_kv_cache_pool_pointers") = std::nullopt, py::arg("host_kv_cache_pool_mapping") = std::nullopt,
-        py::arg("cache_indirection") = std::nullopt, py::arg("kv_scale_orig_quant") = std::nullopt,
-        py::arg("kv_scale_quant_orig") = std::nullopt, py::arg("out_scale") = std::nullopt,
-        py::arg("rotary_inv_freq") = std::nullopt, py::arg("rotary_cos_sin") = std::nullopt,
-        py::arg("latent_cache") = std::nullopt, py::arg("q_pe") = std::nullopt,
-        py::arg("block_ids_per_seq") = std::nullopt, py::arg("attention_sinks") = std::nullopt, py::arg("is_fused_qkv"),
-        py::arg("update_kv_cache"), py::arg("predicted_tokens_per_seq"), py::arg("layer_idx"), py::arg("num_heads"),
-        py::arg("num_kv_heads"), py::arg("head_size"), py::arg("tokens_per_block") = std::nullopt,
+        py::arg("kv_cache_block_offsets").none(true), py::arg("host_kv_cache_block_offsets").none(true),
+        py::arg("host_kv_cache_pool_pointers").none(true), py::arg("host_kv_cache_pool_mapping").none(true),
+        py::arg("cache_indirection").none(true), py::arg("kv_scale_orig_quant").none(true),
+        py::arg("kv_scale_quant_orig").none(true), py::arg("out_scale").none(true),
+        py::arg("rotary_inv_freq").none(true), py::arg("rotary_cos_sin").none(true), py::arg("latent_cache").none(true),
+        py::arg("q_pe").none(true), py::arg("block_ids_per_seq").none(true), py::arg("attention_sinks").none(true),
+        py::arg("is_fused_qkv"), py::arg("update_kv_cache"), py::arg("predicted_tokens_per_seq"), py::arg("layer_idx"),
+        py::arg("num_heads"), py::arg("num_kv_heads"), py::arg("head_size"), py::arg("tokens_per_block").none(true),
         py::arg("max_num_requests"), py::arg("max_context_length"), py::arg("attention_window_size"),
         py::arg("sink_token_length"), py::arg("beam_width"), py::arg("mask_type"), py::arg("quant_mode"),
         py::arg("q_scaling"), py::arg("position_embedding_type"), py::arg("rotary_embedding_dim"),
         py::arg("rotary_embedding_base"), py::arg("rotary_embedding_scale_type"), py::arg("rotary_embedding_scales"),
         py::arg("rotary_embedding_max_position_info"), py::arg("use_paged_context_fmha"),
-        py::arg("attention_input_type") = std::nullopt, py::arg("is_mla_enable"),
-        py::arg("chunked_prefill_buffer_batch_size") = std::nullopt, py::arg("q_lora_rank") = std::nullopt,
-        py::arg("kv_lora_rank") = std::nullopt, py::arg("qk_nope_head_dim") = std::nullopt,
-        py::arg("qk_rope_head_dim") = std::nullopt, py::arg("v_head_dim") = std::nullopt,
-        py::arg("mrope_rotary_cos_sin") = std::nullopt, py::arg("mrope_position_deltas") = std::nullopt,
-        py::arg("mla_tensor_params"), py::arg("attention_chunk_size") = std::nullopt,
-        py::arg("softmax_stats_tensor") = std::nullopt, py::arg("spec_decoding_bool_params"),
-        py::arg("spec_decoding_tensor_params"), py::arg("sparse_kv_indices") = std::nullopt,
-        py::arg("sparse_kv_offsets") = std::nullopt, py::arg("sparse_attn_indices") = std::nullopt,
-        py::arg("sparse_attn_offsets") = std::nullopt, py::arg("sparse_attn_indices_block_size"),
-        py::arg("sparse_mla_topk") = std::nullopt, py::arg("cu_q_seqlens") = std::nullopt,
-        py::arg("cu_kv_seqlens") = std::nullopt, py::arg("fmha_scheduler_counter") = std::nullopt,
-        py::arg("mla_bmm1_scale") = std::nullopt, py::arg("mla_bmm2_scale") = std::nullopt,
-        py::arg("quant_q_buffer") = std::nullopt, "Multi-head attention operation",
+        py::arg("attention_input_type").none(true), py::arg("is_mla_enable"),
+        py::arg("chunked_prefill_buffer_batch_size").none(true), py::arg("q_lora_rank").none(true),
+        py::arg("kv_lora_rank").none(true), py::arg("qk_nope_head_dim").none(true),
+        py::arg("qk_rope_head_dim").none(true), py::arg("v_head_dim").none(true),
+        py::arg("mrope_rotary_cos_sin").none(true), py::arg("mrope_position_deltas").none(true),
+        py::arg("mla_tensor_params"), py::arg("attention_chunk_size").none(true),
+        py::arg("softmax_stats_tensor").none(true), py::arg("spec_decoding_bool_params"),
+        py::arg("spec_decoding_tensor_params"), py::arg("sparse_kv_indices").none(true),
+        py::arg("sparse_kv_offsets").none(true), py::arg("sparse_attn_indices").none(true),
+        py::arg("sparse_attn_offsets").none(true), py::arg("sparse_attn_indices_block_size"),
+        py::arg("sparse_mla_topk").none(true), py::arg("cu_q_seqlens").none(true), py::arg("cu_kv_seqlens").none(true),
+        py::arg("fmha_scheduler_counter").none(true), py::arg("mla_bmm1_scale").none(true),
+        py::arg("mla_bmm2_scale").none(true), py::arg("quant_q_buffer").none(true), "Multi-head attention operation",
         py::call_guard<py::gil_scoped_release>());
 }
 } // namespace tensorrt_llm::pybind::thop