Skip to content

Commit c342c3b

Browse files
jeffbolznvjeffbolznv
authored
vulkan: fix non-contig rope (ggml-org#19299)
1 parent af252d0 commit c342c3b

8 files changed

Lines changed: 100 additions & 95 deletions

File tree

‎ggml/src/ggml-vulkan/ggml-vulkan.cpp‎

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1263,25 +1263,30 @@ struct vk_op_diag_mask_push_constants {
12631263

12641264
struct vk_op_rope_push_constants {
12651265
uint32_t rope_mode;
1266-
uint32_t ncols;
12671266
uint32_t nrows;
12681267
uint32_t n_dims;
12691268
float freq_scale;
1270-
uint32_t p_delta_rows;
12711269
float freq_base;
12721270
float ext_factor;
12731271
float attn_factor;
12741272
float corr_dims[2];
12751273
float theta_scale;
12761274
uint32_t has_ff;
1277-
uint32_t ne02;
1278-
uint32_t s1;
1279-
uint32_t s2;
12801275
int32_t sections[4];
12811276
uint32_t is_imrope;
12821277
uint32_t is_back;
12831278
uint32_t set_rows_stride;
1279+
uint32_t ne00;
1280+
uint32_t ne01;
1281+
uint32_t ne02;
1282+
uint32_t nb01;
1283+
uint32_t nb02;
1284+
uint32_t nb03;
1285+
uint32_t nb11;
1286+
uint32_t nb12;
1287+
uint32_t nb13;
12841288
};
1289+
static_assert(sizeof(vk_op_rope_push_constants) <= 128, "sizeof(vk_op_rope_push_constants) must be <= 128");
12851290

12861291
// For fused rms_norm+mul+rope(+view+set_rows)
12871292
struct vk_op_rms_norm_mul_rope_push_constants {
@@ -10405,12 +10410,22 @@ static vk_op_rope_push_constants ggml_vk_make_rope_constants(const ggml_tensor *
1040510410

1040610411
uint32_t nb01 = src0->nb[1] / ggml_type_size(src0->type);
1040710412
uint32_t nb02 = src0->nb[2] / ggml_type_size(src0->type);
10413+
uint32_t nb03 = src0->nb[3] / ggml_type_size(src0->type);
10414+
10415+
uint32_t nb11 = dst->nb[1] / ggml_type_size(dst->type);
10416+
uint32_t nb12 = dst->nb[2] / ggml_type_size(dst->type);
10417+
uint32_t nb13 = dst->nb[3] / ggml_type_size(dst->type);
1040810418

1040910419
vk_op_rope_push_constants rope {
10410-
(uint32_t)mode, (uint32_t)src0->ne[0], (uint32_t)ggml_nrows(src0), (uint32_t)n_dims, freq_scale, (uint32_t)src0->ne[1],
10411-
freq_base, ext_factor, attn_factor, {corr_dims[0], corr_dims[1]}, theta_scale,
10412-
has_ff, (uint32_t)src0->ne[2], nb01, nb02,
10420+
(uint32_t)mode, (uint32_t)ggml_nrows(src0), (uint32_t)n_dims, freq_scale,
10421+
freq_base, ext_factor, attn_factor, {corr_dims[0], corr_dims[1]}, theta_scale, has_ff,
1041310422
{ sections[0], sections[1], sections[2], sections[3] }, is_imrope, backprop, set_rows_stride,
10423+
10424+
(uint32_t)src0->ne[0],
10425+
(uint32_t)src0->ne[1],
10426+
(uint32_t)src0->ne[2],
10427+
nb01, nb02, nb03,
10428+
nb11, nb12, nb13,
1041410429
};
1041510430

1041610431
return rope;
@@ -14798,6 +14813,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
1479814813
case GGML_OP_REPEAT_BACK:
1479914814
return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32;
1480014815
case GGML_OP_ROPE:
14816+
return ggml_is_contiguous_rows(op) && ggml_is_contiguous_rows(op->src[0]);
1480114817
case GGML_OP_ROPE_BACK:
1480214818
case GGML_OP_NONE:
1480314819
case GGML_OP_RESHAPE:

‎ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -112,12 +112,11 @@ void rms_norm(uint num_iters) {
112112
#if RMS_NORM_ROPE_FUSION
113113
barrier();
114114
rope_params rp = p.rope;
115-
uint rope_row = (samp*nchannels + channel)*nrows + row;
116115
for (uint t = 2*tid; t < ncols; t += 2*BLOCK_SIZE) {
117116
if (rp.rope_mode == GGML_ROPE_TYPE_NEOX) {
118-
rope_neox(t, rope_row, rp);
117+
rope_neox(t, row, channel, samp, rp);
119118
} else if (rp.rope_mode == GGML_ROPE_TYPE_NORMAL) {
120-
rope_norm(t, rope_row, rp);
119+
rope_norm(t, row, channel, samp, rp);
121120
}
122121
}
123122
#endif

‎ggml/src/ggml-vulkan/vulkan-shaders/rope_funcs.glsl‎

Lines changed: 36 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,12 @@ float rope_yarn_ramp(const float low, const float high, const uint i0) {
44
return 1.0f - min(1.0f, max(0.0f, y));
55
}
66

7-
uint rope_a_coord(const uint i0, const uint i01, const uint i02, rope_params p) {
7+
uint rope_a_coord(const uint i0, const uint i01, const uint i02, const uint i03, rope_params p) {
88
#if RMS_NORM_ROPE_FUSION
99
// Per-row offset in shared memory
1010
const uint ix = i0;
1111
#else
12-
const uint ix = i02*p.nb02 + i01*p.nb01 + i0;
12+
const uint ix = i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i0;
1313
#endif
1414
return ix;
1515
}
@@ -34,26 +34,19 @@ void rope_yarn(const float theta_extrap, const uint i0, out float cos_theta, out
3434
sin_theta = sin(theta) * mscale;
3535
}
3636

37-
void rope_norm(const uint i0, const uint i1, rope_params p) {
38-
uint ne0 = p.ncols;
39-
uint ne1 = p.p_delta_rows;
40-
41-
if (i0 >= ne0) {
37+
void rope_norm(const uint i0, const uint i1, const uint i2, const uint i3, rope_params p) {
38+
if (i0 >= p.ne00) {
4239
return;
4340
}
4441

45-
// i1 is actually i2*nb2+i1, but the rows are contiguous
46-
const uint i01 = i1 % ne1;
47-
const uint i02 = i1 / ne1;
48-
49-
uint idst = i1*ne0 + i0;
50-
const uint ix = rope_a_coord(i0, i01, i02, p);
42+
uint idst = i0 + i1 * p.nb11 + i2 * p.nb12 + i3 * p.nb13;
43+
const uint ix = rope_a_coord(i0, i1, i2, i3, p);
5144

5245
// Fusion optimization: ROPE + VIEW + SET_ROWS.
5346
// The rope output is viewed as a 1D tensor and offset based on a row index in rope_data_i.
5447
if (p.set_rows_stride != 0) {
55-
idst = i01*ne0 + i0;
56-
idst += rope_data_i[i02].x * p.set_rows_stride;
48+
idst = i1*p.nb11 + i0;
49+
idst += rope_data_i[i2].x * p.set_rows_stride;
5750
}
5851

5952
if (i0 >= p.n_dims) {
@@ -63,7 +56,7 @@ void rope_norm(const uint i0, const uint i1, rope_params p) {
6356
return;
6457
}
6558

66-
const float theta_base = rope_data_pos[i02] * pow(p.theta_scale, i0/2.0f);
59+
const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, i0/2.0f);
6760

6861
const float freq_factor = p.has_ff != 0 ? rope_data_ff[i0/2] : 1.0f;
6962

@@ -77,25 +70,19 @@ void rope_norm(const uint i0, const uint i1, rope_params p) {
7770
rope_data_d[idst + 1] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta);
7871
}
7972

80-
void rope_neox(const uint i0, const uint i1, rope_params p) {
81-
uint ne0 = p.ncols;
82-
uint ne1 = p.p_delta_rows;
83-
84-
if (i0 >= ne0) {
73+
void rope_neox(const uint i0, const uint i1, const uint i2, const uint i3, rope_params p) {
74+
if (i0 >= p.ne00) {
8575
return;
8676
}
8777

88-
const uint i01 = i1 % ne1;
89-
const uint i02 = i1 / ne1;
90-
91-
uint idst = i1*ne0 + i0/2;
92-
const uint ix = rope_a_coord(i0/2, i01, i02, p);
78+
uint idst = i0/2 + i1 * p.nb11 + i2 * p.nb12 + i3 * p.nb13;
79+
const uint ix = rope_a_coord(i0/2, i1, i2, i3, p);
9380

9481
// Fusion optimization: ROPE + VIEW + SET_ROWS.
9582
// The rope output is viewed as a 1D tensor and offset based on a row index in rope_data_i.
9683
if (p.set_rows_stride != 0) {
97-
idst = i01*ne0 + i0/2;
98-
idst += rope_data_i[i02].x * p.set_rows_stride;
84+
idst = i1*p.nb11 + i0/2;
85+
idst += rope_data_i[i2].x * p.set_rows_stride;
9986
}
10087

10188
if (i0 >= p.n_dims) {
@@ -105,7 +92,7 @@ void rope_neox(const uint i0, const uint i1, rope_params p) {
10592
return;
10693
}
10794

108-
const float theta_base = rope_data_pos[i02] * pow(p.theta_scale, i0/2.0f);
95+
const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, i0/2.0f);
10996

11097
const float freq_factor = p.has_ff != 0 ? rope_data_ff[i0/2] : 1.0f;
11198

@@ -120,26 +107,19 @@ void rope_neox(const uint i0, const uint i1, rope_params p) {
120107
}
121108

122109

123-
void rope_multi(const uint i0, const uint i1, rope_params p) {
124-
uint ne0 = p.ncols;
125-
uint ne1 = p.p_delta_rows;
126-
uint ne2 = p.ne02;
127-
128-
if (i0 >= ne0) {
110+
void rope_multi(const uint i0, const uint i1, const uint i2, const uint i3, rope_params p) {
111+
if (i0 >= p.ne00) {
129112
return;
130113
}
131114

132-
const uint i01 = i1 % ne1;
133-
const uint i02 = i1 / ne1;
134-
135-
uint idst = i1*ne0 + i0/2;
136-
const uint ix = rope_a_coord(i0/2, i01, i02, p);
115+
uint idst = i0/2 + i1 * p.nb11 + i2 * p.nb12 + i3 * p.nb13;
116+
const uint ix = rope_a_coord(i0/2, i1, i2, i3, p);
137117

138118
// Fusion optimization: ROPE + VIEW + SET_ROWS.
139119
// The rope output is viewed as a 1D tensor and offset based on a row index in rope_data_i.
140120
if (p.set_rows_stride != 0) {
141-
idst = i01*ne0 + i0/2;
142-
idst += rope_data_i[i02].x * p.set_rows_stride;
121+
idst = i1*p.nb11 + i0/2;
122+
idst += rope_data_i[i2].x * p.set_rows_stride;
143123
}
144124

145125
if (i0 >= p.n_dims) {
@@ -156,26 +136,26 @@ void rope_multi(const uint i0, const uint i1, rope_params p) {
156136
float theta_base = 0.0;
157137
if (p.is_imrope != 0) {
158138
if (sector % 3 == 1 && sector < 3 * p.sections[1]) {
159-
theta_base = rope_data_pos[i02 + ne2 * 1]*pow(p.theta_scale, i0/2.0f);
139+
theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, i0/2.0f);
160140
} else if (sector % 3 == 2 && sector < 3 * p.sections[2]) {
161-
theta_base = rope_data_pos[i02 + ne2 * 2]*pow(p.theta_scale, i0/2.0f);
141+
theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, i0/2.0f);
162142
} else if (sector % 3 == 0 && sector < 3 * p.sections[0]) {
163-
theta_base = rope_data_pos[i02]*pow(p.theta_scale, i0/2.0f);
143+
theta_base = rope_data_pos[i2]*pow(p.theta_scale, i0/2.0f);
164144
} else {
165-
theta_base = rope_data_pos[i02 + ne2 * 3]*pow(p.theta_scale, i0/2.0f);
145+
theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, i0/2.0f);
166146
}
167147
} else {
168148
if (sector < p.sections[0]) {
169-
theta_base = rope_data_pos[i02]*pow(p.theta_scale, i0/2.0f);
149+
theta_base = rope_data_pos[i2]*pow(p.theta_scale, i0/2.0f);
170150
}
171151
else if (sector >= p.sections[0] && sector < sec_w) {
172-
theta_base = rope_data_pos[i02 + ne2 * 1]*pow(p.theta_scale, i0/2.0f);
152+
theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, i0/2.0f);
173153
}
174154
else if (sector >= sec_w && sector < sec_w + p.sections[2]) {
175-
theta_base = rope_data_pos[i02 + ne2 * 2]*pow(p.theta_scale, i0/2.0f);
155+
theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, i0/2.0f);
176156
}
177157
else if (sector >= sec_w + p.sections[2]) {
178-
theta_base = rope_data_pos[i02 + ne2 * 3]*pow(p.theta_scale, i0/2.0f);
158+
theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, i0/2.0f);
179159
}
180160
}
181161

@@ -191,20 +171,13 @@ void rope_multi(const uint i0, const uint i1, rope_params p) {
191171
rope_data_d[idst + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta);
192172
}
193173

194-
void rope_vision(const uint i0, const uint i1, rope_params p) {
195-
uint ne0 = p.ncols;
196-
uint ne1 = p.p_delta_rows;
197-
uint ne2 = p.ne02;
198-
199-
if (i0 >= ne0) {
174+
void rope_vision(const uint i0, const uint i1, const uint i2, const uint i3, rope_params p) {
175+
if (i0 >= p.ne00) {
200176
return;
201177
}
202178

203-
const uint i01 = i1 % ne1;
204-
const uint i02 = i1 / ne1;
205-
206-
const uint idst = i1*ne0 + i0/2;
207-
const uint ix = rope_a_coord(i0/2, i01, i02, p);
179+
const uint idst = i0/2 + i1 * p.nb11 + i2 * p.nb12 + i3 * p.nb13;
180+
const uint ix = rope_a_coord(i0/2, i1, i2, i3, p);
208181

209182
const int sect_dims = p.sections[0] + p.sections[1];
210183
const int sec_w = p.sections[1] + p.sections[0];
@@ -213,11 +186,11 @@ void rope_vision(const uint i0, const uint i1, rope_params p) {
213186
float theta_base = 0.0;
214187
if (sector < p.sections[0]) {
215188
const uint p0 = sector;
216-
theta_base = rope_data_pos[i02]*pow(p.theta_scale, p0);
189+
theta_base = rope_data_pos[i2]*pow(p.theta_scale, p0);
217190
}
218191
else if (sector >= p.sections[0] && sector < sec_w) {
219192
const uint p0 = sector - p.sections[0];
220-
theta_base = rope_data_pos[i02 + ne2]*pow(p.theta_scale, p0);
193+
theta_base = rope_data_pos[i2 + p.ne02]*pow(p.theta_scale, p0);
221194
}
222195

223196
const float freq_factor = p.has_ff != 0 ? rope_data_ff[i0/2] : 1.0f;

‎ggml/src/ggml-vulkan/vulkan-shaders/rope_multi.comp‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,13 @@
55

66
void main() {
77
const uint i0 = 2*gl_GlobalInvocationID.y;
8-
// i1 is actually i2*nb2+i1, but the rows are contiguous
9-
const uint i1 = gl_GlobalInvocationID.x + 32768 * gl_GlobalInvocationID.z;
10-
if (i1 >= pc.nrows) {
8+
const uint row = gl_GlobalInvocationID.x + 32768 * gl_GlobalInvocationID.z;
9+
if (row >= pc.nrows) {
1110
return;
1211
}
13-
rope_multi(i0, i1, pc);
12+
const uint i3 = row / (pc.ne01*pc.ne02);
13+
const uint i2 = (row - i3 * pc.ne01*pc.ne02) / pc.ne01;
14+
const uint i1 = (row - i3 * pc.ne01*pc.ne02 - i2 * pc.ne01);
15+
16+
rope_multi(i0, i1, i2, i3, pc);
1417
}

‎ggml/src/ggml-vulkan/vulkan-shaders/rope_neox.comp‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,13 @@
55

66
void main() {
77
const uint i0 = 2*gl_GlobalInvocationID.y;
8-
// i1 is actually i2*nb2+i1, but the rows are contiguous
9-
const uint i1 = gl_GlobalInvocationID.x + 32768 * gl_GlobalInvocationID.z;
10-
if (i1 >= pc.nrows) {
8+
const uint row = gl_GlobalInvocationID.x + 32768 * gl_GlobalInvocationID.z;
9+
if (row >= pc.nrows) {
1110
return;
1211
}
13-
rope_neox(i0, i1, pc);
12+
const uint i3 = row / (pc.ne01*pc.ne02);
13+
const uint i2 = (row - i3 * pc.ne01*pc.ne02) / pc.ne01;
14+
const uint i1 = (row - i3 * pc.ne01*pc.ne02 - i2 * pc.ne01);
15+
16+
rope_neox(i0, i1, i2, i3, pc);
1417
}

‎ggml/src/ggml-vulkan/vulkan-shaders/rope_norm.comp‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,13 @@
55

66
void main() {
77
const uint i0 = 2*gl_GlobalInvocationID.y;
8-
// i1 is actually i2*nb2+i1, but the rows are contiguous
9-
const uint i1 = gl_GlobalInvocationID.x + 32768 * gl_GlobalInvocationID.z;
10-
if (i1 >= pc.nrows) {
8+
const uint row = gl_GlobalInvocationID.x + 32768 * gl_GlobalInvocationID.z;
9+
if (row >= pc.nrows) {
1110
return;
1211
}
13-
rope_norm(i0, i1, pc);
12+
const uint i3 = row / (pc.ne01*pc.ne02);
13+
const uint i2 = (row - i3 * pc.ne01*pc.ne02) / pc.ne01;
14+
const uint i1 = (row - i3 * pc.ne01*pc.ne02 - i2 * pc.ne01);
15+
16+
rope_norm(i0, i1, i2, i3, pc);
1417
}

‎ggml/src/ggml-vulkan/vulkan-shaders/rope_params.glsl‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,24 +5,29 @@
55

66
struct rope_params {
77
uint rope_mode;
8-
uint ncols;
98
uint nrows;
109
uint n_dims;
1110
float freq_scale;
12-
uint p_delta_rows;
1311
float freq_base;
1412
float ext_factor;
1513
float attn_factor;
1614
float corr_dims[2];
1715
float theta_scale;
1816
uint has_ff;
19-
uint ne02;
20-
uint nb01;
21-
uint nb02;
2217
int sections[4];
2318
uint is_imrope;
2419
uint is_back;
2520
uint set_rows_stride;
21+
22+
uint ne00;
23+
uint ne01;
24+
uint ne02;
25+
uint nb01;
26+
uint nb02;
27+
uint nb03;
28+
uint nb11;
29+
uint nb12;
30+
uint nb13;
2631
};
2732

2833
#endif // !defined(GGML_ROPE_PARAMS)

‎ggml/src/ggml-vulkan/vulkan-shaders/rope_vision.comp‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,13 @@
55

66
void main() {
77
const uint i0 = 2*gl_GlobalInvocationID.y;
8-
// i1 is actually i2*nb2+i1, but the rows are contiguous
9-
const uint i1 = gl_GlobalInvocationID.x + 32768 * gl_GlobalInvocationID.z;
10-
if (i1 >= pc.nrows) {
8+
const uint row = gl_GlobalInvocationID.x + 32768 * gl_GlobalInvocationID.z;
9+
if (row >= pc.nrows) {
1110
return;
1211
}
13-
rope_vision(i0, i1, pc);
12+
const uint i3 = row / (pc.ne01*pc.ne02);
13+
const uint i2 = (row - i3 * pc.ne01*pc.ne02) / pc.ne01;
14+
const uint i1 = (row - i3 * pc.ne01*pc.ne02 - i2 * pc.ne01);
15+
16+
rope_vision(i0, i1, i2, i3, pc);
1417
}

0 commit comments

Comments
 (0)