| {% macro read_scalar(name, index, dtype) %} |
| {% if dtype == "float16" %} |
| f32({{ name }}[{{ index }}]){% else %}{{ name }}[{{ index }}]{% endif %} |
| {% endmacro %}{% if queryDtype == "float16" %} |
| enable f16; |
| {% endif %} |
| {{ env.wgsl.resourceDeclarations }} |
| |
| // com.microsoft.LinearAttention, chunked prefill: within-chunk decay prefix. |
| // gexp[b, t, c] = exp(sum of decay[b, s, c] for s from the chunk start through t) |
| // Every later pass reads the recurrence's decay through this one buffer. Writing |
| // exp(prefix) rather than the prefix itself keeps the per-element exponential out |
| // of the O(chunk^2) inner loops, where it would cost one transcendental per |
| // multiply-add instead of one per element. |
| const CHUNK: u32 = {{ chunkSize }}u; |
| const WG: u32 = {{ workgroupSize }}u; |
| |
| @compute @workgroup_size(WG, 1, 1) |
| fn main( |
| @builtin(workgroup_id) wg: vec3<u32>, |
| @builtin(local_invocation_id) lid: vec3<u32>, |
| ) { |
| // 2D-folded flat (batch * chunk) index: wg.y carries the high bits past the |
| // per-axis dispatch fold width. |
| let flat = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u; |
| let num_chunks = (params.seqLength + CHUNK - 1u) / CHUNK; |
| let chunk = flat % num_chunks; |
| let batch = flat / num_chunks; |
| if (batch >= params.batchSize) { |
| return; |
| } |
| let t0 = chunk * CHUNK; |
| let t1 = min(t0 + CHUNK, params.seqLength); |
| let packed = params.decayPackedDim; |
| |
| // One thread owns a decay column and walks the chunk in order: the prefix is |
| // serial in t but independent across columns, so the whole chunk grid runs at once. |
| for (var col = lid.x; col < packed; col = col + WG) { |
| var prefix = 0.0; |
| for (var t = t0; t < t1; t = t + 1u) { |
| let idx = (batch * params.seqLength + t) * packed + col; |
| prefix = prefix + {{ read_scalar("decay", "idx", queryDtype) }}; |
| gexp[idx] = exp(prefix); |
| } |
| } |
| } |
| |
