char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Masashi Yoshimuraggml-webgpu : refactor several wgsl files and simplify flash_attn wgsl. (#26134)f401bb139

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
2.8 KiB134 linesraw
1#ifdef Q_F32
2#define Q_TYPE f32
3#else
4#define Q_TYPE f16
5#endif
6
7#ifdef K_F32
8#define K_TYPE f32
9#elif defined(K_Q4_0) || defined(K_Q8_0)
10#define K_TYPE u32
11#else
12#define K_TYPE f16
13#endif
14
15#ifdef V_F32
16#define V_TYPE f32
17#elif defined(V_Q4_0) || defined(V_Q8_0)
18#define V_TYPE u32
19#else
20#define V_TYPE f16
21#endif
22
23#ifdef DST_F32
24#define DST_TYPE f32
25#else
26#define DST_TYPE f16
27#endif
28
29#if defined(FLASH_ATTN_SCALAR_KV) || defined(K_Q4_0) || defined(K_Q8_0)
30#define K_STORAGE_TYPE K_TYPE
31#else
32#define K_STORAGE_TYPE vec4<K_TYPE>
33#endif
34
35#if defined(FLASH_ATTN_SCALAR_KV) || defined(V_Q4_0) || defined(V_Q8_0)
36#define V_STORAGE_TYPE V_TYPE
37#else
38#define V_STORAGE_TYPE vec4<V_TYPE>
39#endif
40
41// Just a very small float value.
42const FLOAT_MIN: f32 = -1.0e9;
43
44struct Params {
45    offset_q: u32,
46    offset_k: u32,
47    offset_v: u32,
48    offset_mask: u32,
49    offset_sinks: u32,
50    offset_dst: u32,
51
52    // shapes of Q/K/V
53    n_heads: u32,
54    seq_len_q: u32,
55    seq_len_kv: u32,
56
57    // strides (in elements)
58    stride_q1: u32,
59    stride_q2: u32,
60    stride_q3: u32,
61    stride_k1: u32,
62    stride_k2: u32,
63    stride_k3: u32,
64    stride_v1: u32,
65    stride_v2: u32,
66    stride_v3: u32,
67    stride_mask3: u32,
68
69    // repeat factors for K/V, e.g., MHA vs. MQA vs. GQA
70    q_per_kv: u32,
71
72    // softmax params
73    scale: f32,
74    max_bias: f32,
75    logit_softcap: f32,
76    n_head_log2: f32,
77    m0: f32,
78    m1: f32,
79
80#ifdef FLASH_ATTN_VEC_SPLIT
81#ifdef BLK
82    blk_base: u32,
83    blk_nblk0: u32,
84    blk_nblk1: u32,
85#endif
86
87    tmp_data_base: u32,
88    tmp_stats_base: u32,
89    nwg: u32,
90#endif
91};
92
93@group(0) @binding(0) var<storage, read_write> Q: array<Q_TYPE>;
94@group(0) @binding(1) var<storage, read_write> K: array<K_STORAGE_TYPE>;
95#ifdef KV_OVERLAP
96#define V K
97#define MASK_BINDING 2
98#else
99@group(0) @binding(2) var<storage, read_write> V: array<V_STORAGE_TYPE>;
100#define MASK_BINDING 3
101#endif // KV_OVERLAP
102
103#ifdef MASK
104@group(0) @binding(MASK_BINDING) var<storage, read_write> mask: array<f16>;
105#define SINKS_BINDING (MASK_BINDING + 1)
106#else
107#define SINKS_BINDING MASK_BINDING
108#endif
109
110#ifdef SINKS
111@group(0) @binding(SINKS_BINDING) var<storage, read_write> sinks: array<f32>;
112#define BLK_BINDING (SINKS_BINDING + 1)
113#else
114#define BLK_BINDING SINKS_BINDING
115#endif
116
117#ifdef FLASH_ATTN_VEC_SPLIT
118#ifdef BLK
119@group(0) @binding(BLK_BINDING) var<storage, read_write> blk: array<u32>;
120#define TMP_BINDING (BLK_BINDING + 1)
121#else
122#define TMP_BINDING BLK_BINDING
123#endif
124
125@group(0) @binding(TMP_BINDING) var<storage, read_write> tmp: array<f32>;
126#define DST_BINDING (TMP_BINDING + 1)
127#else
128#define DST_BINDING BLK_BINDING
129#endif // FLASH_ATTN_VEC_SPLIT
130
131@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<vec4<DST_TYPE>>;
132
133#define PARAMS_BINDING (DST_BINDING + 1)
134@group(0) @binding(PARAMS_BINDING) var<uniform> params: Params;