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, 23:45 UTC
[Kernel] Improve machete memory bound perf (#21556)
Signed-off-by: czhu-cohere <conway.zhu@cohere.com>
czhu-cohere committed
Jul 25, 2025 at 06:53 UTC
136d750f5f421ca5be2e24b0a913e813d99bb831
1 file changed
+6
-2
csrc/quantization/machete/machete_prepacked_layout.cuh
+6
-2
@@ -187,8 +187,12 @@ struct PrepackedLayoutBTemplate {
187
CUTE_HOST_DEVICE static constexpr auto TVbNbKL_to_offset_copy(
188
Shape_NKL shape_mkl) {
189
auto layout = TVbNbKL_to_offset(shape_mkl);
190
- return make_layout(coalesce(get<0>(layout)), get<1>(layout),
191
- get<2>(layout));
190
+ // for 4-bit elements, having >= 64 values per column
191
+ // allows TMA to load full 32-byte sectors
192
+ auto inner_layout =
193
+ make_layout(make_shape(_256{}, size<0>(layout) / _256{}));
194
+
195
+ return make_layout(inner_layout, get<1>(layout), get<2>(layout));
196
}
197
198
// ((BlockN, BlockK), (BlocksN, BlocksK), L) -> (storage_idx)