fix: seed int64_t

Author: CrimsonDump

rtp_llm/cpp/core/torch_utils/BufferTorchUtils.h

@@ -63,7 +63,6 @@ inline std::vector<int64_t> bufferShapeToTorchShape(const Buffer& buffer) {
     F(TYPE_INT16, torch::kShort)                                                                                       \
     F(TYPE_INT32, torch::kInt)                                                                                         \
     F(TYPE_INT64, torch::kLong)                                                                                        \
-    F(TYPE_UINT64, torch::kUInt64)                                                                                        \
     F(TYPE_FP16, torch::kHalf)                                                                                         \
     F(TYPE_FP32, torch::kFloat)                                                                                        \
     F(TYPE_FP64, torch::kDouble)                                                                                       \

rtp_llm/cpp/devices/rocm_impl/ROCmSampleOp.cc

@@ -38,7 +38,6 @@ GreedyOutput ROCmDevice::sampleGreedy(const GreedyParams& params) {
     auto& top_k       = params.top_k;
     auto& top_p       = params.top_p;
     auto& temperature = params.temperature;
-    // auto& random_seed = params.random_seed;
     ROCM_CHECK_VALUE(top_k.size() == batch_size, "top_k.size() != batch_size");
     ROCM_CHECK_VALUE(top_p.size() == batch_size, "top_p.size() != batch_size");
     ROCM_CHECK_VALUE(temperature.size() == batch_size, "temperature.size() != batch_size");
@@ -116,12 +115,14 @@ GreedyOutput ROCmDevice::sampleGreedy(const GreedyParams& params) {
     }
 
     bool deterministic = true;
-    auto seed_h   = allocateBuffer({DataType::TYPE_UINT64, {batch_size}, AllocationType::HOST});
-    auto offset_h = allocateBuffer({DataType::TYPE_UINT64, {batch_size}, AllocationType::HOST});
+    auto seed_h   = allocateBuffer({DataType::TYPE_INT64, {batch_size}, AllocationType::HOST});
+    auto offset_h = allocateBuffer({DataType::TYPE_INT64, {batch_size}, AllocationType::HOST});
     for (int i = 0; i < batch_size; i++) {
-        std::tie(seed_h->data<uint64_t>()[i], offset_h->data<uint64_t>()[i]) = params.generator[i].defined() ?
+        auto [sd, ofst] = params.generator[i].defined() ?
             get_seed_and_offset(batch_size * 32, params.generator[i]) :
             std::make_pair(0ULL, 0ULL);
+        seed_h->data<int64_t>()[i]   = static_cast<int64_t>(sd);
+        offset_h->data<int64_t>()[i] = static_cast<int64_t>(ofst);
     }
     auto seed = Buffer2torchTensor(clone({*seed_h, AllocationType::DEVICE}), false);
     auto offset = Buffer2torchTensor(clone({*offset_h, AllocationType::DEVICE}), false);

rtp_llm/cpp/devices/rocm_impl/test/ROCmSamplerTest.cc

@@ -136,7 +136,6 @@ TEST_F(CudaSamplerTest, testTopP) {
     ASSERT_NEAR(cum_log_probs_host[2], -5.02131, 1e-3);
     ASSERT_NEAR(cum_log_probs_host[3], -5.2682, 1e-3);
 
-    params.random_seed = nullopt;
     for (int i = 0; i < 100; i++) {
         device_->sampleGreedy(params);
         // printBuffer<int32_t>(*output_token_ids, "output_token_ids");

rtp_llm/cpp/kernels/rocm/sampling/BUILD

@@ -38,4 +38,5 @@ py_test(
     ],
     imports = ["."],
     python_version = "PY3",
+    # tags = ["rocm"],
 )