-
Notifications
You must be signed in to change notification settings - Fork 44
Expand file tree
/
Copy pathspeed_int_vs_float.cpp
More file actions
169 lines (140 loc) · 6.51 KB
/
Copy pathspeed_int_vs_float.cpp
File metadata and controls
169 lines (140 loc) · 6.51 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
/**
* Speed benchmark: Integer-domain Q4xQ8 attention vs FP32 dequant+dot
*
* Compares three paths:
* 1. FP32 dot product (baseline, no quantization overhead)
* 2. Dequantize + FP32 dot (old quantized path)
* 3. Integer Q4xQ8 dot (new integer-domain path)
*
* On ARM NEON (Apple M-series), the integer path should be 3-5x faster
* than the dequant path due to avoiding FP32 conversion in the inner loop.
*/
#include <cstdio>
#include <cmath>
#include <cstring>
#include <vector>
#include <chrono>
extern "C" {
#include "turboquant/turboquant.h"
void tq_uniform_4b_quantize_ref(const float* src, void* dst, int n);
void tq_uniform_4b_dequantize_ref(const void* src, float* dst, int n);
void tq_quantize_query_q8(const float* query, int8_t* q8_out,
float* scale_out, float* sum_out, int n);
void tq_uniform_4b_attention_int_ref(const float* query, const void* kv,
float* scores, int seq_len, int head_dim);
#ifdef __ARM_NEON
void tq_uniform_4b_attention_neon(const float* query, const void* kv,
float* scores, int seq_len, int head_dim);
void tq_uniform_4b_attention_int_neon(const float* query, const void* kv,
float* scores, int seq_len, int head_dim);
#endif
}
using Clock = std::chrono::high_resolution_clock;
static double elapsed_ms(Clock::time_point start, Clock::time_point end) {
return std::chrono::duration<double, std::milli>(end - start).count();
}
/* FP32 baseline: no quantization, just dot product */
static void fp32_attention(const float* query, const float* keys,
float* scores, int seq_len, int head_dim) {
for (int s = 0; s < seq_len; s++) {
float dot = 0;
for (int d = 0; d < head_dim; d++)
dot += query[d] * keys[s * head_dim + d];
scores[s] = dot;
}
}
/* Dequant + dot: the old path */
static void dequant_attention(const float* query, const block_tq_uniform_4b* blocks,
float* scores, int seq_len, int head_dim) {
for (int s = 0; s < seq_len; s++) {
float deq[256];
tq_uniform_4b_dequantize_ref(&blocks[s], deq, head_dim);
float dot = 0;
for (int d = 0; d < head_dim; d++)
dot += query[d] * deq[d];
scores[s] = dot;
}
}
static void benchmark_config(int head_dim, int seq_len, int warmup_iters, int bench_iters) {
printf("\n--- head_dim=%d, seq_len=%d ---\n", head_dim, seq_len);
/* Setup data */
std::vector<float> query(head_dim);
for (int i = 0; i < head_dim; i++) query[i] = cosf(i * 0.05f);
std::vector<float> fp32_keys(seq_len * head_dim);
std::vector<block_tq_uniform_4b> q4_blocks(seq_len);
for (int s = 0; s < seq_len; s++) {
for (int d = 0; d < head_dim; d++) {
fp32_keys[s * head_dim + d] = sinf(s * 0.1f + d * 0.05f);
}
tq_uniform_4b_quantize_ref(&fp32_keys[s * head_dim], &q4_blocks[s], head_dim);
}
std::vector<float> scores(seq_len);
/* --- Benchmark FP32 baseline --- */
for (int i = 0; i < warmup_iters; i++)
fp32_attention(query.data(), fp32_keys.data(), scores.data(), seq_len, head_dim);
auto t0 = Clock::now();
for (int i = 0; i < bench_iters; i++)
fp32_attention(query.data(), fp32_keys.data(), scores.data(), seq_len, head_dim);
auto t1 = Clock::now();
double fp32_ms = elapsed_ms(t0, t1) / bench_iters;
/* --- Benchmark dequant+dot --- */
for (int i = 0; i < warmup_iters; i++)
dequant_attention(query.data(), q4_blocks.data(), scores.data(), seq_len, head_dim);
t0 = Clock::now();
for (int i = 0; i < bench_iters; i++)
dequant_attention(query.data(), q4_blocks.data(), scores.data(), seq_len, head_dim);
t1 = Clock::now();
double dequant_ms = elapsed_ms(t0, t1) / bench_iters;
/* --- Benchmark integer Q4xQ8 (ref) --- */
for (int i = 0; i < warmup_iters; i++)
tq_uniform_4b_attention_int_ref(query.data(), q4_blocks.data(),
scores.data(), seq_len, head_dim);
t0 = Clock::now();
for (int i = 0; i < bench_iters; i++)
tq_uniform_4b_attention_int_ref(query.data(), q4_blocks.data(),
scores.data(), seq_len, head_dim);
t1 = Clock::now();
double int_ref_ms = elapsed_ms(t0, t1) / bench_iters;
printf(" FP32 baseline: %8.3f ms\n", fp32_ms);
printf(" Dequant+dot (old): %8.3f ms\n", dequant_ms);
printf(" Int Q4xQ8 (ref): %8.3f ms (%.2fx vs dequant)\n",
int_ref_ms, dequant_ms / int_ref_ms);
#ifdef __ARM_NEON
/* --- Benchmark NEON dequant+dot --- */
for (int i = 0; i < warmup_iters; i++)
tq_uniform_4b_attention_neon(query.data(), q4_blocks.data(),
scores.data(), seq_len, head_dim);
t0 = Clock::now();
for (int i = 0; i < bench_iters; i++)
tq_uniform_4b_attention_neon(query.data(), q4_blocks.data(),
scores.data(), seq_len, head_dim);
t1 = Clock::now();
double neon_dequant_ms = elapsed_ms(t0, t1) / bench_iters;
/* --- Benchmark NEON integer Q4xQ8 --- */
for (int i = 0; i < warmup_iters; i++)
tq_uniform_4b_attention_int_neon(query.data(), q4_blocks.data(),
scores.data(), seq_len, head_dim);
t0 = Clock::now();
for (int i = 0; i < bench_iters; i++)
tq_uniform_4b_attention_int_neon(query.data(), q4_blocks.data(),
scores.data(), seq_len, head_dim);
t1 = Clock::now();
double neon_int_ms = elapsed_ms(t0, t1) / bench_iters;
printf(" NEON dequant+dot: %8.3f ms (%.2fx vs ref dequant)\n",
neon_dequant_ms, dequant_ms / neon_dequant_ms);
printf(" NEON Int Q4xQ8: %8.3f ms (%.2fx vs NEON dequant, %.2fx vs FP32)\n",
neon_int_ms, neon_dequant_ms / neon_int_ms, fp32_ms / neon_int_ms);
#endif
printf("\n");
}
int main() {
printf("=== Integer-Domain Attention Speed Benchmark ===\n");
printf("Comparing: FP32 | Dequant+Dot | Integer Q4xQ8\n");
/* Typical LLM configurations */
benchmark_config(128, 64, 10, 1000); /* Small: GQA head_dim=128, short ctx */
benchmark_config(128, 512, 10, 200); /* Medium: head_dim=128, medium ctx */
benchmark_config(128, 2048, 5, 50); /* Large: head_dim=128, long ctx */
benchmark_config(128, 8192, 3, 20); /* XL: head_dim=128, very long ctx */
printf("=== Benchmark Complete ===\n");
return 0;
}