A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
Make sure that vectorize_with_alignment produced vectorized global loads (#23182)
Elvir Crnčević committed
Aug 21, 2025 at 22:06 UTC
044931f97b39975cce6dbef3df94586d83893758
1 file changed
+10
-5
csrc/quantization/vectorization_utils.cuh
+10
-5
@@ -41,8 +41,10 @@ __device__ inline void vectorize_with_alignment(
41
42
for (int i = tid; i < num_vec; i += stride) {
43
vout_t tmp;
44
- vec_op(tmp, v_in[i]);
45
- v_out[i] = tmp;
44
+ // Make a local copy of the entire pack
45
+ vin_t src = v_in[i]; // <- encourages a single vector ld
46
+ vec_op(tmp, src);
47
+ v_out[i] = tmp; // <- encourages a single vector st
48
}
49
return;
50
}
@@ -71,8 +73,10 @@ __device__ inline void vectorize_with_alignment(
73
// 2. vectorize the main part
74
for (int i = tid; i < num_vec; i += stride) {
75
vout_t tmp;
74
- vec_op(tmp, v_in[i]);
75
- v_out[i] = tmp;
76
+ // Make a local copy of the entire pack
77
+ vin_t src = v_in[i]; // <- encourages a single vector ld
78
+ vec_op(tmp, src);
79
+ v_out[i] = tmp; // <- encourages a single vector st
80
}
81
82
// 3. handle the tail
@@ -125,7 +129,8 @@ __device__ inline void vectorize_read_with_alignment(const InT* in, int len,
129
auto* v_in = reinterpret_cast<const vin_t*>(in);
130
131
for (int i = tid; i < num_vec; i += stride) {
128
- vec_op(v_in[i]);
132
+ vin_t tmp = v_in[i];
133
+ vec_op(tmp);
134
}
135
return;
136
}