feat: Add cpu random sample op

src/base/random_sample.h

@@ -0,0 +1,160 @@
+#ifndef INFINI_OPS_BASE_RANDOM_SAMPLE_H_
+#define INFINI_OPS_BASE_RANDOM_SAMPLE_H_
+
+#include <cstdint>
+#include <optional>
+
+#include "operator.h"
+
+namespace infini::ops {
+
+class RandomSample : public Operator<RandomSample> {
+ public:
+  // clang-format off
+  //
+  // logits:  [batch_size, vocab_size] or [vocab_size] (batch_size=1)
+  // out:     [batch_size]  sampled token ids (int32/int64)
+  // valid:   [batch_size]  uint8 (0 or 1), whether sample is valid
+  //
+  // Per-batch parameters support two modes:
+  //   - optional<Tensor> has value: per-batch tensor of shape [batch_size]
+  //   - optional<Tensor> is nullopt: use the scalar _val for all requests
+  //
+  // When both (optional<Tensor> == nullopt) and (_val == default),
+  // the corresponding filtering is disabled.
+  //
+  // seed:   per-request RNG seed. Same seed + offset produces identical
+  //         results (reproducibility). Different requests should use
+  //         different seeds.
+  // offset: per-step counter, increments each decode step within a request.
+  //         Ensures different steps produce different samples even with the
+  //         same seed.
+  //
+  // clang-format on
+  RandomSample(const Tensor logits, Tensor out, Tensor valid,
+               std::optional<Tensor> temperature, float temperature_val,
+               std::optional<Tensor> top_k, int top_k_val,
+               std::optional<Tensor> top_p, float top_p_val,
+               std::optional<Tensor> min_p, float min_p_val, std::uint64_t seed,
+               std::uint64_t offset, bool deterministic)
+      : logits_dtype_{logits.dtype()},
+        out_dtype_{out.dtype()},
+        ndim_{logits.ndim()},
+        batch_size_{ndim_ == 2 ? logits.size(-2) : 1},
+        vocab_size_{logits.size(-1)},
+        logits_strides_{logits.strides()},
+        temperature_{temperature},
+        temperature_val_{temperature_val},
+        top_k_{top_k},
+        top_k_val_{top_k_val},
+        top_p_{top_p},
+        top_p_val_{top_p_val},
+        min_p_{min_p},
+        min_p_val_{min_p_val},
+        seed_{seed},
+        offset_{offset},
+        deterministic_{deterministic} {
+    assert((ndim_ == 1 || ndim_ == 2) &&
+           "`RandomSample` requires 1D [vocab_size] or 2D [batch, vocab_size] "
+           "logits");
+    assert(out.ndim() == 1 && out.size(0) == batch_size_ &&
+           "`RandomSample` requires 1D output [batch_size]");
+    assert(valid.ndim() == 1 && valid.size(0) == batch_size_ &&
+           "`RandomSample` requires 1D valid [batch_size]");
+    assert((out_dtype_ == DataType::kInt32 || out_dtype_ == DataType::kInt64) &&
+           "`RandomSample` requires int32 or int64 output");
+    ValidateParams(temperature, top_k, top_p, min_p);
+  }
+
+  // Simplified constructor: no filtering, default temperature.
+  RandomSample(const Tensor logits, Tensor out, Tensor valid,
+               std::uint64_t seed, std::uint64_t offset)
+      : RandomSample{logits, out,          valid, std::nullopt,
+                     1.0f,   std::nullopt, 0,     std::nullopt,
+                     1.0f,   std::nullopt, 0.0f,  seed,
+                     offset, false} {}
+
+  virtual void operator()(const Tensor logits, Tensor out, Tensor valid,
+                          std::optional<Tensor> temperature,
+                          float temperature_val, std::optional<Tensor> top_k,
+                          int top_k_val, std::optional<Tensor> top_p,
+                          float top_p_val, std::optional<Tensor> min_p,
+                          float min_p_val, std::uint64_t seed,
+                          std::uint64_t offset, bool deterministic) const = 0;
+
+  virtual void operator()(const Tensor logits, Tensor out, Tensor valid,
+                          std::uint64_t seed, std::uint64_t offset) const {
+    return operator()(logits, out, valid, temperature_, temperature_val_,
+                      top_k_, top_k_val_, top_p_, top_p_val_, min_p_,
+                      min_p_val_, seed, offset, deterministic_);
+  }
+
+ protected:
+  static void ValidateIntParam(std::optional<Tensor> t,
+                               Tensor::Size batch_size) {
+    if (!t.has_value()) return;
+    const auto& tensor = *t;
+    assert(tensor.ndim() == 1 && tensor.size(0) == batch_size &&
+           "per-batch int param must be 1D [batch_size]");
+    assert((tensor.dtype() == DataType::kInt32 ||
+            tensor.dtype() == DataType::kInt64) &&
+           "per-batch int param must be int32 or int64");
+  }
+
+  static void ValidateFloatParam(std::optional<Tensor> t,
+                                 Tensor::Size batch_size) {
+    if (!t.has_value()) return;
+    const auto& tensor = *t;
+    assert(tensor.ndim() == 1 && tensor.size(0) == batch_size &&
+           "per-batch float param must be 1D [batch_size]");
+    assert((tensor.dtype() == DataType::kFloat32 ||
+            tensor.dtype() == DataType::kFloat64 ||
+            tensor.dtype() == DataType::kFloat16 ||
+            tensor.dtype() == DataType::kBFloat16) &&
+           "per-batch float param must be float16/bfloat16/float32/float64");
+  }
+
+  void ValidateParams(std::optional<Tensor> temperature,
+                      std::optional<Tensor> top_k, std::optional<Tensor> top_p,
+                      std::optional<Tensor> min_p) const {
+    ValidateFloatParam(temperature, batch_size_);
+    ValidateIntParam(top_k, batch_size_);
+    ValidateFloatParam(top_p, batch_size_);
+    ValidateFloatParam(min_p, batch_size_);
+  }
+
+  const DataType logits_dtype_;
+
+  const DataType out_dtype_;
+
+  Tensor::Size ndim_{0};
+
+  Tensor::Size batch_size_{1};
+
+  Tensor::Size vocab_size_{0};
+
+  Tensor::Strides logits_strides_;
+
+  // Per-batch or scalar sampling parameters.
+  std::optional<Tensor> temperature_;
+  float temperature_val_{1.0f};
+
+  std::optional<Tensor> top_k_;
+  int top_k_val_{0};
+
+  std::optional<Tensor> top_p_;
+  float top_p_val_{1.0f};
+
+  std::optional<Tensor> min_p_;
+  float min_p_val_{0.0f};
+
+  std::uint64_t seed_{0};
+
+  std::uint64_t offset_{0};
+
+  bool deterministic_{false};
+};
+
+}  // namespace infini::ops
+
+#endif