@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cpp 151 lines 4.81 KB
Raw
1 #include "cpu_types.hpp"
2
3 #include <torch/library.h>
4
5 #include <cmath>
6 #include <cstdint>
7 #include <limits>
8 #include <mutex>
9 #include <random>
10
11 namespace {
12
13 constexpr int GUMBEL_TABLE_SIZE = 1 << 20;
14 constexpr int GUMBEL_TABLE_MASK = GUMBEL_TABLE_SIZE - 1;
15
16 static float g_gumbel_table[GUMBEL_TABLE_SIZE];
17 static std::once_flag g_gumbel_init_flag;
18
19 static void init_gumbel_table() {
20 std::mt19937 rng(42);
21 std::uniform_real_distribution<float> uniform(1.0e-7f, 1.0f - 1.0e-7f);
22 for (int i = 0; i < GUMBEL_TABLE_SIZE; ++i) {
23 float u = uniform(rng);
24 g_gumbel_table[i] = -std::log(-std::log(u));
25 }
26 }
27
28 static inline void ensure_gumbel_table() {
29 std::call_once(g_gumbel_init_flag, init_gumbel_table);
30 }
31
32 // Fused Gumbel-max kernel
33
34 static void fused_gumbel_argmax_kernel(int64_t* __restrict__ output,
35 const float* __restrict__ logits,
36 const int64_t* __restrict__ seeds,
37 const int64_t batch_size,
38 const int64_t vocab_size) {
39 ensure_gumbel_table();
40 constexpr int VEC_ELEM_NUM = vec_op::FP32Vec8::VEC_ELEM_NUM;
41 const int64_t vec_end = vocab_size - (vocab_size % VEC_ELEM_NUM);
42
43 #pragma omp parallel for schedule(static)
44 for (int64_t b = 0; b < batch_size; ++b) {
45 const float* row = logits + b * vocab_size;
46 const uint32_t seed = static_cast<uint32_t>(seeds[b]);
47
48 float best_score = -std::numeric_limits<float>::infinity();
49 int64_t best_idx = 0;
50
51 for (int64_t i = 0; i < vec_end; i += VEC_ELEM_NUM) {
52 vec_op::FP32Vec8 logit_vec(row + i);
53
54 float gumbel_buf[VEC_ELEM_NUM];
55 for (int k = 0; k < VEC_ELEM_NUM; ++k)
56 gumbel_buf[k] = g_gumbel_table[(seed + i + k) & GUMBEL_TABLE_MASK];
57 vec_op::FP32Vec8 gumbel_vec(gumbel_buf);
58
59 vec_op::FP32Vec8 score = logit_vec + gumbel_vec;
60
61 float score_buf[VEC_ELEM_NUM];
62 score.save(score_buf);
63 for (int k = 0; k < VEC_ELEM_NUM; ++k) {
64 if (score_buf[k] > best_score) {
65 best_score = score_buf[k];
66 best_idx = i + k;
67 }
68 }
69 }
70
71 for (int64_t i = vec_end; i < vocab_size; ++i) {
72 float s = row[i] + g_gumbel_table[(seed + i) & GUMBEL_TABLE_MASK];
73 if (s > best_score) {
74 best_score = s;
75 best_idx = i;
76 }
77 }
78
79 output[b] = best_idx;
80 }
81 }
82
83 static void greedy_argmax_kernel(int64_t* __restrict__ output,
84 const float* __restrict__ logits,
85 const int64_t batch_size,
86 const int64_t vocab_size) {
87 constexpr int VEC_ELEM_NUM = vec_op::FP32Vec16::VEC_ELEM_NUM;
88 const int64_t vec_end = vocab_size - (vocab_size % VEC_ELEM_NUM);
89
90 #pragma omp parallel for schedule(static)
91 for (int64_t b = 0; b < batch_size; ++b) {
92 const float* row = logits + b * vocab_size;
93
94 vec_op::FP32Vec16 vmax(-std::numeric_limits<float>::infinity());
95 for (int64_t i = 0; i < vec_end; i += VEC_ELEM_NUM) {
96 vmax = vmax.max(vec_op::FP32Vec16(row + i));
97 }
98 float best_val = vmax.reduce_max();
99 for (int64_t i = vec_end; i < vocab_size; ++i) {
100 if (row[i] > best_val) {
101 best_val = row[i];
102 }
103 }
104
105 int64_t best_idx = 0;
106 for (int64_t i = 0; i < vocab_size; ++i) {
107 if (row[i] == best_val) {
108 best_idx = i;
109 break;
110 }
111 }
112 output[b] = best_idx;
113 }
114 }
115
116 } // namespace
117
118 torch::Tensor fused_gumbel_argmax(const torch::Tensor& logits,
119 const torch::Tensor& seeds) {
120 TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [batch, vocab]");
121 TORCH_CHECK(logits.scalar_type() == torch::kFloat32,
122 "logits must be float32");
123 TORCH_CHECK(seeds.dim() == 1 && seeds.size(0) == logits.size(0),
124 "seeds must be 1-D with batch_size elements");
125
126 auto logits_contig = logits.contiguous();
127 auto output = torch::empty({logits_contig.size(0)}, torch::kInt64);
128 fused_gumbel_argmax_kernel(
129 output.data_ptr<int64_t>(), logits_contig.data_ptr<float>(),
130 seeds.data_ptr<int64_t>(), logits_contig.size(0), logits_contig.size(1));
131 return output;
132 }
133
134 torch::Tensor greedy_argmax(const torch::Tensor& logits) {
135 TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [batch, vocab]");
136 TORCH_CHECK(logits.scalar_type() == torch::kFloat32,
137 "logits must be float32");
138
139 // A single row is faster in ATen: OpenMP fork/join dominates the scan
140 // (~1 ms vs ~50 µs on s390x at vocab=32k).
141 if (logits.size(0) <= 1) {
142 return logits.argmax(/*dim=*/-1);
143 }
144
145 auto logits_contig = logits.contiguous();
146 auto output = torch::empty({logits_contig.size(0)}, torch::kInt64);
147 greedy_argmax_kernel(output.data_ptr<int64_t>(),
148 logits_contig.data_ptr<float>(), logits_contig.size(0),
149 logits_contig.size(1));
150 return output;
151 }