Skip to content

Commit 74a363e

Browse files
Clarity256Clarity256
authored
[XPU][OP] Add build_sampling_params kernel for MTP speculative decoding (#8032)
* [XPU] Add build_sampling_params kernel for MTP speculative decoding Add a new XPU custom operator `build_sampling_params` that constructs sampling parameters (top_p, top_k, topp_seed) on device for MTP speculative decoding verification. This replaces the previous Python-level `padding_sampling_params` approach with a more efficient XPU kernel implementation that supports CudaGraph capture. Key components: - XPU kernel implementation (build_sampling_params.xpu) - C++ wrapper and op registration - Plugin header declaration - Unit tests with comprehensive coverage * [XPU] Fix build_sampling_params seed offset stride to 32 Align the per-position seed offset stride with the Python padding_sampling_params implementation it replaces: XPU requires a stride of 32 (not 4) so that the generated topp_seed sequence matches the original reference. Update both the kernel and CPU wrapper, and the unit test reference accordingly.
1 parent c46930d commit 74a363e

5 files changed

Lines changed: 678 additions & 0 deletions

File tree

‎custom_ops/xpu_ops/src/ops/mtp/build_sampling_params.cc‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include <paddle/phi/backends/xpu/xpu_context.h>
16+
#include "paddle/extension.h"
17+
#include "xpu/plugin.h"
18+
19+
#ifndef PD_BUILD_STATIC_OP
20+
#define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
21+
#endif
22+
23+
namespace api = baidu::xpu::api;
24+
25+
std::vector<paddle::Tensor> BuildSamplingParams(
26+
const paddle::Tensor& top_p,
27+
const paddle::Tensor& top_k,
28+
paddle::Tensor& infer_seed,
29+
const paddle::Tensor& seq_lens_this_time,
30+
const paddle::Tensor& seq_lens_encoder,
31+
const int64_t token_num_output_cpu,
32+
const int64_t increment_value) {
33+
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
34+
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
35+
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
36+
api::Context* ctx = xpu_ctx->x_context();
37+
std::unique_ptr<api::Context> cpu_ctx;
38+
if (top_p.is_cpu()) {
39+
cpu_ctx = std::make_unique<api::Context>(api::kCPU);
40+
ctx = cpu_ctx.get();
41+
}
42+
43+
int real_bsz = static_cast<int>(seq_lens_this_time.shape()[0]);
44+
45+
auto top_p_padding = paddle::empty(
46+
{token_num_output_cpu, 1}, paddle::DataType::FLOAT32, top_p.place());
47+
auto top_k_padding = paddle::empty(
48+
{token_num_output_cpu, 1}, paddle::DataType::INT64, top_p.place());
49+
auto topp_seed = paddle::empty(
50+
{token_num_output_cpu, 1}, paddle::DataType::INT64, top_p.place());
51+
52+
int r =
53+
fastdeploy::plugin::build_sampling_params(ctx,
54+
top_p_padding.data<float>(),
55+
top_k_padding.data<int64_t>(),
56+
topp_seed.data<int64_t>(),
57+
top_p.data<float>(),
58+
top_k.data<int64_t>(),
59+
infer_seed.data<int64_t>(),
60+
seq_lens_this_time.data<int>(),
61+
seq_lens_encoder.data<int>(),
62+
real_bsz,
63+
token_num_output_cpu,
64+
increment_value);
65+
PD_CHECK(r == 0, "fastdeploy::plugin::build_sampling_params failed.");
66+
67+
return {top_p_padding, top_k_padding, topp_seed};
68+
}
69+
70+
std::vector<std::vector<int64_t>> BuildSamplingParamsInferShape(
71+
const std::vector<int64_t>& top_p_shape,
72+
const std::vector<int64_t>& top_k_shape,
73+
const std::vector<int64_t>& infer_seed_shape,
74+
const std::vector<int64_t>& seq_lens_this_time_shape,
75+
const std::vector<int64_t>& seq_lens_encoder_shape) {
76+
// token_num is dynamic; return a placeholder shape of [-1, 1]
77+
return {{-1, 1}, {-1, 1}, {-1, 1}};
78+
}
79+
80+
std::vector<paddle::DataType> BuildSamplingParamsInferDtype(
81+
const paddle::DataType& top_p_dtype,
82+
const paddle::DataType& top_k_dtype,
83+
const paddle::DataType& infer_seed_dtype,
84+
const paddle::DataType& seq_lens_this_time_dtype,
85+
const paddle::DataType& seq_lens_encoder_dtype) {
86+
return {paddle::DataType::FLOAT32,
87+
paddle::DataType::INT64,
88+
paddle::DataType::INT64};
89+
}
90+
91+
PD_BUILD_STATIC_OP(build_sampling_params)
92+
.Inputs({"top_p",
93+
"top_k",
94+
"infer_seed",
95+
"seq_lens_this_time",
96+
"seq_lens_encoder"})
97+
.Outputs({"top_p_padding", "top_k_padding", "topp_seed"})
98+
.Attrs({"token_num_output_cpu: int64_t", "increment_value: int64_t"})
99+
.SetKernelFn(PD_KERNEL(BuildSamplingParams))
100+
.SetInferShapeFn(PD_INFER_SHAPE(BuildSamplingParamsInferShape))
101+
.SetInferDtypeFn(PD_INFER_DTYPE(BuildSamplingParamsInferDtype));

‎custom_ops/xpu_ops/src/plugin/include/xpu/plugin.h‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -834,6 +834,19 @@ DLL_EXPORT int reasoning_phase_token_constraint(
834834
int max_seq_len,
835835
int allowed_tokens_len);
836836

837+
DLL_EXPORT int build_sampling_params(api::Context* ctx,
838+
float* top_p_padding,
839+
int64_t* top_k_padding,
840+
int64_t* topp_seed,
841+
const float* top_p,
842+
const int64_t* top_k,
843+
int64_t* infer_seed,
844+
const int* seq_lens_this_time,
845+
const int* seq_lens_encoder,
846+
int bs,
847+
int64_t token_num,
848+
int64_t increment_value);
849+
837850
/*--------------------------------------- MTP end
838851
* --------------------------------------------*/
839852

‎custom_ops/xpu_ops/src/plugin/src/kernel/kunlun3cpp/mtp_kernel/build_sampling_params.xpu‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include "xpu/kernel/cluster.h"
16+
#include "xpu/kernel/cluster_partition.h"
17+
#include "xpu/kernel/cluster_primitive.h"
18+
19+
namespace fd_xpu3 {
20+
21+
constexpr int64_t BUILD_SAMPLING_MAX_INFER_SEED = 2147483646LL;
22+
23+
// Each cluster handles one batch item (bi = cluster_id).
24+
// Within the cluster, core 0 reads the per-batch scalars and broadcasts via
25+
// shared memory; all cores then fill their assigned token slots in parallel.
26+
__global__ void build_sampling_params_kernel(
27+
__global_ptr__ float* top_p_padding,
28+
__global_ptr__ int64_t* top_k_padding,
29+
__global_ptr__ int64_t* topp_seed,
30+
__global_ptr__ const float* top_p,
31+
__global_ptr__ const int64_t* top_k,
32+
__global_ptr__ int64_t* infer_seed,
33+
__global_ptr__ const int* seq_lens_this_time,
34+
__global_ptr__ const int* seq_lens_encoder,
35+
int bs,
36+
int64_t token_num,
37+
int64_t increment_value) {
38+
int cid = core_id();
39+
int ncores = core_num();
40+
int clusterid = cluster_id();
41+
int nclusters = cluster_num();
42+
43+
// Shared scalars broadcast from core 0 to all cores in the cluster.
44+
__shared__ float sm_top_p;
45+
__shared__ int64_t sm_top_k;
46+
__shared__ int64_t sm_seed;
47+
__shared__ int sm_repeat; // number of tokens this batch produces
48+
__shared__ int sm_pad_start; // starting index in the output buffer
49+
50+
// Shared prefix-sum buffer: each cluster computes its own pad_start via
51+
// a two-pass scan over seq_lens_this_time / seq_lens_encoder.
52+
// We use a simple approach: core 0 of cluster 0 writes per-batch start
53+
// offsets into a global scratch area is not available here, so instead we
54+
// compute pad_start with a sequential scan in core 0 of each cluster.
55+
// Because clusters run concurrently we cannot share a global accumulator;
56+
// instead each cluster independently sums the first `bi` entries.
57+
// This is O(bs) per cluster but bs is typically small (<=512).
58+
59+
for (int bi = clusterid; bi < bs; bi += nclusters) {
60+
if (cid == 0) {
61+
// Read per-batch parameters from global memory.
62+
float lm_top_p;
63+
int64_t lm_top_k;
64+
int64_t lm_seed;
65+
int lm_slt; // seq_lens_this_time[bi]
66+
int lm_sle; // seq_lens_encoder[bi]
67+
68+
GM2LM_ASYNC(top_p + bi, &lm_top_p, sizeof(float));
69+
GM2LM_ASYNC(top_k + bi, &lm_top_k, sizeof(int64_t));
70+
GM2LM_ASYNC(infer_seed + bi, &lm_seed, sizeof(int64_t));
71+
GM2LM_ASYNC(seq_lens_this_time + bi, &lm_slt, sizeof(int));
72+
GM2LM(seq_lens_encoder + bi, &lm_sle, sizeof(int)); // sync barrier
73+
74+
bool is_decoder = (lm_sle == 0);
75+
int repeat = is_decoder ? lm_slt : 1;
76+
77+
// Compute pad_start = sum of token counts for batches [0, bi).
78+
int pad_start = 0;
79+
for (int k = 0; k < bi; k++) {
80+
int slt_k, sle_k;
81+
GM2LM_ASYNC(seq_lens_this_time + k, &slt_k, sizeof(int));
82+
GM2LM(seq_lens_encoder + k, &sle_k, sizeof(int));
83+
pad_start += (sle_k == 0) ? slt_k : 1;
84+
}
85+
86+
sm_top_p = lm_top_p;
87+
sm_top_k = lm_top_k;
88+
sm_seed = lm_seed;
89+
sm_repeat = repeat;
90+
sm_pad_start = pad_start;
91+
}
92+
mfence();
93+
sync_all();
94+
95+
// All cores fill token slots [sm_pad_start, sm_pad_start + sm_repeat).
96+
float bi_top_p = sm_top_p;
97+
int64_t bi_top_k = sm_top_k;
98+
int64_t bi_seed = sm_seed;
99+
int repeat = sm_repeat;
100+
int pad_start = sm_pad_start;
101+
102+
for (int local_pos = cid; local_pos < repeat; local_pos += ncores) {
103+
int pad_idx = pad_start + local_pos;
104+
float lm_top_p_out = bi_top_p;
105+
int64_t lm_top_k_out = bi_top_k;
106+
// Decoder tokens: offset seed by position; encoder token: no offset.
107+
int64_t offset = static_cast<int64_t>(local_pos) * 32;
108+
int64_t lm_seed_out = (bi_seed + offset) % BUILD_SAMPLING_MAX_INFER_SEED;
109+
110+
LM2GM_ASYNC(&lm_top_p_out, top_p_padding + pad_idx, sizeof(float));
111+
LM2GM_ASYNC(&lm_top_k_out, top_k_padding + pad_idx, sizeof(int64_t));
112+
LM2GM(&lm_seed_out, topp_seed + pad_idx, sizeof(int64_t));
113+
}
114+
115+
// Core 0 updates infer_seed in-place.
116+
if (cid == 0) {
117+
int64_t new_seed =
118+
(bi_seed + increment_value) % BUILD_SAMPLING_MAX_INFER_SEED;
119+
LM2GM(&new_seed, infer_seed + bi, sizeof(int64_t));
120+
}
121+
122+
mfence();
123+
sync_all();
124+
}
125+
}
126+
127+
} // namespace fd_xpu3

0 commit comments

Comments
 (0)