bugfix: fix runtime error for mlu.

Author: XuZhang99

xllm/core/framework/batch/batch_input_builder.cpp

@@ -557,6 +557,14 @@ ForwardInput BatchInputBuilder::state_to_forward_input() {
   input_params.decode_seq_range =
       util::find_ones_indices(input_params.q_seq_lens_vec);
 
+  // for flashinfer
+  input_params.paged_kv_indptr =
+      torch::tensor(state_.paged_kv_indptr, torch::kInt);
+  input_params.paged_kv_indices =
+      torch::tensor(state_.paged_kv_indices, torch::kInt);
+  input_params.paged_kv_last_page_len =
+      torch::tensor(state_.paged_kv_last_page_len, torch::kInt);
+
   // Setup multimodal data
   input_params.mm_data = MMData::batch(mm_data_vec_);
 
@@ -631,6 +639,12 @@ RawForwardInput BatchInputBuilder::state_to_raw_forward_input() {
   raw_forward_input.transfer_kv_infos = std::move(state_.transfer_kv_infos);
   raw_forward_input.prefill_seq_len = state_.prefill_seq_len;
 
+  // for flashinfer
+  raw_forward_input.paged_kv_indptr = std::move(state_.paged_kv_indptr);
+  raw_forward_input.paged_kv_indices = std::move(state_.paged_kv_indices);
+  raw_forward_input.paged_kv_last_page_len =
+      std::move(state_.paged_kv_last_page_len);
+
   raw_forward_input.embedding_ids = std::move(state_.embedding_ids);
   raw_forward_input.extra_token_ids = std::move(state_.extra_token_ids);
 

xllm/core/kernels/ops_api.cpp

@@ -139,7 +139,7 @@ void batch_decode(AttentionParams& params) {
   mlu::batch_decode(params.query,
                     params.k_cache,
                     params.output,
-                    params.block_table,
+                    params.block_table.value(),
                     params.kv_seq_lens,
                     params.v_cache,
                     params.output_lse,

xllm/core/kernels/param.h

@@ -72,7 +72,7 @@ struct AttentionParams {
   std::optional<torch::Tensor> alibi_slope;
   std::optional<torch::Tensor> q_quant_scale;
   std::optional<torch::Tensor> out_quant_scale;
-  torch::Tensor block_table;
+  std::optional<torch::Tensor> block_table;
   std::string compute_dtype;
   int max_seq_len;
   int window_size_left;

xllm/core/layers/mlu/attention.cpp

@@ -139,6 +139,7 @@ std::tuple<torch::Tensor, std::optional<torch::Tensor>> AttentionImpl::forward(
 
     // for mlu
     attention_params.block_table = attn_metadata.block_table;
+    attention_params.kv_seq_lens = attn_metadata.kv_seq_lens;
 
     // for flashinfer
     attention_params.paged_kv_indptr = attn_metadata.paged_kv_indptr;