File size: 3,526 Bytes
d882115 | 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 | {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
// 2D-folded flat index: gid.y carries the high bits past the dispatch's
// per-axis workgroup fold width.
let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
if ({{ name }} >= {{ bound }}) {
return;
}{% endmacro %}
{{ env.wgsl.resourceDeclarations }}
// Cache maintenance for decode self attention, in two roles from one template.
//
// retain — copy the whole past buffer into the present buffer at capacity.
// The rows beyond the live length are copied too: the operator's
// contract is that present equals past everywhere except the one
// appended row, and upstream's own kernel memcpys the full buffer.
// The pass is a pure copy, so it declares `viewAlias`: it is elided
// when `past_key` and `present_key` resolve to one buffer.
// append — write this token's key and value rows at index past_sequence_length,
// which arrives as tensor data and so is read here rather than baked.
//
// Key and value share one dispatch: they have identical geometry, and one
// invocation touching both halves keeps the index arithmetic in registers.
const WG: u32 = {{ workgroupSize }}u;
{% if appendRow %}
const HEAD_SIZE_V4: u32 = {{ headSizeV4 }}u;
const NUM_HEADS: u32 = {{ numHeads }}u;
const HIDDEN_V4: u32 = NUM_HEADS * HEAD_SIZE_V4;
@compute @workgroup_size(WG)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
// One invocation owns one vec4 of one (batch, head) row.
{{ flat_index_2d("WG", "index", "params.rowCountV4") }}
let d4 = index % HEAD_SIZE_V4;
let head = (index / HEAD_SIZE_V4) % NUM_HEADS;
let batch = index / HIDDEN_V4;
// past_sequence_length is a single value shared by every batch row.
let slot = u32(past_sequence_length[0]);
if (slot >= params.capacity) {
return;
}
let cacheIndex = ((batch * NUM_HEADS + head) * params.capacity + slot) * HEAD_SIZE_V4 + d4;
present_key[cacheIndex] = key[index];
present_value[cacheIndex] = value[index];
}
{% else %}
// Each lane stages ITEMS words of each cache before storing any of them, so that
// many loads are in flight before it has to wait. One word per lane leaves the
// copy latency-bound wherever the caches stay resident in cache; where the copy is
// DRAM-bound the depth is neutral. Consecutive lanes still read consecutive words;
// the stride between a lane's own items is a whole workgroup. ITEMS is bounded so
// the dispatch keeps a floor number of workgroups, since raising it divides them.
const ITEMS: u32 = {{ itemsPerLane }}u;
@compute @workgroup_size(WG)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let block = (wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u) * WG * ITEMS + lid.x;
var stagedKey: array<vec4<{{ cacheScalar }}>, ITEMS>;
var stagedValue: array<vec4<{{ cacheScalar }}>, ITEMS>;
{% for i in range(itemsPerLane) %}
{
let index = block + {{ i }}u * WG;
if (index < params.countV4) {
stagedKey[{{ i }}] = past_key[index];
stagedValue[{{ i }}] = past_value[index];
}
}
{% endfor %}
{% for i in range(itemsPerLane) %}
{
let index = block + {{ i }}u * WG;
if (index < params.countV4) {
present_key[index] = stagedKey[{{ i }}];
present_value[index] = stagedValue[{{ i }}];
}
}
{% endfor %}
}
{% endif %}
|