feat - support mla cudagraph

Author: Nancheng-11

rtp_llm/cpp/devices/cuda_impl/CudaFlashInfer.cc

@@ -142,7 +142,8 @@ void FlashInferAttnParams::fillParams(torch::Tensor sequence_lengths,
                                       torch::Tensor input_lengths,
                                       torch::Tensor kv_cache_block_id_host,
                                       int           batch_size,
-                                      int           seq_size_per_block) {
+                                      int           seq_size_per_block,
+                                      torch::Tensor prefix_lengths) {
     fillFlashInfer(nullptr,
                    torchTensor2Buffer(sequence_lengths),
                    torchTensor2Buffer(input_lengths),

rtp_llm/cpp/devices/cuda_impl/CudaFlashInfer.h

@@ -109,7 +109,8 @@ struct FlashInferAttnParams: ParamsBase {
                                             torch::Tensor input_lengths,
                                             torch::Tensor kv_cache_block_id_host,
                                             int           batch_size,
-                                            int           seq_size_per_block) override;
+                                            int           seq_size_per_block,
+                                            torch::Tensor prefix_lengths = torch::Tensor()) override;
     void                         fillFlashInfer(const BufferPtr& prefix_lengths_host,
                                                 const BufferPtr& sequence_lengths_host,
                                                 const BufferPtr& input_lengths_host,

rtp_llm/cpp/devices/cuda_impl/CudaGraphRunner.cc

@@ -40,7 +40,7 @@ GraphBase* CudaDevice::getDeviceGraphRunner(const DeviceInitParams& params,
 
 py::object CudaGraphRunner::normalForward(PyModelInputs& inputs) {
     auto attn_pyobj = py_attn_pyobj_method_(inputs, false);
-    attn_pyobj.attr("prepare")(inputs);
+    attn_pyobj.attr("prepare")(inputs.attention_inputs);
     return py_forward_method_(inputs, attn_pyobj);
 }
 
@@ -100,7 +100,7 @@ void CudaGraphRunner::prepareInputs(PyModelInputs& inputs) {
                       py_model_inputs_.attention_inputs.padding_offset,
                       inputs.attention_inputs.padding_offset.size(0) * sizeof(int));
         auto attn_pyobj = graph_instances_[state_.current_real_graph_bs].mem_hold_.attn_pyobj_;
-        attn_pyobj.attr("prepare_replay")(inputs);
+        attn_pyobj.attr("prepare_replay")(inputs.attention_inputs);
     } else {
         auto& py_model_inputs_ = graph_instances_[state_.current_real_graph_seq_len].mem_hold_.py_model_inputs_;
 
@@ -342,7 +342,7 @@ void CudaGraphRunner::initCapture() {
         initKernelInternalMemory();
         // get real output data type
         auto attn_pyobj = py_attn_pyobj_method_(capture_mem_hold_.py_model_inputs_, true);
-        attn_pyobj.attr("prepare")(capture_mem_hold_.py_model_inputs_);
+        attn_pyobj.attr("prepare")(capture_mem_hold_.py_model_inputs_.attention_inputs);
         RTP_LLM_LOG_INFO("initCapture forward for output datatype start");
         auto py_outputs_obj = py_forward_method_(capture_mem_hold_.py_model_inputs_, attn_pyobj);
         RTP_LLM_LOG_INFO("initCapture forward for output datatype end");
@@ -383,7 +383,7 @@ void CudaGraphRunner::captureOneGraphInstance(int key, const char* key_type) {
     // WarmUp twice
     RTP_LLM_LOG_INFO("WarmUp for %s %d start.", key_type, key);
     auto attn_pyobj = graph_instances_[key].mem_hold_.attn_pyobj_;
-    attn_pyobj.attr("prepare")(inputs);
+    attn_pyobj.attr("prepare")(inputs.attention_inputs);
     py_forward_method_(inputs, attn_pyobj);
     py_forward_method_(inputs, attn_pyobj);
     RTP_LLM_LOG_INFO("WarmUp for %s %d successfully.", key_type, key);

rtp_llm/cpp/models/PyWrappedModel.cc

@@ -257,7 +257,7 @@ GptModelOutputs PyWrappedModel::forward(const GptModelInputs& inputs) {
         } else {
             DevicePerfWrapper wrapper(device_, "normal forward");
             auto              attn_pyobj = py_model_.attr("prepare_fmha_impl")(py_model_inputs, false);
-            attn_pyobj.attr("prepare")(py_model_inputs);
+            attn_pyobj.attr("prepare")(py_model_inputs.attention_inputs);
             auto py_model_forward = py_model_.attr("forward");
             auto outputs          = py_model_forward(py_model_inputs, attn_pyobj);
             py_model_outputs      = outputs.cast<PyModelOutputs>();

rtp_llm/models/deepseek_v2.py

@@ -523,6 +523,9 @@ def _create_config(cls, ckpt_path: str):
         DeepSeekV2._from_hf(config, ckpt_path)
         return config
 
+    def support_cuda_graph(self) -> bool:
+        return True
+
     def _create_python_model(self) -> Optional[GptModelBase]:
         self.py_model = GenericMoeModel(self.config, self.weight)
 

rtp_llm/models_py/bindings/OpDefs.cc

@@ -46,19 +46,6 @@ void registerPyOpDefs(pybind11::module& m) {
             pybind11::arg("seq_size_per_block"),
             "Fill parameters for CUDA graph execution");
 
-    pybind11::class_<MlaParams, std::shared_ptr<MlaParams>, rtp_llm::ParamsBase>(m, "MlaParams")
-        .def(pybind11::init<>())
-        .def_readonly("batch_indice", &MlaParams::batch_indice)
-        .def_readonly("positions", &MlaParams::positions)
-        .def_readonly("paged_kv_last_page_len", &MlaParams::paged_kv_last_page_len)
-        .def_readonly("kvlen", &MlaParams::kvlen)
-        .def_readonly("page_indice", &MlaParams::page_indice)
-        .def_readonly("reuse_cache_page_indice", &MlaParams::reuse_cache_page_indice)
-        .def_readonly("decode_page_indptr", &MlaParams::decode_page_indptr)
-        .def_readonly("prefill_page_indptr", &MlaParams::prefill_page_indptr)
-        .def_readonly("qo_indptr", &MlaParams::qo_indptr)
-        .def_readonly("batch_reuse_info_vec", &MlaParams::batch_reuse_info_vec);
-
     pybind11::class_<PyPrefillCudaGaphCopyParams>(m, "PyPrefillCudaGaphCopyParams")
         .def(pybind11::init<>())
         .def_readonly("cuda_graph_prefill_batch_size", &PyPrefillCudaGaphCopyParams::cuda_graph_prefill_batch_size)

rtp_llm/models_py/bindings/OpDefs.h

@@ -8,22 +8,6 @@
 #include "rtp_llm/models_py/bindings/ParamsBase.h"
 #include "rtp_llm/cpp/utils/Logger.h"
 namespace torch_ext {
-struct MlaParams: public rtp_llm::ParamsBase {
-    torch::Tensor batch_indice;
-    torch::Tensor positions;
-    torch::Tensor paged_kv_last_page_len;
-    torch::Tensor kvlen;
-    torch::Tensor page_indice;
-    torch::Tensor reuse_cache_page_indice;
-    torch::Tensor decode_page_indptr;
-    torch::Tensor prefill_page_indptr;
-    torch::Tensor qo_indptr;
-    torch::Tensor batch_reuse_info_vec;
-
-    // Hidden field to keep FlashInferMlaAttnParams object alive
-    // This ensures the underlying buffers (buf_d, buf_h) are not deallocated
-    std::shared_ptr<void> _params_holder;
-};
 
 struct KVCache {
     torch::Tensor k_cache_base;
@@ -96,7 +80,7 @@ struct PyAttentionInputs {
     std::optional<PyCacheStoreInputs> cache_store_inputs;
 
     std::optional<PyPrefillCudaGaphCopyParams> prefill_cuda_graph_copy_params;
-    bool                              is_s_padded = false;
+    bool                                       is_s_padded = false;
 };
 
 struct BertEmbeddingInputs {

rtp_llm/models_py/bindings/ParamsBase.h

@@ -12,7 +12,8 @@ class ParamsBase {
                             torch::Tensor input_lengths,
                             torch::Tensor kv_cache_block_id_host,
                             int           batch_size,
-                            int           seq_size_per_block) {};
+                            int           seq_size_per_block,
+                            torch::Tensor prefix_lengths = torch::Tensor()) {};
     // check whether the parmas can be recycled automatically.
     virtual bool check_recycle() {
         return false;

rtp_llm/models_py/bindings/cuda/FlashInferMlaParams.cc

@@ -11,6 +11,9 @@ using namespace torch_ext;
 
 namespace rtp_llm {
 
+static const int MIN_CACHE_PAGE_NUM = 1024 * 1024;
+// static const int MIN_CACHE_BATCH_SIZE      = 256;
+// static const int MIN_CACHE_INPUT_TOKEN_NUM = 512;
 std::tuple<torch::Tensor, std::vector<torch::Tensor>>
 FlashInferMlaAttnParams::allocateManyBuffer(const std::vector<std::vector<int64_t>>& shapes, bool is_device) {
     std::vector<torch::Tensor> tensors;
@@ -65,7 +68,7 @@ void FlashInferMlaAttnParams::ensureTensorSize(
     // Update max sizes
     max_batch_size_       = std::max(max_batch_size_, batch_size);
     max_input_token_num_  = std::max(max_input_token_num_, input_token_num);
-    max_page_num_         = std::max(max_page_num_, page_num);
+    max_page_num_         = std::max(max_page_num_, MIN_CACHE_PAGE_NUM);
     max_reuse_page_num_   = std::max(max_reuse_page_num_, reuse_page_num);
     max_batch_reuse_info_ = std::max(max_batch_reuse_info_, batch_reuse_info_size);
 
@@ -317,11 +320,12 @@ void FlashInferMlaAttnParams::refreshBuffer(
     batch_reuse_info_vec_h.unsafeGetTensorImpl()->set_sizes_contiguous(shape);
 }
 
-MlaParams FlashInferMlaAttnParams::fillParams(torch::Tensor t_prefix_lengths,
-                                              torch::Tensor t_sequence_lengths,
-                                              torch::Tensor t_input_lengths,
-                                              torch::Tensor t_kv_cache_block_id_host,
-                                              int           seq_size_per_block) {
+void FlashInferMlaAttnParams::fillParams(torch::Tensor t_sequence_lengths,
+                                         torch::Tensor t_input_lengths,
+                                         torch::Tensor t_kv_cache_block_id_host,
+                                         int           t_batch_size,
+                                         int           seq_size_per_block,
+                                         torch::Tensor t_prefix_lengths) {
     const int batch_size = t_input_lengths.size(0);
 
     // First pass: calculate required sizes accurately
@@ -370,54 +374,77 @@ MlaParams FlashInferMlaAttnParams::fillParams(torch::Tensor t_prefix_lengths,
     // Refresh buffer (copy to DEVICE and update shapes)
     refreshBuffer(batch_size, input_token_num, page_num, reuse_page_num, batch_reuse_info_size);
 
-    batch_indice            = batch_indice_d;
-    page_indice             = page_indice_d;
-    reuse_cache_page_indice = reuse_page_num > 0 ? reuse_cache_page_indice_d : torch::Tensor();
-    decode_page_indptr      = decode_page_indptr_d;
-    prefill_page_indptr     = prefill_page_indptr_d;
-    paged_kv_last_page_len  = paged_kv_last_page_len_d;
-    qo_indptr               = qo_indptr_d;
-    kvlen                   = kvlen_d;
-    positions               = positions_d;
-    batch_reuse_info_vec    = batch_size > 0 ? batch_reuse_info_vec_d : torch::Tensor();
-
-    // Return MlaParams with DEVICE tensors
-    MlaParams params;
-    params.batch_indice            = batch_indice_d;
-    params.page_indice             = page_indice_d;
-    params.reuse_cache_page_indice = reuse_page_num > 0 ? reuse_cache_page_indice_d : torch::Tensor();
-    params.decode_page_indptr      = decode_page_indptr_d;
-    params.prefill_page_indptr     = prefill_page_indptr_d;
-    params.paged_kv_last_page_len  = paged_kv_last_page_len_d;
-    params.qo_indptr               = qo_indptr_d;
-    params.kvlen                   = kvlen_d;
-    params.positions               = positions_d;
-    params.batch_reuse_info_vec    = batch_size > 0 ? batch_reuse_info_vec_d : torch::Tensor();
-
-    return params;
+    return;
 }
 
 void registerPyFlashInferMlaParams(pybind11::module& m) {
+    pybind11::class_<FlashInferMlaAttnParams, std::shared_ptr<FlashInferMlaAttnParams>, rtp_llm::ParamsBase>(
+        m, "FlashInferMlaAttnParams")
+        .def(pybind11::init<>())
+        // HOST tensors (_h suffix)
+        .def_readonly("batch_indice_h", &FlashInferMlaAttnParams::batch_indice_h, "Batch indices on HOST")
+        .def_readonly("page_indice_h", &FlashInferMlaAttnParams::page_indice_h, "Page indices on HOST")
+        .def_readonly("reuse_cache_page_indice_h",
+                      &FlashInferMlaAttnParams::reuse_cache_page_indice_h,
+                      "Reuse cache page indices on HOST")
+        .def_readonly(
+            "decode_page_indptr_h", &FlashInferMlaAttnParams::decode_page_indptr_h, "Decode page indptr on HOST")
+        .def_readonly(
+            "prefill_page_indptr_h", &FlashInferMlaAttnParams::prefill_page_indptr_h, "Prefill page indptr on HOST")
+        .def_readonly("paged_kv_last_page_len_h",
+                      &FlashInferMlaAttnParams::paged_kv_last_page_len_h,
+                      "Paged KV last page length on HOST")
+        .def_readonly("qo_indptr_h", &FlashInferMlaAttnParams::qo_indptr_h, "Query/output indptr on HOST")
+        .def_readonly("kvlen_h", &FlashInferMlaAttnParams::kvlen_h, "KV length on HOST")
+        .def_readonly("positions_h", &FlashInferMlaAttnParams::positions_h, "Positions on HOST")
+        .def_readonly("batch_reuse_info_vec_h",
+                      &FlashInferMlaAttnParams::batch_reuse_info_vec_h,
+                      "Batch reuse info vector on HOST")
+        // DEVICE tensors (_d suffix)
+        .def_readonly("batch_indice_d", &FlashInferMlaAttnParams::batch_indice_d, "Batch indices on DEVICE")
+        .def_readonly("page_indice_d", &FlashInferMlaAttnParams::page_indice_d, "Page indices on DEVICE")
+        .def_readonly("reuse_cache_page_indice_d",
+                      &FlashInferMlaAttnParams::reuse_cache_page_indice_d,
+                      "Reuse cache page indices on DEVICE")
+        .def_readonly(
+            "decode_page_indptr_d", &FlashInferMlaAttnParams::decode_page_indptr_d, "Decode page indptr on DEVICE")
+        .def_readonly(
+            "prefill_page_indptr_d", &FlashInferMlaAttnParams::prefill_page_indptr_d, "Prefill page indptr on DEVICE")
+        .def_readonly("paged_kv_last_page_len_d",
+                      &FlashInferMlaAttnParams::paged_kv_last_page_len_d,
+                      "Paged KV last page length on DEVICE")
+        .def_readonly("qo_indptr_d", &FlashInferMlaAttnParams::qo_indptr_d, "Query/output indptr on DEVICE")
+        .def_readonly("kvlen_d", &FlashInferMlaAttnParams::kvlen_d, "KV length on DEVICE")
+        .def_readonly("positions_d", &FlashInferMlaAttnParams::positions_d, "Positions on DEVICE")
+        .def_readonly("batch_reuse_info_vec_d",
+                      &FlashInferMlaAttnParams::batch_reuse_info_vec_d,
+                      "Batch reuse info vector on DEVICE");
+
     m.def(
         "fill_mla_params",
-        [](torch::Tensor t_prefill_lengths,
-           torch::Tensor t_sequence_lengths,
+        [](torch::Tensor t_sequence_lengths,
            torch::Tensor t_input_lengths,
            torch::Tensor t_kv_cache_block_id_host,
-           int           seq_size_per_block) {
-            auto params     = std::make_shared<rtp_llm::FlashInferMlaAttnParams>();
-            auto mla_params = params->fillParams(
-                t_prefill_lengths, t_sequence_lengths, t_input_lengths, t_kv_cache_block_id_host, seq_size_per_block);
+           int           batch_size,
+           int           seq_size_per_block,
+           torch::Tensor t_prefix_lengths) {
+            auto params = std::make_shared<rtp_llm::FlashInferMlaAttnParams>();
+            params->fillParams(t_sequence_lengths,
+                               t_input_lengths,
+                               t_kv_cache_block_id_host,
+                               batch_size,
+                               seq_size_per_block,
+                               t_prefix_lengths);
             // Store the params object in _params_holder to keep it alive
             // This ensures the underlying buffers (buf_d, buf_h) are not deallocated
-            mla_params._params_holder = std::static_pointer_cast<void>(params);
-            return mla_params;
+            return params;
         },
-        pybind11::arg("t_prefill_lengths"),
         pybind11::arg("t_sequence_lengths"),
         pybind11::arg("t_input_lengths"),
         pybind11::arg("t_kv_cache_block_id_host"),
-        pybind11::arg("seq_size_per_block"));
+        pybind11::arg("batch_size"),
+        pybind11::arg("seq_size_per_block"),
+        pybind11::arg("t_prefix_lengths"));
 }
 
 }  // namespace rtp_llm

rtp_llm/models_py/bindings/cuda/FlashInferMlaParams.h

@@ -17,40 +17,6 @@ class FlashInferMlaAttnParams: public ParamsBase {
     torch::Tensor buf_h;  // Large continuous HOST buffer (pinned memory)
     torch::Tensor buf_d;  // Large continuous DEVICE buffer
 
-    // Tensor views into buf_h and buf_d
-    torch::Tensor batch_indice_h;
-    torch::Tensor page_indice_h;
-    torch::Tensor reuse_cache_page_indice_h;
-    torch::Tensor decode_page_indptr_h;
-    torch::Tensor prefill_page_indptr_h;
-    torch::Tensor paged_kv_last_page_len_h;
-    torch::Tensor qo_indptr_h;
-    torch::Tensor kvlen_h;
-    torch::Tensor positions_h;
-    torch::Tensor batch_reuse_info_vec_h;
-
-    torch::Tensor batch_indice_d;
-    torch::Tensor page_indice_d;
-    torch::Tensor reuse_cache_page_indice_d;
-    torch::Tensor decode_page_indptr_d;
-    torch::Tensor prefill_page_indptr_d;
-    torch::Tensor paged_kv_last_page_len_d;
-    torch::Tensor qo_indptr_d;
-    torch::Tensor kvlen_d;
-    torch::Tensor positions_d;
-    torch::Tensor batch_reuse_info_vec_d;
-
-    torch::Tensor batch_indice;
-    torch::Tensor positions;
-    torch::Tensor paged_kv_last_page_len;
-    torch::Tensor kvlen;
-    torch::Tensor page_indice;
-    torch::Tensor reuse_cache_page_indice;
-    torch::Tensor decode_page_indptr;
-    torch::Tensor prefill_page_indptr;
-    torch::Tensor qo_indptr;
-    torch::Tensor batch_reuse_info_vec;
-
     // Reserved sizes
     int max_batch_size_       = 0;
     int max_input_token_num_  = 0;
@@ -83,11 +49,35 @@ class FlashInferMlaAttnParams: public ParamsBase {
     ensureTensorSize(int batch_size, int input_token_num, int page_num, int reuse_page_num, int batch_reuse_info_size);
 
 public:
-    MlaParams fillParams(torch::Tensor t_prefix_lengths,
-                         torch::Tensor t_sequence_lengths,
-                         torch::Tensor t_input_lengths,
-                         torch::Tensor t_kv_cache_block_id_host,
-                         int           seq_size_per_block);
+    // Tensor views into buf_h and buf_d
+    torch::Tensor batch_indice_h;
+    torch::Tensor page_indice_h;
+    torch::Tensor reuse_cache_page_indice_h;
+    torch::Tensor decode_page_indptr_h;
+    torch::Tensor prefill_page_indptr_h;
+    torch::Tensor paged_kv_last_page_len_h;
+    torch::Tensor qo_indptr_h;
+    torch::Tensor kvlen_h;
+    torch::Tensor positions_h;
+    torch::Tensor batch_reuse_info_vec_h;
+
+    torch::Tensor batch_indice_d;
+    torch::Tensor page_indice_d;
+    torch::Tensor reuse_cache_page_indice_d;
+    torch::Tensor decode_page_indptr_d;
+    torch::Tensor prefill_page_indptr_d;
+    torch::Tensor paged_kv_last_page_len_d;
+    torch::Tensor qo_indptr_d;
+    torch::Tensor kvlen_d;
+    torch::Tensor positions_d;
+    torch::Tensor batch_reuse_info_vec_d;
+
+    void fillParams(torch::Tensor sequence_lengths,
+                    torch::Tensor input_lengths,
+                    torch::Tensor kv_cache_block_id_host,
+                    int           batch_size,
+                    int           seq_size_per_block,
+                    torch::Tensor prefix_lengths = torch::Tensor()) override;
 };
 void registerPyFlashInferMlaParams(pybind11::module& m);