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
[Bug] Fix Cutlass Scaled MM Compilation Error (#24887)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Sep 15, 2025 at 17:21 UTC
e757a629e7336c4e436dac62daee20fd81776aeb
3 files changed
+53
-41
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh
+23
-19
@@ -146,6 +146,7 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
146
147
using ElementAB = typename Gemm::ElementAB;
148
using ElementD = typename Gemm::ElementD;
149
+ using ElementBlockScale = typename Gemm::ElementBlockScale;
150
151
int32_t m = a.size(0), n = b.size(1), k = a.size(1);
152
@@ -166,26 +167,29 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
167
ScaleConfig::tile_atom_to_shape_SFB(make_shape(n, m, k, 1)) :
168
ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
169
169
- auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
170
- auto b_ptr = static_cast<ElementAB*>(b.data_ptr());
171
- auto a_scales_ptr = static_cast<float*>(a_scales.data_ptr());
172
- auto b_scales_ptr = static_cast<float*>(b_scales.data_ptr());
170
+ auto a_ptr = static_cast<ElementAB const*>(a.data_ptr());
171
+ auto b_ptr = static_cast<ElementAB const*>(b.data_ptr());
172
+ auto a_scales_ptr = static_cast<ElementBlockScale const*>(a_scales.data_ptr());
173
+ auto b_scales_ptr = static_cast<ElementBlockScale const*>(b_scales.data_ptr());
174
174
- auto mainloop_args = [&](){
175
- // layout_SFA and layout_SFB cannot be swapped since they are deduced.
176
- if (swap_ab) {
177
- return typename GemmKernel::MainloopArguments{
178
- b_ptr, b_stride, a_ptr, a_stride,
179
- b_scales_ptr, layout_SFA, a_scales_ptr, layout_SFB
180
- };
181
- }
182
- else {
183
- return typename GemmKernel::MainloopArguments{
184
- a_ptr, a_stride, b_ptr, b_stride,
185
- a_scales_ptr, layout_SFA, b_scales_ptr, layout_SFB
186
- };
187
- }
188
- }();
175
+ typename GemmKernel::MainloopArguments mainloop_args{};
176
+ mainloop_args.layout_SFA = layout_SFA;
177
+ mainloop_args.layout_SFB = layout_SFB;
178
+ if (swap_ab) {
179
+ mainloop_args.ptr_A = b_ptr;
180
+ mainloop_args.dA = b_stride;
181
+ mainloop_args.ptr_B = a_ptr;
182
+ mainloop_args.dB = a_stride;
183
+ mainloop_args.ptr_SFA = b_scales_ptr;
184
+ mainloop_args.ptr_SFB = a_scales_ptr;
185
+ } else {
186
+ mainloop_args.ptr_A = a_ptr;
187
+ mainloop_args.dA = a_stride;
188
+ mainloop_args.ptr_B = b_ptr;
189
+ mainloop_args.dB = b_stride;
190
+ mainloop_args.ptr_SFA = a_scales_ptr;
191
+ mainloop_args.ptr_SFB = b_scales_ptr;
192
+ }
193
auto prob_shape = swap_ab ? cute::make_shape(n, m, k, 1) : cute::make_shape(m, n, k, 1);
194
195
auto c_ptr = static_cast<ElementD*>(out.data_ptr());
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh
+15
-11
@@ -125,6 +125,7 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
125
126
using ElementAB = typename Gemm::ElementAB;
127
using ElementD = typename Gemm::ElementD;
128
+ using ElementBlockScale = typename Gemm::ElementBlockScale;
129
130
int32_t m = a.size(0), n = b.size(1), k = a.size(1);
131
@@ -143,17 +144,20 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
144
LayoutSFB layout_SFB =
145
ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
146
146
- auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
147
- auto b_ptr = static_cast<ElementAB*>(b.data_ptr());
148
- auto a_scales_ptr = static_cast<float*>(a_scales.data_ptr());
149
- auto b_scales_ptr = static_cast<float*>(b_scales.data_ptr());
150
-
151
- auto mainloop_args = [&](){
152
- return typename GemmKernel::MainloopArguments{
153
- a_ptr, a_stride, b_ptr, b_stride,
154
- a_scales_ptr, layout_SFA, b_scales_ptr, layout_SFB
155
- };
156
- }();
147
+ auto a_ptr = static_cast<ElementAB const*>(a.data_ptr());
148
+ auto b_ptr = static_cast<ElementAB const*>(b.data_ptr());
149
+ auto a_scales_ptr = static_cast<ElementBlockScale const*>(a_scales.data_ptr());
150
+ auto b_scales_ptr = static_cast<ElementBlockScale const*>(b_scales.data_ptr());
151
+
152
+ typename GemmKernel::MainloopArguments mainloop_args{};
153
+ mainloop_args.ptr_A = a_ptr;
154
+ mainloop_args.dA = a_stride;
155
+ mainloop_args.ptr_B = b_ptr;
156
+ mainloop_args.dB = b_stride;
157
+ mainloop_args.ptr_SFA = a_scales_ptr;
158
+ mainloop_args.layout_SFA = layout_SFA;
159
+ mainloop_args.ptr_SFB = b_scales_ptr;
160
+ mainloop_args.layout_SFB = layout_SFB;
161
auto prob_shape = cute::make_shape(m, n, k, 1);
162
163
auto c_ptr = static_cast<ElementD*>(out.data_ptr());
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh
+15
-11
@@ -115,6 +115,7 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
115
116
using ElementAB = typename Gemm::ElementAB;
117
using ElementD = typename Gemm::ElementD;
118
+ using ElementBlockScale = typename Gemm::ElementBlockScale;
119
120
int32_t m = a.size(0), n = b.size(1), k = a.size(1);
121
@@ -135,17 +136,20 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
136
LayoutSFB layout_SFB =
137
ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
138
138
- auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
139
- auto b_ptr = static_cast<ElementAB*>(b.data_ptr());
140
- auto a_scales_ptr = static_cast<float*>(a_scales.data_ptr());
141
- auto b_scales_ptr = static_cast<float*>(b_scales.data_ptr());
142
-
143
- auto mainloop_args = [&](){
144
- return typename GemmKernel::MainloopArguments{
145
- a_ptr, a_stride, b_ptr, b_stride,
146
- a_scales_ptr, layout_SFA, b_scales_ptr, layout_SFB
147
- };
148
- }();
139
+ auto a_ptr = static_cast<ElementAB const*>(a.data_ptr());
140
+ auto b_ptr = static_cast<ElementAB const*>(b.data_ptr());
141
+ auto a_scales_ptr = static_cast<ElementBlockScale const*>(a_scales.data_ptr());
142
+ auto b_scales_ptr = static_cast<ElementBlockScale const*>(b_scales.data_ptr());
143
+
144
+ typename GemmKernel::MainloopArguments mainloop_args{};
145
+ mainloop_args.ptr_A = a_ptr;
146
+ mainloop_args.dA = a_stride;
147
+ mainloop_args.ptr_B = b_ptr;
148
+ mainloop_args.dB = b_stride;
149
+ mainloop_args.ptr_SFA = a_scales_ptr;
150
+ mainloop_args.layout_SFA = layout_SFA;
151
+ mainloop_args.ptr_SFB = b_scales_ptr;
152
+ mainloop_args.layout_SFB = layout_SFB;
153
auto prob_shape = cute::make_shape(m, n, k, 1);
154
155
auto c_ptr = static_cast<ElementD*>(out.data_ptr());