C3x SM100 FP8 blockwise - Pad activation scales to multiple of 4 (#57377)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Codex <noreply@openai.com>
Thien Tran committed
Sep 18, 2026 at 16:39 UTC
8da75d61904cd0b2c5c322746ea15251773ec51f
3 files changed
+39
-6
.buildkite/test_areas/kernels.yaml
+3
@@ -367,6 +367,7 @@ steps:
367
- csrc/quantization/fp4/
368
- csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu
369
- csrc/libtorch_stable/quantization/fp4/
370
+ - csrc/libtorch_stable/quantization/w8a8/cutlass/
371
- csrc/attention/mla/
372
- csrc/quantization/cutlass_w8a8/moe/
373
- vllm/model_executor/layers/fusion/
@@ -408,6 +409,7 @@ steps:
409
- tests/kernels/test_ll_bf16_gemm.py
410
- tests/kernels/test_top_k_per_row.py
411
- tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
412
+ - tests/kernels/quantization/test_block_fp8.py
413
- tests/models/quantization/test_nvfp4.py
414
commands:
415
- nvidia-smi
@@ -442,6 +444,7 @@ steps:
444
pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2'
445
pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py
446
pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
447
+ pytest -v -s tests/kernels/quantization/test_block_fp8.py::test_w8a8_block_fp8_cutlass_matmul
448
pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
449
pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py
450
pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh
+29
-4
@@ -152,6 +152,29 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
152
153
int32_t m = a.size(0), n = b.size(1), k = a.size(1);
154
155
+ // The SM100 blockwise SF copy requires the SFA M extent to be a multiple
156
+ // of 4 (4-row copy atom). For misaligned M on the non-swapAB path, pad only
157
+ // the per-token activation scales (column-major (m, k/128), 1/64 the size
158
+ // of the bf16 activation) and build layout_SFA from the padded extent.
159
+ int32_t m_sf = m;
160
+ std::optional<torch::stable::Tensor> a_scales_pad;
161
+ void const* a_scales_data = a_scales.data_ptr();
162
+ if constexpr (!swap_ab) {
163
+ if (m % 4 != 0) {
164
+ m_sf = (m + 3) & ~3;
165
+ int64_t num_groups = a_scales.numel() / m;
166
+ a_scales_pad = torch::stable::new_empty(
167
+ a_scales, {num_groups * m_sf},
168
+ torch::headeronly::ScalarType::Float);
169
+ cudaMemcpy2DAsync(a_scales_pad->data_ptr(), m_sf * sizeof(float),
170
+ a_scales.data_ptr(), m * sizeof(float),
171
+ m * sizeof(float), num_groups,
172
+ cudaMemcpyDeviceToDevice,
173
+ get_current_cuda_stream(a.get_device()));
174
+ a_scales_data = a_scales_pad->data_ptr();
175
+ }
176
+ }
177
+
178
StrideA a_stride;
179
StrideB b_stride;
180
StrideC c_stride;
@@ -162,16 +185,16 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
185
c_stride =
186
cutlass::make_cute_packed_stride(StrideC{}, swap_ab ? cute::make_shape(n, m, 1) : cute::make_shape(m, n, 1));
187
165
- LayoutSFA layout_SFA = swap_ab ?
188
+ LayoutSFA layout_SFA = swap_ab ?
189
ScaleConfig::tile_atom_to_shape_SFA(make_shape(n, m, k, 1)) :
167
- ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1));
190
+ ScaleConfig::tile_atom_to_shape_SFA(make_shape(m_sf, n, k, 1));
191
LayoutSFB layout_SFB = swap_ab ?
192
ScaleConfig::tile_atom_to_shape_SFB(make_shape(n, m, k, 1)) :
193
ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
194
195
auto a_ptr = static_cast<ElementAB const*>(a.data_ptr());
196
auto b_ptr = static_cast<ElementAB const*>(b.data_ptr());
174
- auto a_scales_ptr = static_cast<ElementBlockScale const*>(a_scales.data_ptr());
197
+ auto a_scales_ptr = static_cast<ElementBlockScale const*>(a_scales_data);
198
auto b_scales_ptr = static_cast<ElementBlockScale const*>(b_scales.data_ptr());
199
200
typename GemmKernel::MainloopArguments mainloop_args{};
@@ -212,7 +235,9 @@ void cutlass_gemm_blockwise_sm100_fp8_dispatch(torch::stable::Tensor& out,
235
236
constexpr int TILE_K = 128;
237
// TODO: better heuristics
215
- bool swap_ab = (m < 16) || (m % 4 != 0);
238
+ // swapAB only wins for small M; misaligned M >= 65 is faster on the
239
+ // non-swapAB path with padded activation scales (measured on GB300).
240
+ bool swap_ab = (m <= 64);
241
bool use_tma_epilogue = (m * n) % 4 == 0;
242
if (!swap_ab) {
243
constexpr int TILE_N = 128;
tests/kernels/quantization/test_block_fp8.py
+7
-2
@@ -162,11 +162,16 @@ def test_w8a8_block_fp8_matmul(M, N, K, block_size, out_dtype, seed):
162
@pytest.mark.skipif(
163
not current_platform.is_cuda(), reason="CUTLASS only supported on CUDA platform."
164
)
165
+@pytest.mark.parametrize(
166
+ # 65/66/67 cover all M%4 residue classes above the SM100 swapAB
167
+ # threshold (m <= 64); 1026 crosses multiple 128-row SF atoms.
168
+ "M",
169
+ [32, 65, 66, 67, 1026],
170
+)
171
@torch.inference_mode()
166
-def test_w8a8_block_fp8_cutlass_matmul():
172
+def test_w8a8_block_fp8_cutlass_matmul(M):
173
# Test simple case where weight.shape % 128 != 0,
174
# like in DSV3 kv_a_proj_with_mqa
169
- M = 32
175
N = 576
176
K = 7168
177
block_size = [128, 128]