| 44 | |
| 45 | template <typename T, typename Context> |
| 46 | void TopPSamplingKernel(const Context& dev_ctx, |
| 47 | const DenseTensor& x, |
| 48 | const DenseTensor& ps, |
| 49 | const optional<DenseTensor>& threshold, |
| 50 | const optional<DenseTensor>& topp_seed, |
| 51 | int64_t random_seed, |
| 52 | int k, |
| 53 | const std::string& mode, |
| 54 | DenseTensor* out, |
| 55 | DenseTensor* ids, |
| 56 | DenseTensor* topk_scores, |
| 57 | DenseTensor* topk_ids) { |
| 58 | using XPUType = typename XPUTypeTrait<T>::Type; |
| 59 | |
| 60 | const XPUType* x_ptr = reinterpret_cast<const XPUType*>(x.data<T>()); |
| 61 | const XPUType* ps_ptr = reinterpret_cast<const XPUType*>(ps.data<T>()); |
| 62 | XPUType* out_ptr = reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(out)); |
| 63 | int64_t* ids_ptr = dev_ctx.template Alloc<int64_t>(ids); |
| 64 | auto x_dims = x.dims(); |
| 65 | int64_t bs = x_dims[0]; |
| 66 | int64_t vocab_size = x_dims[1]; |
| 67 | |
| 68 | XPUType* topk_scores_data = nullptr; |
| 69 | int64_t* topk_ids_data = nullptr; |
| 70 | if (k > 0) { |
| 71 | topk_scores_data = |
| 72 | reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(topk_scores)); |
| 73 | topk_ids_data = dev_ctx.template Alloc<int64_t>(topk_ids); |
| 74 | int r = xpu::topk<XPUType, int64_t>(dev_ctx.x_context(), |
| 75 | x_ptr, |
| 76 | topk_scores_data, |
| 77 | topk_ids_data, |
| 78 | {bs, vocab_size}, |
| 79 | k, |
| 80 | 1, |
| 81 | true, |
| 82 | true); |
| 83 | PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu::topk"); |
| 84 | } |
| 85 | std::vector<int64_t> infer_seed(bs, random_seed); |
| 86 | if (topp_seed.get_ptr() != nullptr) { |
| 87 | TensorToVector(*topp_seed, dev_ctx, &infer_seed); |
| 88 | } |
| 89 | |
| 90 | std::uniform_real_distribution<float> dist(0.0, 1.0); |
| 91 | std::vector<float> rand_coeff_cpu; |
| 92 | for (int64_t i = 0; i < bs; i++) { |
| 93 | if (infer_seed[i] == -1) { |
| 94 | std::shared_ptr<std::mt19937_64> engine = |
| 95 | dev_ctx.GetGenerator()->GetCPUEngine(); |
| 96 | rand_coeff_cpu.push_back(dist(*engine)); |
| 97 | } else { |
| 98 | std::mt19937_64 engine(infer_seed[i]); |
| 99 | rand_coeff_cpu.push_back(dist(engine)); |
| 100 | } |
| 101 | } |
| 102 | uint64_t seed_now = rand_coeff_cpu.empty() ? random_seed : rand_coeff_cpu[0]; |
| 103 | uint64_t offset = 0; |
nothing calls this directly
no test coverage detected