-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_transformer.cpp
More file actions
84 lines (74 loc) · 2.89 KB
/
Copy pathtest_transformer.cpp
File metadata and controls
84 lines (74 loc) · 2.89 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
#include <cstdio>
#include <cmath>
#include <vector>
#include "transformer.h"
static int g_fails = 0;
#define CHECK(cond) do { if (!(cond)) { std::fprintf(stderr, "FAIL %s:%d %s\n", __FILE__, __LINE__, #cond); ++g_fails; } } while(0)
int main() {
using namespace muscriptor;
// Sin embedding shape
{
std::vector<int64_t> pos = {0, 1, 2, 3};
auto emb = create_sin_embedding(pos, 64);
CHECK(emb.size() == 4 * 64);
// position 0 → cos(0)=1 for first half-dim entry? phase=0 for i=0
CHECK(std::abs(emb[0] - 1.f) < 1e-5f);
}
// Tiny transformer offline vs streaming shapes
{
TransformerConfig cfg;
cfg.dim = 32;
cfg.num_heads = 4;
cfg.num_layers = 1;
cfg.dim_ff = 64;
// Random-ish weights (fixed)
std::vector<float> in_proj(3 * 32 * 32, 0.01f);
std::vector<float> out_proj(32 * 32, 0.01f);
for (int i = 0; i < 32; ++i) out_proj[i * 32 + i] = 1.f;
std::vector<float> n1w(32, 1.f), n1b(32, 0.f), n2w(32, 1.f), n2b(32, 0.f);
std::vector<float> l1(64 * 32, 0.01f), l2(32 * 64, 0.01f);
LayerWeights lw;
lw.attn.in_proj = in_proj.data();
lw.attn.out_proj = out_proj.data();
lw.norm1_w = n1w.data();
lw.norm1_b = n1b.data();
lw.norm2_w = n2w.data();
lw.norm2_b = n2b.data();
lw.linear1 = l1.data();
lw.linear2 = l2.data();
std::vector<float> x(8 * 32, 0.1f);
auto offline = transformer_forward(x.data(), 8, cfg, {lw}, nullptr, 0);
CHECK(offline.size() == 8 * 32);
KVCache cache;
cache.init(1, 4, 8, 16);
std::vector<float> streamed;
for (int t = 0; t < 8; ++t) {
auto y = transformer_forward(x.data() + t * 32, 1, cfg, {lw}, &cache, 0);
// position_offset comes from cache->offset before increment...
// actually transformer_forward uses position_offset arg; we pass 0 always
// but cache offset advances. Need to pass cache.offset before call.
streamed.insert(streamed.end(), y.begin(), y.end());
}
// Re-run streaming correctly with position offset
cache.reset();
streamed.clear();
for (int t = 0; t < 8; ++t) {
int pos = cache.offset;
auto y = transformer_forward(x.data() + t * 32, 1, cfg, {lw}, &cache, pos);
streamed.insert(streamed.end(), y.begin(), y.end());
}
CHECK(streamed.size() == offline.size());
// Values should be close (same causal path)
float max_diff = 0.f;
for (size_t i = 0; i < offline.size(); ++i) {
max_diff = std::max(max_diff, std::abs(offline[i] - streamed[i]));
}
CHECK(max_diff < 1e-3f);
}
if (g_fails) {
std::fprintf(stderr, "%d failures\n", g_fails);
return 1;
}
std::printf("ok transformer\n");
return 0;
}