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
[Compile] Fix Compile Warning for `w4a8_mm_entry.cu` (#23660)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Wentao Ye committed
Sep 2, 2025 at 23:45 UTC
c4ed78b14f7f63cfa65722ad21deffe964441fd2
1 file changed
+9
-3
csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu
+9
-3
@@ -11,6 +11,7 @@
11
#include "core/registration.h"
12
13
#include "cutlass/cutlass.h"
14
+#include <limits>
15
16
#include "cute/tensor.hpp"
17
#include "cutlass/gemm/collective/collective_builder.hpp"
@@ -169,6 +170,11 @@ struct W4A8GemmKernel {
170
int k = A.size(1);
171
int n = B.size(1);
172
173
+ // safely cast group_size to int
174
+ TORCH_CHECK(group_size > 0 && group_size <= std::numeric_limits<int>::max(),
175
+ "group_size out of supported range for int: ", group_size);
176
+ int const group_size_int = static_cast<int>(group_size);
177
+
178
// Allocate output
179
const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
180
auto device = A.device();
@@ -192,7 +198,7 @@ struct W4A8GemmKernel {
198
cute::tile_to_shape(LayoutAtomQuant{}, shape_B);
199
200
// strides
195
- int const scale_k = cutlass::ceil_div(k, group_size);
201
+ int const scale_k = cutlass::ceil_div(k, group_size_int);
202
StrideA stride_A =
203
cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1));
204
// Reverse stride here due to swap and transpose
@@ -211,8 +217,8 @@ struct W4A8GemmKernel {
217
using EpilogueArguments = typename GemmKernelShuffled::EpilogueArguments;
218
219
MainloopArguments mainloop_arguments{
214
- B_ptr, layout_B_reordered, A_ptr, stride_A,
215
- S_ptr, stride_S, group_size};
220
+ B_ptr, layout_B_reordered, A_ptr, stride_A,
221
+ S_ptr, stride_S, group_size_int};
222
223
EpilogueArguments epilogue_arguments{
224
ChTokScalesEpilogue::prepare_args(channel_scales, token_scales),