wehub-resource-sync
64 lines
1.8 KiB
Plaintext
64 lines
1.8 KiB
Plaintext
layout(set=0, binding=0) readonly buffer s0 {
|
|
FLOAT4 data[];
|
|
} uKey;
|
|
|
|
layout(set=0, binding=1) readonly buffer s1 {
|
|
FLOAT4 data[];
|
|
} uValue;
|
|
|
|
layout(set=0, binding=2) writeonly buffer s2 {
|
|
FLOAT4 data[];
|
|
} uCacheKey;
|
|
|
|
layout(set=0, binding=3) writeonly buffer s3 {
|
|
FLOAT4 data[];
|
|
} uCacheValue;
|
|
|
|
layout(set=0, binding=4) readonly uniform constBuffer {
|
|
ivec4 s0; // qLen, kLen, headNum, kvHeadNum
|
|
ivec4 s1; // headDim, group, pastLen, totalLen
|
|
ivec4 s2; // maskQlen, maskKvlen, hasMask, cacheMaxLen
|
|
vec4 f0; // scale, 0, 0, 0
|
|
} uConst;
|
|
|
|
layout(local_size_x = 8, local_size_y = 1, local_size_z = 1) in;
|
|
|
|
void main() {
|
|
const int d4 = int(gl_GlobalInvocationID.x);
|
|
const int token = int(gl_GlobalInvocationID.y);
|
|
const int kvh = int(gl_GlobalInvocationID.z);
|
|
|
|
const int headDim = uConst.s1.x;
|
|
const int kLen = uConst.s0.y;
|
|
const int kvHeadNum = uConst.s0.w;
|
|
const int pastLen = uConst.s1.z;
|
|
const int maxLen = uConst.s2.w;
|
|
if (token >= kLen || kvh >= kvHeadNum) {
|
|
return;
|
|
}
|
|
const int d = d4 * 4;
|
|
if (d + 3 >= headDim) {
|
|
return;
|
|
}
|
|
if (maxLen <= 0) {
|
|
return;
|
|
}
|
|
|
|
// input: [1, kLen, kvHeadNum, headDim]
|
|
const int inBase = (token * kvHeadNum + kvh) * headDim + d;
|
|
const int cacheToken = pastLen + token;
|
|
if (cacheToken >= maxLen) {
|
|
return;
|
|
}
|
|
|
|
// cacheK (packed): [kvHeadNum, headDim/4, maxLen, 4]
|
|
const int d4Size = headDim / 4;
|
|
const int outBase = (kvh * d4Size + d4) * maxLen + cacheToken;
|
|
uCacheKey.data[outBase] = uKey.data[inBase / 4];
|
|
|
|
// cacheV (kvh-major, packed-D): [kvHeadNum, maxLen, headDim/4, 4]
|
|
// Index: ((kvh * maxLen + token) * d4Size + d4)
|
|
const int outBaseV4 = (kvh * maxLen + cacheToken) * d4Size + d4;
|
|
uCacheValue.data[outBaseV4] = uValue.data[inBase / 4];
|
|
}
|