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 5, 08:30 UTC
[BugFix] Fix OOB read in CUTLASS grouped GEMM with epilogue (#38571)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Lucas Wilkinson committed
Apr 9, 2026 at 23:52 UTC
f83de7196f0d589a1a5d71dd20e2d5867c361dd4
2 files changed
+34
-18
csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp
+17
-9
@@ -389,20 +389,28 @@ struct Sm90ColOrScalarBroadcastArray {
389
390
CUTLASS_DEVICE void
391
begin() {
392
- cute::Tensor pred = make_tensor<bool>(shape(tCgCol));
393
- CUTLASS_PRAGMA_UNROLL
394
- for (int i = 0; i < size(pred); ++i) {
395
- pred(i) = get<0>(tCcCol(i)) < m;
396
- }
397
-
392
if (!params.col_broadcast) {
393
fill(tCrCol, *(params.ptr_col_array[group]));
394
return;
395
}
396
403
- // Filter so we don't issue redundant copies over stride-0 modes
404
- // (only works if 0-strides are in same location, which is by construction)
405
- copy_if(pred, filter(tCgCol), filter(tCrCol));
397
+ // tCgCol has layout (CPY,CPY_M,CPY_N,EPI_M,EPI_N) where CPY_N and
398
+ // EPI_N are stride-0 for the column broadcast. Slice those modes at
399
+ // index 0 to avoid redundant copies AND ensure pred/data consistency
400
+ static_assert(decltype(stride<2>(tCgCol))::value == 0, "Expected stride-0 CPY_N for col broadcast");
401
+ static_assert(decltype(stride<4>(tCgCol))::value == 0, "Expected stride-0 EPI_N for col broadcast");
402
+
403
+ auto tCgCol_s = tCgCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M)
404
+ auto tCrCol_s = tCrCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M)
405
+ auto tCcCol_s = tCcCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M)
406
+
407
+ cute::Tensor pred = make_tensor<bool>(shape(tCgCol_s));
408
+ CUTLASS_PRAGMA_UNROLL
409
+ for (int i = 0; i < size(pred); ++i) {
410
+ pred(i) = get<0>(tCcCol_s(i)) < m;
411
+ }
412
+
413
+ copy_if(pred, tCgCol_s, tCrCol_s);
414
}
415
416
template <typename ElementAccumulator, int FragmentSize>
csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp
+17
-9
@@ -382,20 +382,28 @@ struct Sm90ColOrScalarBroadcast {
382
383
CUTLASS_DEVICE void
384
begin() {
385
- cute::Tensor pred = make_tensor<bool>(shape(tCgCol));
386
- CUTLASS_PRAGMA_UNROLL
387
- for (int i = 0; i < size(pred); ++i) {
388
- pred(i) = get<0>(tCcCol(i)) < m;
389
- }
390
-
385
if (!params.col_broadcast) {
386
fill(tCrCol, *(params.ptr_col));
387
return;
388
}
389
396
- // Filter so we don't issue redundant copies over stride-0 modes
397
- // (only works if 0-strides are in same location, which is by construction)
398
- copy_if(pred, filter(tCgCol), filter(tCrCol));
390
+ // tCgCol has layout (CPY,CPY_M,CPY_N,EPI_M,EPI_N) where CPY_N and
391
+ // EPI_N are stride-0 for the column broadcast. Slice those modes at
392
+ // index 0 to avoid redundant copies AND ensure pred/data consistency
393
+ static_assert(decltype(stride<2>(tCgCol))::value == 0, "Expected stride-0 CPY_N for col broadcast");
394
+ static_assert(decltype(stride<4>(tCgCol))::value == 0, "Expected stride-0 EPI_N for col broadcast");
395
+
396
+ auto tCgCol_s = tCgCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M)
397
+ auto tCrCol_s = tCrCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M)
398
+ auto tCcCol_s = tCcCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M)
399
+
400
+ cute::Tensor pred = make_tensor<bool>(shape(tCgCol_s));
401
+ CUTLASS_PRAGMA_UNROLL
402
+ for (int i = 0; i < size(pred); ++i) {
403
+ pred(i) = get<0>(tCcCol_s(i)) < m;
404
+ }
405
+
406
+ copy_if(pred, tCgCol_s, tCrCol_s);
407
}
408
409
template <typename ElementAccumulator, int FragmentSize>