xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-07-10 10:37:16 +08:00
Code Issues Packages Projects Releases Wiki Activity

Make sure that vectorize_with_alignment produced vectorized global loads (#23182)

Browse Source
This commit is contained in:
Elvir Crnčević 2025-08-21 22:06:54 +02:00 committed by GitHub
parent 1d353b6352
commit 044931f97b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 10 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
csrc/quantization/vectorization_utils.cuh
View File

@ -41,8 +41,10 @@ __device__ inline void vectorize_with_alignment(
for (int i = tid; i < num_vec; i += stride) { for (int i = tid; i < num_vec; i += stride) {
vout_t tmp; vout_t tmp;
vec_op(tmp, v_in[i]); // Make a local copy of the entire pack
v_out[i] = tmp; vin_t src = v_in[i]; // <- encourages a single vector ld
vec_op(tmp, src);
v_out[i] = tmp; // <- encourages a single vector st
} }
return; return;
} }
@ -71,8 +73,10 @@ __device__ inline void vectorize_with_alignment(
// 2. vectorize the main part // 2. vectorize the main part
for (int i = tid; i < num_vec; i += stride) { for (int i = tid; i < num_vec; i += stride) {
vout_t tmp; vout_t tmp;
vec_op(tmp, v_in[i]); // Make a local copy of the entire pack
v_out[i] = tmp; vin_t src = v_in[i]; // <- encourages a single vector ld
vec_op(tmp, src);
v_out[i] = tmp; // <- encourages a single vector st
} }
// 3. handle the tail // 3. handle the tail
@ -125,7 +129,8 @@ __device__ inline void vectorize_read_with_alignment(const InT* in, int len,
auto* v_in = reinterpret_cast<const vin_t*>(in); auto* v_in = reinterpret_cast<const vin_t*>(in);
for (int i = tid; i < num_vec; i += stride) { for (int i = tid; i < num_vec; i += stride) {
vec_op(v_in[i]); vin_t tmp = v_in[i];
vec_op(tmp);
} }
return; return;
} }