diff --git a/engine/compile_m.py b/engine/compile_m.py index b35e38c..66b4ac7 100644 --- a/engine/compile_m.py +++ b/engine/compile_m.py @@ -32,6 +32,8 @@ from test_m import pack_layer, defines EXACT = os.environ.get('M_EXACT') == '1' FAST = [] if EXACT or '-DM_RELAXED' in os.environ.get('M_DEFS', '').split() else ['-DM_RELAXED'] FAST_WPW = [] if EXACT or os.environ.get('M_WPW') or 'WPW=' in os.environ.get('M_DEFS', '') else ['-DWPW=2'] +# preset L's expansion stage (dec0l) two windows per wave as well; its input stage (enc0l) is slower that way +L_WPW = [] if EXACT or 'WPW=' in os.environ.get('M_DEFS', '') + os.environ.get('M_DEFS_DEC0L', '') else ['-DWPW=2'] SHADERS = {'enc1': 'm_enc1', 'enc2': 'm_enc2', 'enc3_tube': 'm_tube', 'dec2': 'm_dec2', 'dec1': 'm_dec1'} @@ -159,7 +161,7 @@ def compile_l(a, one, R): for name, out in SHADERS.items(): glsl('layer_m.comp', out + '.spv', defines(name) + ['-DTEXL=1'] + FAST + FAST_WPW) for planes, suffix in ((0, ''), (1, '_planes')): glsl('layer_m.comp', f'enc0_l{suffix}.spv', defines('enc0l') + ldr + [f'-DPLANES={planes}'] + FAST) - glsl('layer_m.comp', f'dec0_l{suffix}.spv', defines('dec0l') + [f'-DPLANES={planes}'] + FAST) + glsl('layer_m.comp', f'dec0_l{suffix}.spv', defines('dec0l') + [f'-DPLANES={planes}'] + FAST + L_WPW) # the reconstruction one render texel wider per tier than M's (4x4 taps): 11 (reconstruction/render >= 3), # 16 (the whole post_3_2 range), 12 (the compact tier: render/reconstruction <= 0.45, jitter within half a texel) for out, tile in (('post_l', 11), ('post_l_q', 16), ('post_l_c', 12)): diff --git a/engine/layer_m.comp b/engine/layer_m.comp index 36bd42c..84e673c 100644 --- a/engine/layer_m.comp +++ b/engine/layer_m.comp @@ -64,8 +64,8 @@ #ifndef WPW #define WPW 1 #endif -#if WPW > 1 && KIND >= 3 -#error layer_m.comp: L's input and expansion stages stage their pixels per thread (WPW = 1) +#if WPW > 2 && KIND >= 3 +#error layer_m.comp: L's input and expansion stages handle their 64 pixels one per invocation (WPW <= 2) #endif layout(local_size_x = 128 / WPW) in; // o: 0 g1, 1 bo, 2 g2, 3 b2 (f16); per head h at 4 + 4h: q, v (FP8), bias (f16), wo (FP8); @@ -496,8 +496,7 @@ void main() #endif #if KIND == 3 // ------------------------------------------------------------ L's input: features, embedding -> x0 - FORW { - WINV + { const float e1 = ep.expo == 0.0 ? 1.0 : ep.expo, e2 = ep.expoPrev == 0.0 ? 1.0 : ep.expoPrev; const uint li = gl_LocalInvocationID.x; if (li < 64u) @@ -510,7 +509,10 @@ void main() const uint at = (((py >> 2) * 2u + (px >> 2)) * 16u + (py & 3u) * 4u + (px & 3u)) * FS; for (uint k = 0u; k < 32u; ++k) FEW[at + k] = fe[k]; } - SYNC + } + SYNC + FORW { + WINV // the embedding has the shape of an MLP chunk: 32 features -> 32 hidden (ReLU, FP8) -> C B8 f0, f1; coopMatLoad(f0, FEW, mt * 16u * FS, FS, CM); @@ -532,8 +534,8 @@ void main() for (int i = 0; i < 8; ++i) X0S[xb + 16u * (N) + uint(i)] = float16_t(x_[i]); } EACH_T(EMB_N) #undef EMB_N - SYNC } + SYNC #endif // ------------------------------------------------------------ norm 1: sum of squares, residual and h1 FORW { @@ -768,7 +770,9 @@ void main() const uint hb = (mt * 16u + row) * HS + ch8; for (uint nt = 0u; nt < NT; ++nt) for (int i = 0; i < 8; ++i) HO[hb + 16u * nt + uint(i)] = float16_t(res[nt][i]); - SYNC + } + SYNC + { const uint li = gl_LocalInvocationID.x; const uint px = li & 7u, py = li >> 3; const ivec2 at = ivec2(X0 + int(px), Y0 + int(py)); diff --git a/engine/test_m.py b/engine/test_m.py index 257e4c0..ffa273e 100644 --- a/engine/test_m.py +++ b/engine/test_m.py @@ -90,7 +90,7 @@ def defines(name): # compile_m.py's arithmetic (ZLUDA's lowering); M_TEXL=0 gives the native accuracy kernels' for the network layers. # M_DEFS: extra defines for every layer_m.comp stage; M_DEFS_ENC1, M_DEFS_ENC3_TUBE, ...: for one layer texl = ['-DTEXL=1'] if kind >= 3 or os.environ.get('M_TEXL', '1') != '0' else [] - # M_WPW=2: two 4x4-token windows per wave in the network layers (layer_m.comp WPW; not L's input / expansion) + # M_WPW=2: two 4x4-token windows per wave in the network layers (layer_m.comp WPW; L's input / expansion: M_DEFS_*) wpw = [f"-DWPW={os.environ['M_WPW']}"] if kind < 3 and os.environ.get('M_WPW') else [] return [f'-DC={C}', f'-DNH={heads}', f'-DKIND={kind}', f'-DXC={NPM or CIN or 16}'] + texl + wpw + os.environ.get('M_DEFS', '').split() + os.environ.get('M_DEFS_' + name.upper(), '').split()