| {% if foldKind == "subgroup_full" or foldKind == "subgroup_butterfly" %} |
| enable subgroups; |
| {% endif %} |
| {{ env.wgsl.resourceDeclarations }} |
| |
| // The Engram gate, fused per (batch, sequence, hc_mult) row: |
| // dot = sum(RMSNorm(key) * RMSNorm(query)) / sqrt(HIDDEN) |
| // out = sigmoid(sign(dot) * sqrt(max(abs(dot), 1e-6))) * value |
| // Both inverse-RMS factors are row scalars, so they come out of the sum and the |
| // two sums of squares and the scaled cross term ride ONE pass as a single vec3 |
| // accumulator. Normalizing per element first would need two passes and a |
| // different rounding. |
| // |
| // Each row owns a fixed LANE GROUP rather than a whole workgroup. LANES_PER_ROW |
| // is the smallest power of two that covers the row, capped by the workgroup |
| // (and, on the subgroup routes, by the subgroup) so several rows are resident at |
| // once, and widened toward the workgroup when there are too few rows to fill the |
| // dispatch. The value words stay in registers across the fold. |
| // |
| // Every accumulation is float32 whatever the tensor type; only the store narrows. |
| const ROW_WORDS: u32 = {{ rowWords }}u; |
| const HIDDEN: f32 = {{ hidden }}.0; |
| const HC_MULT: u32 = {{ hcMult }}u; |
| const EPS: f32 = {{ epsilon }}; |
| const GATE_FLOOR: f32 = 0.000001; |
| const WG: u32 = {{ workgroupSize }}u; |
| const LANES_PER_ROW: u32 = {{ lanesPerRow }}u; |
| const ROWS_PER_GROUP: u32 = {{ rowsPerGroup }}u; |
| {% if foldKind == "workgroup_tree" %} |
| |
| // One contiguous LANES_PER_ROW slice per row, carrying all three accumulators as |
| // one vector. The fold never crosses a slice boundary, so the rows sharing this |
| // array do not interfere. |
| var<workgroup> partials: array<vec3<f32>, WG>; |
| {% endif %} |
| |
| fn sigmoid_safe(x: f32) -> f32 { |
| if (x >= 0.0) { |
| let z = exp(-x); |
| return 1.0 / (1.0 + z); |
| } |
| let z = exp(x); |
| return z / (1.0 + z); |
| } |
| |
| // sign(0.0) is 0.0 in WGSL, so a zero dot product gives an argument of exactly |
| // zero and therefore a gate of exactly 0.5. A copysign spelling would map it to |
| // +sqrt(GATE_FLOOR) instead, which is the divergence upstream calls out. |
| fn engram_gate(dotValue: f32) -> f32 { |
| return sigmoid_safe(sign(dotValue) * sqrt(max(abs(dotValue), GATE_FLOOR))); |
| } |
| |
| @compute @workgroup_size(WG) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) { |
| let lane = lid.x % LANES_PER_ROW; |
| let localRow = lid.x / LANES_PER_ROW; |
| // 2D-folded row index: wg.y carries the high bits past the dispatch's per-axis |
| // workgroup fold width. |
| let row = (wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u) * ROWS_PER_GROUP + localRow; |
| let rowValid = row < params.rows; |
| // A lane group past the last row folds an empty row and stores nothing, so the |
| // collectives and barriers below stay in uniform control flow. |
| let limit = select(0u, ROW_WORDS, rowValid); |
| let base = row * ROW_WORDS; |
| // The hyper-connection stream selects the scale rows; the token selects the one |
| // value row that every stream of that token shares. |
| let scaleBase = (row % HC_MULT) * ROW_WORDS; |
| let valueBase = (row / HC_MULT) * ROW_WORDS; |
| |
| // x: sum(key^2), y: sum(query^2), z: sum(key * key_scale * query * query_scale). |
| var acc = vec3<f32>(0.0); |
| {% if stageRow %} |
| {% for i in range(wordsPerLane) %} |
| {% if useVec4 %} |
| var vv{{ i }} = vec4<f32>(0.0); |
| {% else %} |
| var vv{{ i }} = 0.0; |
| {% endif %} |
| let c{{ i }} = lane + {{ i }}u * LANES_PER_ROW; |
| if (c{{ i }} < limit) { |
| {% if useVec4 %} |
| let kv{{ i }} = vec4<f32>(key[base + c{{ i }}]); |
| let qv{{ i }} = vec4<f32>(query[base + c{{ i }}]); |
| let ks{{ i }} = vec4<f32>(kscale[scaleBase + c{{ i }}]); |
| let qs{{ i }} = vec4<f32>(qscale[scaleBase + c{{ i }}]); |
| vv{{ i }} = vec4<f32>(value[valueBase + c{{ i }}]); |
| acc = acc + vec3<f32>(dot(kv{{ i }}, kv{{ i }}), dot(qv{{ i }}, qv{{ i }}), dot(kv{{ i }} * ks{{ i }}, qv{{ i }} * qs{{ i }})); |
| {% else %} |
| let kv{{ i }} = f32(key[base + c{{ i }}]); |
| let qv{{ i }} = f32(query[base + c{{ i }}]); |
| let ks{{ i }} = f32(kscale[scaleBase + c{{ i }}]); |
| let qs{{ i }} = f32(qscale[scaleBase + c{{ i }}]); |
| vv{{ i }} = f32(value[valueBase + c{{ i }}]); |
| acc = acc + vec3<f32>(kv{{ i }} * kv{{ i }}, qv{{ i }} * qv{{ i }}, (kv{{ i }} * ks{{ i }}) * (qv{{ i }} * qs{{ i }})); |
| {% endif %} |
| } |
| {% endfor %} |
| {% else %} |
| // The row's value words do not fit in registers at this width, so the statistic |
| // pass and the store pass each walk it. |
| for (var c = lane; c < limit; c = c + LANES_PER_ROW) { |
| {% if useVec4 %} |
| let kv = vec4<f32>(key[base + c]); |
| let qv = vec4<f32>(query[base + c]); |
| let ks = vec4<f32>(kscale[scaleBase + c]); |
| let qs = vec4<f32>(qscale[scaleBase + c]); |
| acc = acc + vec3<f32>(dot(kv, kv), dot(qv, qv), dot(kv * ks, qv * qs)); |
| {% else %} |
| let kv = f32(key[base + c]); |
| let qv = f32(query[base + c]); |
| let ks = f32(kscale[scaleBase + c]); |
| let qs = f32(qscale[scaleBase + c]); |
| acc = acc + vec3<f32>(kv * kv, qv * qv, (kv * ks) * (qv * qs)); |
| {% endif %} |
| } |
| {% endif %} |
| {% if foldKind == "subgroup_full" %} |
| let sums = subgroupAdd(acc); |
| {% elif foldKind == "subgroup_butterfly" %} |
| // Segmented butterfly: only the low bits of the lane index are exchanged, so each |
| // aligned group of LANES_PER_ROW lanes folds its own row and one subgroup carries |
| // several rows without a barrier or any workgroup memory. All three accumulators |
| // ride one shuffle. |
| var folded = acc; |
| {% for step in range(foldSteps) %} |
| folded = folded + subgroupShuffleXor(folded, 1u << {{ step }}u); |
| {% endfor %} |
| let sums = folded; |
| {% elif foldKind == "workgroup_tree" %} |
| partials[lid.x] = acc; |
| workgroupBarrier(); |
| // LANES_PER_ROW is a power of two, so the halving stride is uniform across the |
| // whole workgroup and every lane reaches every barrier. |
| var stride = LANES_PER_ROW / 2u; |
| loop { |
| if (stride == 0u) { |
| break; |
| } |
| if (lane < stride) { |
| partials[lid.x] = partials[lid.x] + partials[lid.x + stride]; |
| } |
| workgroupBarrier(); |
| stride = stride / 2u; |
| } |
| let sums = partials[localRow * LANES_PER_ROW]; |
| {% else %} |
| // One lane owns the whole row, so the accumulator already holds the row sums. |
| let sums = acc; |
| {% endif %} |
| let keyInv = inverseSqrt(sums.x / HIDDEN + EPS); |
| let queryInv = inverseSqrt(sums.y / HIDDEN + EPS); |
| let gate = engram_gate(sums.z * keyInv * queryInv / sqrt(HIDDEN)); |
| {% if stageRow %} |
| {% for i in range(wordsPerLane) %} |
| if (c{{ i }} < limit) { |
| {% if useVec4 %} |
| out[base + c{{ i }}] = {{ vectorScalar }}(vv{{ i }} * gate); |
| {% else %} |
| out[base + c{{ i }}] = {{ scalar }}(vv{{ i }} * gate); |
| {% endif %} |
| } |
| {% endfor %} |
| {% else %} |
| for (var c = lane; c < limit; c = c + LANES_PER_ROW) { |
| {% if useVec4 %} |
| out[base + c] = {{ vectorScalar }}(vec4<f32>(value[valueBase + c]) * gate); |
| {% else %} |
| out[base + c] = {{ scalar }}(f32(value[valueBase + c]) * gate); |
| {% endif %} |
| } |
| {% endif %} |
| } |
| |
