Skip to content

[XPU] get_infer_param use inplace copy, remove block_tables abundant d2h copy #7431

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
RuohengMa wants to merge 23 commits into
base: develop
Choose a base branch
from
Open

[XPU] get_infer_param use inplace copy, remove block_tables abundant d2h copy #7431

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
786c9b6
inplace_copy: encoder_batch_idx/decoder_batch_idx bs == 9 ok
RuohengMa Apr 14, 2026
e984971
inplace_copy: encoder_seq_lod/decoder_seq_lod bs == 9 ok
RuohengMa Apr 14, 2026
ad80381
inplace_copy: all bs == 9 ok
RuohengMa Apr 14, 2026
27bd39c
inplace_copy: all cpu bs == 9 ok
RuohengMa Apr 14, 2026
385a48f
inplace_copy: len_info_cpu bs == 9 ok
RuohengMa Apr 14, 2026
461d554
finished and rm unused code
RuohengMa Apr 14, 2026
8ca7538
prefix_block_tables reuse
RuohengMa Apr 15, 2026
b972890
refine
RuohengMa Apr 15, 2026
888e91b
improve performance
RuohengMa Apr 15, 2026
a38f2ba
remove block_table copy to cpu
RuohengMa Apr 15, 2026
3aed097
fix unit test
RuohengMa Apr 16, 2026
7bc1ec7
fix
RuohengMa Apr 16, 2026
0498456
resolve conflict
RuohengMa Apr 17, 2026
7d2b8e9
refine code
RuohengMa Apr 16, 2026
ddeb1b2
fix
RuohengMa Apr 16, 2026
6bc55a6
fix
RuohengMa Apr 17, 2026
3283716
fix
RuohengMa Apr 17, 2026
cbec29b
fix
RuohengMa Apr 17, 2026
fb0c429
fix
RuohengMa Apr 17, 2026
d2430e5
try fix unit tests
RuohengMa Apr 17, 2026
0d537d4
fix
RuohengMa Apr 17, 2026
d3b903a
tmp save
RuohengMa Apr 17, 2026
7d3c075
fix unit test
RuohengMa Apr 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
322 changes: 178 additions & 144 deletions custom_ops/xpu_ops/src/ops/get_infer_param.cc
View file
Open in desktop

Large diffs are not rendered by default.

44 changes: 44 additions & 0 deletions custom_ops/xpu_ops/src/ops/pybind/pybind.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -478,6 +478,28 @@ std::vector<paddle::Tensor> GetInferParam(
const paddle::Tensor& seq_lens_decoder,
const paddle::Tensor& seq_lens_this_time,
const paddle::Tensor& block_tables,
paddle::Tensor& encoder_batch_map,
paddle::Tensor& decoder_batch_map,
paddle::Tensor& encoder_batch_idx,
paddle::Tensor& decoder_batch_idx,
paddle::Tensor& encoder_seq_lod,
paddle::Tensor& decoder_seq_lod,
paddle::Tensor& encoder_kv_lod,
paddle::Tensor& prefix_len,
paddle::Tensor& decoder_context_len,
paddle::Tensor& decoder_context_len_cache,
paddle::Tensor& prefix_block_tables,
paddle::Tensor& encoder_batch_map_cpu,
paddle::Tensor& decoder_batch_map_cpu,
paddle::Tensor& encoder_batch_idx_cpu,
paddle::Tensor& decoder_batch_idx_cpu,
paddle::Tensor& encoder_seq_lod_cpu,
paddle::Tensor& decoder_seq_lod_cpu,
paddle::Tensor& encoder_kv_lod_cpu,
paddle::Tensor& prefix_len_cpu,
paddle::Tensor& decoder_context_len_cpu,
paddle::Tensor& decoder_context_len_cache_cpu,
paddle::Tensor& len_info_cpu,
int block_size,
int num_speculative_tokens);

Expand Down Expand Up @@ -1052,6 +1074,28 @@ PYBIND11_MODULE(fastdeploy_ops, m) {
py::arg("seq_lens_decoder"),
py::arg("seq_lens_this_time"),
py::arg("block_tables"),
py::arg("encoder_batch_map"),
py::arg("decoder_batch_map"),
py::arg("encoder_batch_idx"),
py::arg("decoder_batch_idx"),
py::arg("encoder_seq_lod"),
py::arg("decoder_seq_lod"),
py::arg("encoder_kv_lod"),
py::arg("prefix_len"),
py::arg("decoder_context_len"),
py::arg("decoder_context_len_cache"),
py::arg("prefix_block_tables"),
py::arg("encoder_batch_map_cpu"),
py::arg("decoder_batch_map_cpu"),
py::arg("encoder_batch_idx_cpu"),
py::arg("decoder_batch_idx_cpu"),
py::arg("encoder_seq_lod_cpu"),
py::arg("decoder_seq_lod_cpu"),
py::arg("encoder_kv_lod_cpu"),
py::arg("prefix_len_cpu"),
py::arg("decoder_context_len_cpu"),
py::arg("decoder_context_len_cache_cpu"),
py::arg("len_info_cpu"),
py::arg("block_size"),
py::arg("num_speculative_tokens"),
"Get infer parameters for block attention in XPU");
Expand Down
73 changes: 61 additions & 12 deletions custom_ops/xpu_ops/test/test_adjust_batch_and_gather_next_token.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

import numpy as np
import paddle
from utils import init_inplace_tensor

from fastdeploy.model_executor.ops.xpu import (
adjust_batch,
Expand All @@ -33,35 +34,87 @@ def _run_test_base(seq_lens_this_time_data, is_speculative):
seq_lens_this_time = paddle.to_tensor(seq_lens_this_time_data, dtype="int32")

bsz = seq_lens_this_time.shape[0]
cum_offsets = paddle.zeros(bsz, dtype="int32")
block_table = paddle.arange(0, 56, dtype="int32").reshape((bsz, 8))

infer_params = get_infer_param(seq_lens_encoder, seq_lens_decoder, seq_lens_this_time, block_table, 64)

(
encoder_batch_map,
decoder_batch_map,
encoder_batch_idx,
decoder_batch_idx,
encoder_seq_lod,
decoder_seq_lod,
encoder_kv_lod,
prefix_len,
decoder_context_len,
decoder_context_len_cache,
prefix_block_tables,
encoder_batch_map_cpu,
decoder_batch_map_cpu,
encoder_batch_idx_cpu,
decoder_batch_idx_cpu,
encoder_seq_lod_cpu,
decoder_seq_lod_cpu,
encoder_kv_lod_cpu,
prefix_len_cpu,
decoder_context_len_cpu,
decoder_context_len_cache_cpu,
len_info_cpu,
) = init_inplace_tensor(seq_lens_encoder.shape[0], block_table.shape)
(
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
slot_mapping_enc,
slot_mapping_dec,
) = get_infer_param(
seq_lens_encoder,
seq_lens_decoder,
seq_lens_this_time,
block_table,
encoder_batch_map,
decoder_batch_map,
encoder_batch_idx,
decoder_batch_idx,
encoder_seq_lod,
decoder_seq_lod,
encoder_kv_lod,
prefix_len,
decoder_context_len,
decoder_context_len_cache,
prefix_block_tables,
encoder_batch_map_cpu,
decoder_batch_map_cpu,
encoder_batch_idx_cpu,
decoder_batch_idx_cpu,
encoder_seq_lod_cpu,
decoder_seq_lod_cpu,
_,
_,
_,
_,
encoder_kv_lod_cpu,
prefix_len_cpu,
decoder_context_len_cpu,
decoder_context_len_cache_cpu,
len_info_cpu,
) = infer_params
64,
0,
)

token_num = seq_lens_this_time.sum().cpu().item()
hidden_dim = 8192
Expand All @@ -72,7 +125,6 @@ def _run_test_base(seq_lens_this_time_data, is_speculative):
# 测试 adjust_batch
adjusted_output = adjust_batch(
input_tensor,
cum_offsets,
encoder_seq_lod,
decoder_seq_lod,
encoder_batch_idx,
Expand All @@ -88,7 +140,6 @@ def _run_test_base(seq_lens_this_time_data, is_speculative):

adjusted_output_cpu = adjust_batch(
input_tensor.cpu(),
cum_offsets,
encoder_seq_lod,
decoder_seq_lod,
encoder_batch_idx,
Expand All @@ -110,7 +161,6 @@ def _run_test_base(seq_lens_this_time_data, is_speculative):
# 测试 gather_next_token
gather_out = gather_next_token(
adjusted_output,
cum_offsets,
encoder_seq_lod,
decoder_seq_lod,
encoder_batch_map,
Expand All @@ -126,7 +176,6 @@ def _run_test_base(seq_lens_this_time_data, is_speculative):

gather_out_cpu = gather_next_token(
adjusted_output.cpu(),
cum_offsets,
encoder_seq_lod,
decoder_seq_lod,
encoder_batch_map,
Expand Down
70 changes: 61 additions & 9 deletions custom_ops/xpu_ops/test/test_adjust_batch_and_recover_batch_sequence.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

import numpy as np
import paddle
from utils import init_inplace_tensor

from fastdeploy.model_executor.ops.xpu import (
adjust_batch,
Expand All @@ -33,32 +34,85 @@ def _run_test_base(seq_lens_this_time_data):
cum_offsets = paddle.zeros(bsz, dtype="int32")
block_table = paddle.arange(0, 56, dtype="int32").reshape((bsz, 8))

infer_params = get_infer_param(seq_lens_encoder, seq_lens_decoder, seq_lens_this_time, block_table, 64)

(
encoder_batch_map,
decoder_batch_map,
encoder_batch_idx,
decoder_batch_idx,
encoder_seq_lod,
decoder_seq_lod,
encoder_kv_lod,
prefix_len,
decoder_context_len,
decoder_context_len_cache,
prefix_block_tables,
encoder_batch_map_cpu,
decoder_batch_map_cpu,
encoder_batch_idx_cpu,
decoder_batch_idx_cpu,
encoder_seq_lod_cpu,
decoder_seq_lod_cpu,
encoder_kv_lod_cpu,
prefix_len_cpu,
decoder_context_len_cpu,
decoder_context_len_cache_cpu,
len_info_cpu,
) = init_inplace_tensor(seq_lens_encoder.shape[0], block_table.shape)
(
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
_,
slot_mapping_enc,
slot_mapping_dec,
) = get_infer_param(
seq_lens_encoder,
seq_lens_decoder,
seq_lens_this_time,
block_table,
encoder_batch_map,
decoder_batch_map,
encoder_batch_idx,
decoder_batch_idx,
encoder_seq_lod,
decoder_seq_lod,
encoder_kv_lod,
prefix_len,
decoder_context_len,
decoder_context_len_cache,
prefix_block_tables,
encoder_batch_map_cpu,
decoder_batch_map_cpu,
encoder_batch_idx_cpu,
decoder_batch_idx_cpu,
encoder_seq_lod_cpu,
decoder_seq_lod_cpu,
_,
_,
_,
_,
encoder_kv_lod_cpu,
prefix_len_cpu,
decoder_context_len_cpu,
decoder_context_len_cache_cpu,
len_info_cpu,
) = infer_params
64,
0,
)

token_num = seq_lens_this_time.sum().cpu().item()
hidden_dim = 8192
Expand All @@ -68,7 +122,6 @@ def _run_test_base(seq_lens_this_time_data):
# test adjust_batch
adjusted_output = adjust_batch(
input_tensor,
cum_offsets,
encoder_seq_lod,
decoder_seq_lod,
encoder_batch_idx,
Expand All @@ -84,7 +137,6 @@ def _run_test_base(seq_lens_this_time_data):

adjusted_output_cpu = adjust_batch(
input_tensor.cpu(),
cum_offsets,
encoder_seq_lod,
decoder_seq_lod,
encoder_batch_idx,
Expand Down
Loading
Loading