vulkan: Use L2-friendly workgroup scheduling for cm2 mulmat (#30215)
Copilot Setup Steps / copilot-setup-steps (push) Has been cancelled
Check Pre-Tokenizer Hashes / pre-tokenizer-hashes (push) Has been cancelled
Python check requirements.txt / check-requirements (push) Has been cancelled
Python Type-Check / python type-check (push) Has been cancelled
Update Operations Documentation / update-ops-docs (push) Has been cancelled
Copilot Setup Steps / copilot-setup-steps (push) Has been cancelled
Check Pre-Tokenizer Hashes / pre-tokenizer-hashes (push) Has been cancelled
Python check requirements.txt / check-requirements (push) Has been cancelled
Python Type-Check / python type-check (push) Has been cancelled
Update Operations Documentation / update-ops-docs (push) Has been cancelled
This commit is contained in:
@@ -45,6 +45,7 @@ const uint BNover2 = BN / 2;
|
||||
const uint BNover4 = enable_smaller_matrices ? (BN / 4) : (BN / 2);
|
||||
layout (constant_id = 5) const uint ALIGNED = 0;
|
||||
layout (constant_id = 6) const uint subgroup_size = 32;
|
||||
#define GROUP_A_BUDGET (8u * 1024u * 1024u)
|
||||
|
||||
#ifdef MULMAT_QUANT
|
||||
|
||||
@@ -258,8 +259,6 @@ void main() {
|
||||
const uint row_begin = tile_idx * BN;
|
||||
_ne1 = data_row_map[2 + 2 * tile_idx];
|
||||
const uint ic = 0;
|
||||
#else
|
||||
const uint ic = gl_WorkGroupID.y;
|
||||
#endif
|
||||
|
||||
#ifdef MUL_MAT_ID
|
||||
@@ -289,9 +288,33 @@ void main() {
|
||||
#endif
|
||||
|
||||
const uint blocks_m = (p.M + BM - 1) / BM;
|
||||
const uint ir = gl_WorkGroupID.x % blocks_m;
|
||||
const uint ik = gl_WorkGroupID.x / blocks_m;
|
||||
|
||||
#ifdef MUL_MAT_ID
|
||||
const uint ir = gl_WorkGroupID.x % blocks_m;
|
||||
#else
|
||||
// L2-friendly workgroup scheduling
|
||||
const uint blocks_n = (p.N + BN - 1) / BN;
|
||||
#if defined(MULMAT_QUANT)
|
||||
const uint a_panel_bytes = (BM * p.K / mm_quant_k()) * MmABlockBytes;
|
||||
#elif defined(DATA_A_F32)
|
||||
const uint a_panel_bytes = BM * p.K * 4;
|
||||
#elif QUANT_K == 1
|
||||
const uint a_panel_bytes = BM * p.K * 2;
|
||||
#else
|
||||
const uint a_panel_bytes = (BM * p.K) / 2 + (BM * p.K) / 16;
|
||||
#endif
|
||||
const uint group_m = clamp(GROUP_A_BUDGET / max(a_panel_bytes, 1u), 1u, min(blocks_m, 32u));
|
||||
const uint tiles_per_group = group_m * blocks_n;
|
||||
const uint lin = gl_WorkGroupID.y * blocks_m + (gl_WorkGroupID.x % blocks_m);
|
||||
const uint group_id = lin / tiles_per_group;
|
||||
const uint first_m = group_id * group_m;
|
||||
const uint gsize = min(blocks_m - first_m, group_m);
|
||||
const uint in_group = lin - group_id * tiles_per_group;
|
||||
const uint ir = first_m + in_group % gsize;
|
||||
const uint ic = in_group / gsize;
|
||||
#endif
|
||||
|
||||
#ifdef MUL_MAT_ID
|
||||
load_row_ids(row_begin);
|
||||
#endif
|
||||
|
||||
Reference in New Issue
Block a user