Source browser
examples/shaders/transformer_infer.cpp
Read the complete file below or download its original bytes. Line numbers provide stable links to a declaration or processing step.
File fingerprint
1
2// Generated by scripts/generate_transformer.py. Edit its shared function body.
3// Four-token attention, RMS normalization, feed-forward residual and classifier.
4#ifdef __cplusplus // Include the C++ runtime only for CPU compilation.
5#include "shdr.h" // Use the portable math types and kernel registration API.
6#endif // Keep GPU compilation free of C++ header contents.
7layout(local_size_x = 64) in; // Process independent sample or parameter indices in groups of 64.
8layout(std430, binding = 0) readonly buffer Inputs { // Declare the readonly model resource named by this block.
9 float inputs[]; // Read sixteen feature values for each four-token sample.
10}; // Close the input feature storage block.
11layout(std430, binding = 1) readonly buffer Parameters { // Declare the readonly model resource named by this block.
12 float parameters[]; // Read all 134 trainable parameter values.
13}; // Close the model parameter storage block.
14layout(std430, binding = 2) readonly buffer Labels { // Declare the readonly model resource named by this block.
15 uint labels[]; // Read the binary target class for each sequence.
16}; // Close the target label storage block.
17uniform uint batch_count; // Limit every dispatch to the declared number of sequences.
18
19layout(std430, binding = 3) writeonly buffer Predictions { // Declare the inference output.
20 vec4 predictions[]; // Store class probabilities, loss and the selected class.
21}; // Close the prediction output block.
22
23
24// Each vec2 carries a value and one analytic parameter derivative.
25vec2 ad_mul(vec2 a, vec2 b) { // Apply the product rule to values and their partial derivatives.
26 return vec2(a.x * b.x, a.y * b.x + a.x * b.y); // Keep the value product and both derivative contributions.
27} // End the product-rule helper.
28vec2 ad_div(vec2 a, vec2 b) { // Apply the quotient rule to one partial derivative.
29 return vec2(a.x / b.x, (a.y * b.x - a.x * b.y) / (b.x * b.x)); // Divide values and propagate the numerator and denominator derivatives.
30} // End the quotient-rule helper.
31vec2 ad_exp(vec2 a) { // Differentiate the exponential using its computed value.
32 float value = exp(a.x); // Compute the scalar value once before propagating its derivative.
33 return vec2(value, value * a.y); // Multiply the input derivative by the exponential value.
34} // End the exponential helper.
35vec2 ad_log(vec2 a) { // Differentiate the natural logarithm.
36 return vec2(log(a.x), a.y / a.x); // Divide the input derivative by its value.
37} // End the logarithm helper.
38vec2 ad_invsqrt(vec2 a) { // Differentiate inverse square root for the positive normalization input.
39 float value = inversesqrt(a.x); // Compute the scalar value once before propagating its derivative.
40 return vec2(value, -0.5 * value * value * value * a.y); // Use the derivative of a power with exponent minus one half.
41} // End the inverse-square-root helper.
42vec2 ad_relu(vec2 a) { // Keep the derivative only where the activation is positive.
43 return a.x > 0.0 ? a : vec2(0.0); // Use a zero value and derivative on the inactive branch.
44} // End the ReLU helper.
45vec2 weight(uint index, uint seed) { // Read a parameter and seed its selected partial derivative.
46 return vec2(parameters[index], index == seed ? 1.0 : 0.0); // Seed one for this parameter and zero for every other parameter.
47} // End the parameter derivative seed helper.
48
49// Parameter ranges: position0:16,Q16:32,K32:48,V48:64,O64:80,
50// RMSscale80:84,FFN1weight84:100,bias100:104,FFN2weight104:120,
51// bias120:124,classifierweight124:132,bias132:134.
52// The returned vec4 stores two logits followed by their two derivatives.
53vec4 transformer(uint row, uint seed) { // Evaluate two classifier logits and one parameter derivative.
54 vec2 x[16]; // Hold four tokens after adding their learned positional features.
55 vec2 q[16]; // Hold the four query vectors and their derivatives.
56 vec2 k[16]; // Hold the four key vectors and their derivatives.
57 vec2 v[16]; // Hold the four value vectors and their derivatives.
58 vec2 attended[16]; // Hold attention output after its projection and residual.
59 vec2 normalized[16]; // Hold each normalized feature and its derivative.
60 vec2 features[16]; // Hold the feed-forward output after its residual.
61 for (uint i = 0u; i < 16u; ++i) // Add learned positions to all sixteen input features.
62 x[i] = vec2(inputs[row * 16u + i], 0.0) + weight(i, seed); // Add the learned position to this input value.
63 for (uint token = 0u; token < 4u; ++token) // Project each of the four tokens into query, key and value features.
64 for (uint c = 0u; c < 4u; ++c) { // Accumulate query, key and value matrix products.
65 vec2 qvalue = vec2(0.0); // Start the query dot product at zero.
66 vec2 kvalue = vec2(0.0); // Start the key dot product at zero.
67 vec2 vvalue = vec2(0.0); // Start the value dot product at zero.
68 for (uint j = 0u; j < 4u; ++j) { // Accumulate query, key and value matrix products.
69 qvalue += ad_mul(x[token * 4u + j], weight(16u + j * 4u + c, seed)); // Accumulate the query projection and its derivative.
70 kvalue += ad_mul(x[token * 4u + j], weight(32u + j * 4u + c, seed)); // Accumulate the key projection and its derivative.
71 vvalue += ad_mul(x[token * 4u + j], weight(48u + j * 4u + c, seed)); // Accumulate the value projection and its derivative.
72 } // The query, key and value dot products are complete.
73 q[token * 4u + c] = qvalue; // Store the projected query feature.
74 k[token * 4u + c] = kvalue; // Store the projected key feature.
75 v[token * 4u + c] = vvalue; // Store the projected value feature.
76 } // All projected token features are now stored.
77 for (uint token = 0u; token < 4u; ++token) { // Evaluate this token's attention and feed-forward features.
78 vec2 scores[4]; // Keep one query-key score for each source token.
79 vec2 probabilities[4]; // Keep exponentiated attention scores before normalization.
80 float maximum = -1000000000.0; // Find a score shift that prevents exponential overflow.
81 for (uint source = 0u; source < 4u; ++source) { // Create an attention score for every source token.
82 scores[source] = vec2(0.0); // Start this scaled query-key dot product at zero.
83 for (uint c = 0u; c < 4u; ++c) // Sum the query-key product across four feature channels.
84 scores[source] += ad_mul(q[token * 4u + c], k[source * 4u + c]) * 0.5; // Accumulate the scaled query-key dot product.
85 maximum = max(maximum, scores[source].x); // Retain the largest attention score for stable softmax.
86 } // Every source token now has an attention score.
87 vec2 total = vec2(0.0); // Accumulate the denominator of the softmax.
88 for (uint source = 0u; source < 4u; ++source) { // Exponentiate shifted scores for the softmax denominator.
89 probabilities[source] = ad_exp(scores[source] - vec2(maximum, 0.0)); // Keep the shifted exponential and its derivative.
90 total += probabilities[source]; // Sum the exponentials for probability normalization.
91 } // The softmax denominator now includes every source score.
92 vec2 context[4]; // Combine source value vectors using their attention probabilities.
93 for (uint c = 0u; c < 4u; ++c) { // Build each channel of the attention value mixture.
94 context[c] = vec2(0.0); // Start this output feature mixture at zero.
95 for (uint source = 0u; source < 4u; ++source) // Mix source value vectors with their normalized attention probabilities.
96 context[c] += ad_mul(ad_div(probabilities[source], total), v[source * 4u + c]); // Weight source values by their normalized attention probabilities.
97 } // The attention mixture now contains four feature values.
98 for (uint c = 0u; c < 4u; ++c) { // Project attention output and preserve the input residual.
99 attended[token * 4u + c] = x[token * 4u + c]; // Start with the input token as a residual connection.
100 for (uint j = 0u; j < 4u; ++j) // Apply the learned attention output matrix to this feature.
101 attended[token * 4u + c] += ad_mul(context[j], weight(64u + j * 4u + c, seed)); // Add the learned attention output projection.
102 } // The attention projection and residual are complete.
103 vec2 square = vec2(0.00001, 0.0); // Start the mean squared feature value with a small positive epsilon.
104 for (uint c = 0u; c < 4u; ++c) // Accumulate the mean squared feature value for RMS normalization.
105 square += ad_mul(attended[token * 4u + c], attended[token * 4u + c]) * 0.25; // Add one quarter of this squared feature to the RMS estimate.
106 vec2 inverse = ad_invsqrt(square); // Compute the inverse RMS value and its derivative.
107 for (uint c = 0u; c < 4u; ++c) // Apply inverse RMS and the learned channel scale.
108 normalized[token * 4u + c] = // Scale the normalized feature with its learned RMS parameter.
109 ad_mul(ad_mul(attended[token * 4u + c], inverse), weight(80u + c, seed)); // Propagate both the inverse RMS and learned scale derivatives.
110 vec2 hidden[4]; // Hold four feed-forward hidden activations.
111 for (uint c = 0u; c < 4u; ++c) { // Evaluate each feed-forward hidden feature.
112 hidden[c] = weight(100u + c, seed); // Start the hidden neuron with its learned bias.
113 for (uint j = 0u; j < 4u; ++j) // Apply the first feed-forward matrix to normalized features.
114 hidden[c] += ad_mul(normalized[token * 4u + j], weight(84u + j * 4u + c, seed)); // Accumulate normalized input through the first feed-forward matrix.
115 hidden[c] = ad_relu(hidden[c]); // Apply the ReLU value and derivative together.
116 } // The four hidden features now include ReLU derivatives.
117 for (uint c = 0u; c < 4u; ++c) { // Start each feed-forward output with its residual and bias.
118 features[token * 4u + c] = normalized[token * 4u + c] + weight(120u + c, seed); // Start the residual output with its bias and normalized input.
119 for (uint j = 0u; j < 4u; ++j) // Project the hidden activations through the second feed-forward matrix.
120 features[token * 4u + c] += ad_mul(hidden[j], weight(104u + j * 4u + c, seed)); // Add the second feed-forward matrix projection.
121 } // The second feed-forward projection and residual are complete.
122 } // This token now has its final learned features.
123 vec2 pooled[4]; // Average token features before classification.
124 for (uint c = 0u; c < 4u; ++c) { // Mean-pool each feature across the four tokens.
125 pooled[c] = vec2(0.0); // Start this pooled feature at zero.
126 for (uint token = 0u; token < 4u; ++token) // Average this feature across the four token outputs.
127 pooled[c] += features[token * 4u + c] * 0.25; // Add one quarter of this token feature.
128 } // All token features have been mean-pooled.
129 vec2 logits[2]; // Hold the two class scores and their derivatives.
130 for (uint c = 0u; c < 2u; ++c) { // Build each of the two classifier scores.
131 logits[c] = weight(132u + c, seed); // Start the classifier score with its learned bias.
132 for (uint j = 0u; j < 4u; ++j) // Apply the classifier matrix to the pooled feature values.
133 logits[c] += ad_mul(pooled[j], weight(124u + j * 2u + c, seed)); // Accumulate pooled features through the classifier matrix.
134 } // Both class scores and their partials are complete.
135 return vec4(logits[0].x, logits[1].x, logits[0].y, logits[1].y); // Return both logits followed by their two partial derivatives.
136} // The shared transformer forward pass is complete.
137
138// Stable two-class cross entropy. The maximum shift cancels analytically.
139vec2 cross_entropy(vec4 logits, uint label) { // Evaluate the selected class loss and its derivative.
140 float shift = max(logits.x, logits.y); // Use the larger logit as a safe exponential shift.
141 vec2 first = vec2(logits.x, logits.z); // Keep the first class value or its shifted exponential.
142 vec2 second = vec2(logits.y, logits.w); // Keep the second class value or its shifted exponential.
143 vec2 total = ad_exp(first - vec2(shift, 0.0)) + ad_exp(second - vec2(shift, 0.0)); // Accumulate the denominator of the softmax.
144 return ad_log(total) + vec2(shift, 0.0) - (label == 0u ? first : second); // Subtract the selected class score from the stable log-sum-exp.
145} // End the stable cross-entropy helper.
146
147
148void main() { // Write the result belonging to this compute invocation.
149 uint row = gl_GlobalInvocationID.x; // Locate this sample in the batch.
150 if (row >= batch_count) // Ignore invocations outside the allocated batch.
151 return; // Stop padding invocations before processing a sample.
152 vec4 logits = transformer(row, 4294967295u); // Hold the two class scores and their derivatives.
153 float shift = max(logits.x, logits.y); // Use the larger logit as a safe exponential shift.
154 float first = exp(logits.x - shift); // Keep the first class value or its shifted exponential.
155 float second = exp(logits.y - shift); // Keep the second class value or its shifted exponential.
156 float total = first + second; // Accumulate the denominator of the softmax.
157 predictions[row] = vec4(first / total, // Publish both probabilities, the loss and the predicted class.
158 second / total, // Publish the probability of class one.
159 cross_entropy(logits, labels[row]).x, // Publish the selected class loss.
160 logits.y > logits.x ? 1.0 : 0.0); // Select the class with the larger score.
161} // Finish this prediction invocation.