Spaces:
Running on Zero
Running on Zero
| layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; | |
| layout (binding = 0) readonly buffer A {block_nvfp4 data_a[];}; | |
| layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; | |
| void main() { | |
| const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; | |
| init_iq_shmem(gl_WorkGroupSize); | |
| const uint tid = gl_LocalInvocationID.x % 64; | |
| const uint sub = tid / 16; | |
| const uint ir = tid % 16; | |
| const uint ib = 16 * i + ir; | |
| if (ib >= p.nel / 64) { | |
| return; | |
| } | |
| const uint q_idx = 8 * sub; | |
| const uint b_idx = 1024 * i + 64 * ir + 16 * sub; | |
| const float d = ue4m3_to_fp32(data_a[ib].d[sub]); | |
| [[unroll]] for (uint l = 0; l < 8; ++l) { | |
| data_b[b_idx + l + 0] = D_TYPE(d * 0.5 * float(kvalues_mxfp4[data_a[ib].qs[q_idx + l] & 0xF])); | |
| data_b[b_idx + l + 8] = D_TYPE(d * 0.5 * float(kvalues_mxfp4[data_a[ib].qs[q_idx + l] >> 4])); | |
| } | |
| } | |