diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index b2f77ecd66..db7cff7d42 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -928,3 +928,91 @@ def test_mxfp8_dequantize_columnwise_only_quantized_separately( # Make sure we are not trivially passing the test with pytest.raises(AssertionError): torch.testing.assert_close(x_deq, -x_ref, **_tols[fp8_dtype]) + + +@pytest.mark.parametrize("quantization", _quantization_list) +@pytest.mark.parametrize("usage", ["rowwise", "columnwise", "both"]) +@pytest.mark.parametrize("shape", [(256, 512), (128, 320)], ids=lambda s: f"{s[0]}x{s[1]}") +def test_skip_quantization_with_noop_flag( + quantization: str, usage: str, shape: Tuple[int, int] +) -> None: + """ + Test if the quantization honors the noop flag and skips quantization. + This test only verifies that the kernel skips quantization where it's expected to do so. + It doesn't verify otherwise (kernel correctly quantizes when noop is false) since that is already covered by other tests. + + Parametrized over output usage and shape so that every kernel reachable from + dispatch/quantize.cuh with a noop flag gets exercised, not just the ones recently fixed. + """ + if usage == "columnwise" and quantization in ("fp8", "fp8_delayed_scaling"): + pytest.skip("Delayed scaling does not support columnwise-only quantization") + if usage != "rowwise" and quantization == "nvfp4_row_scaled": + pytest.skip("Row-scaled NVFP4 does not produce columnwise output") + + quantizer = make_quantizer(quantization) + quantizer.set_usage( + rowwise=usage in ("rowwise", "both"), + columnwise=usage in ("columnwise", "both"), + ) + + first_batch = torch.rand(shape, dtype=torch.bfloat16, device="cuda") + second_batch = torch.rand_like(first_batch) + x = first_batch.clone() + + noop = torch.zeros(1, dtype=torch.float32, device="cuda") + + # Allocate and populate the destination outside the graph + quantized = quantizer(x) + + # CUDA graph capture requires the work to be warmed up on a side stream first. + side_stream = torch.cuda.Stream() + side_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side_stream): + for _ in range(3): + quantized.quantize_(x, noop_flag=noop) + torch.cuda.current_stream().wait_stream(side_stream) + + # Capture the CUDA graph + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + quantized.quantize_(x, noop_flag=noop) + + # Clone the underlying buffers of the quantized tensor's output so we can compare them later. + # Note that we can't clone the quantized tensor, but its underlying buffers are torch tensors and can be cloned. + def clone_quantized_output(): + output_buffers = {} + for attr in ( + "_data", + "_transpose", + "_rowwise_data", + "_columnwise_data", + "_scale_inv", + "_rowwise_scale_inv", + "_columnwise_scale_inv", + ): + buf = getattr(quantized, attr, None) + if buf is not None: + output_buffers[attr] = buf.clone() + assert output_buffers, f"{quantization}: quantized tensor exposes no data buffer" + return output_buffers + + # Obtain the quantized x by setting the noop flag to false (zero) + noop.zero_() + graph.replay() + torch.cuda.synchronize() + quantized_without_noop = clone_quantized_output() + + # Reset x to different values and set the noop flag to true (one). + # The quantization should be skipped so the output should not be different. + x.copy_(second_batch) + noop.fill_(1.0) + graph.replay() + torch.cuda.synchronize() + quantized_with_noop = clone_quantized_output() + + for attr in quantized_without_noop: + buf_without_noop = quantized_without_noop[attr] + buf_with_noop = quantized_with_noop[attr] + assert torch.equal( + buf_without_noop, buf_with_noop + ), f"{quantization}/{usage}: noop flag fails to take effect because {attr} changed." diff --git a/transformer_engine/common/cast/fp8/group_quantize_fp8.cuh b/transformer_engine/common/cast/fp8/group_quantize_fp8.cuh index 5186794dd7..082ee0ca54 100644 --- a/transformer_engine/common/cast/fp8/group_quantize_fp8.cuh +++ b/transformer_engine/common/cast/fp8/group_quantize_fp8.cuh @@ -198,8 +198,11 @@ __global__ void __launch_bounds__(THREADS_PER_TILE) const int64_t *__restrict__ offsets_ptr, const int64_t *__restrict__ first_dims_ptr, const int64_t *__restrict__ last_dims_ptr) { - if (noop != nullptr && noop[0] == 1.0f) { - return; + // Activation-fused quantize ignores the noop flag. + if constexpr (!IS_ACT) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } } constexpr bool ROWWISE_OUTPUT = @@ -527,8 +530,11 @@ __global__ void __launch_bounds__(ROWWISE_FLAT_THREADS) const float *__restrict__ scale_ptr, const float *__restrict__ noop, const size_t rows_per_tensor, const size_t cols, const size_t vecs_per_tensor) { - if (noop != nullptr && noop[0] == 1.0f) { - return; + // Activation-fused quantize ignores the noop flag. + if constexpr (!IS_ACT) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } } constexpr size_t nvec = ROWWISE_FLAT_LOAD_SIZE_BYTES / sizeof(IType); @@ -564,8 +570,11 @@ __global__ void __launch_bounds__(ROWWISE_FLAT_THREADS) const float *__restrict__ scale_ptr, const float *__restrict__ noop, const size_t num_tensors, const size_t total_elements, const int64_t *__restrict__ offsets_ptr, const size_t target_blocks) { - if (noop != nullptr && noop[0] == 1.0f) { - return; + // Activation-fused quantize ignores the noop flag. + if constexpr (!IS_ACT) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } } constexpr size_t nvec = ROWWISE_FLAT_LOAD_SIZE_BYTES / sizeof(IType); @@ -666,8 +675,11 @@ __global__ void __launch_bounds__(ROWWISE_FLAT_THREADS) group_cast_fp8_variable_ const IType *__restrict__ input, OType *__restrict__ output_rowwise, const float *__restrict__ scale_ptr, const float *__restrict__ noop, const size_t num_tensors, const size_t total_elements, const int64_t *__restrict__ offsets_ptr) { - if (noop != nullptr && noop[0] == 1.0f) { - return; + // Activation-fused quantize ignores the noop flag. + if constexpr (!IS_ACT) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } } constexpr size_t nvec = ROWWISE_FLAT_LOAD_SIZE_BYTES / sizeof(IType); @@ -732,8 +744,11 @@ __global__ void __launch_bounds__(THREADS_PER_TILE) const float *__restrict__ scale_ptr, const float *__restrict__ noop, const size_t rows_per_tensor, const size_t cols) { - if (noop != nullptr && noop[0] == 1.0f) { - return; + // Activation-fused quantize ignores the noop flag. + if constexpr (!IS_ACT) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } } constexpr bool ROWWISE_OUTPUT = @@ -825,8 +840,11 @@ __global__ void __launch_bounds__(THREADS_PER_TILE) group_cast_fp8_varying_first const float *__restrict__ noop, const size_t num_tensors, const size_t rows_upper_bound, const size_t cols, const size_t total_elements, const int64_t *__restrict__ offsets_ptr, const int64_t *__restrict__ first_dims_ptr) { - if (noop != nullptr && noop[0] == 1.0f) { - return; + // Activation-fused quantize ignores the noop flag. + if constexpr (!IS_ACT) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } } constexpr bool ROWWISE_OUTPUT = diff --git a/transformer_engine/common/cast/fp8/quantize_fp8.cuh b/transformer_engine/common/cast/fp8/quantize_fp8.cuh index bad10c954e..6cf0a4816f 100644 --- a/transformer_engine/common/cast/fp8/quantize_fp8.cuh +++ b/transformer_engine/common/cast/fp8/quantize_fp8.cuh @@ -251,9 +251,17 @@ template __global__ void __launch_bounds__(THREADS_PER_BLOCK) cast_fp8_1D_kernel(const IType *input_ptr, OType *output_ptr, float *const amax_ptr, - float *const scale_inv_ptr, const float *const scale_ptr, const size_t N) { + float *const scale_inv_ptr, const float *const scale_ptr, const size_t N, + const float *const noop) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + // Activation-fused quantize ignores the noop flag. + if constexpr (!IS_ACT) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + } + const size_t block_offset = blockIdx.x * ELEMS_PER_BLOCK; const IType *input = input_ptr + block_offset; OType *output = output_ptr + block_offset; @@ -353,7 +361,7 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) } // namespace quantize_1D_kernel template -void quantize_1D(const Tensor &input, Tensor *output, cudaStream_t stream) { +void quantize_1D(const Tensor &input, const Tensor *noop, Tensor *output, cudaStream_t stream) { using namespace quantize_1D_kernel; const size_t N = product(input.data.shape); @@ -368,6 +376,7 @@ void quantize_1D(const Tensor &input, Tensor *output, cudaStream_t stream) { float *const amax_ptr = reinterpret_cast(output->amax.dptr); float *const scale_inv_ptr = reinterpret_cast(output->scale_inv.dptr); const float *const scale_ptr = reinterpret_cast(output->scale.dptr); + const float *noop_ptr = reinterpret_cast(noop->data.dptr); const dim3 block(THREADS_PER_BLOCK); const dim3 grid(blocks); @@ -379,9 +388,10 @@ void quantize_1D(const Tensor &input, Tensor *output, cudaStream_t stream) { const IType *input_ptr = reinterpret_cast(input.data.dptr); OType *output_ptr = reinterpret_cast(output->data.dptr); - cast_fp8_1D_kernel<<>>( - input_ptr, output_ptr, amax_ptr, scale_inv_ptr, scale_ptr, N);); // NOLINT(*) - ); // NOLINT(*) + cast_fp8_1D_kernel + <<>>(input_ptr, output_ptr, amax_ptr, scale_inv_ptr, scale_ptr, N, + noop_ptr);); // NOLINT(*) + ); // NOLINT(*) NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -463,6 +473,9 @@ template void CastVectorizedUnaryKernelLauncher(const Tensor &input, const Tensor *noop, Tensor *output, cudaStream_t stream) { constexpr float (*UnaryOP)(float, const ParamOP &) = (OP == nullptr) ? detail::identity : OP; + // If fusion is involved, ignore the noop ptr because we should compute the activation anyway. + const float *const noop_ptr = + (OP == nullptr) ? reinterpret_cast(noop->data.dptr) : nullptr; const size_t N = product(input.data.shape); TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( input.data.dtype, IType, @@ -471,8 +484,7 @@ void CastVectorizedUnaryKernelLauncher(const Tensor &input, const Tensor *noop, if (!is_fp8_dtype(output->data.dtype) || is_tensor_scaling(output->scaling_mode)) { constexpr int nvec = 32 / sizeof(IType); VectorizedUnaryKernelLauncher( - reinterpret_cast(input.data.dptr), - reinterpret_cast(noop->data.dptr), + reinterpret_cast(input.data.dptr), noop_ptr, reinterpret_cast(output->data.dptr), reinterpret_cast(output->scale.dptr), reinterpret_cast(output->amax.dptr), @@ -537,7 +549,7 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, is_aligned_tensor_data(input, TMA_GMEM_ALIGNMENT) && is_aligned_tensor_data(*output, TMA_GMEM_ALIGNMENT)) { // Aligned AND FP8 - quantize_1D(input, output, stream); + quantize_1D(input, noop, output, stream); } else { // Unaligned CastVectorizedUnaryKernelLauncher(input, noop, output, stream); diff --git a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh index 33aadb5bfa..203f569471 100644 --- a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh +++ b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh @@ -287,7 +287,9 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma const float epsilon, const bool pow_2_scales, const float* __restrict__ noop_ptr, float* __restrict__ dbias_workspace) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 - if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; + // Skipping is only safe without dbias: the grouped_reduce_dbias launch is unconditional, so an + // early return would leave it reducing a workspace this kernel never wrote. + if (dbias_workspace == nullptr && noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; const size_t tile_x = blockIdx.x; const size_t tile_y_global = blockIdx.y; @@ -574,7 +576,9 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke const size_t R_total, const float epsilon, const bool pow_2_scales, const float* __restrict__ noop_ptr, float* __restrict__ dbias_workspace) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 - if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; + // Skipping is only safe without dbias: the grouped_reduce_dbias launch is unconditional, so an + // early return would leave it reducing a workspace this kernel never wrote. + if (dbias_workspace == nullptr && noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; const size_t tile_x = blockIdx.x; const size_t tile_y_global = blockIdx.y; diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 14832573d7..c02d6b891a 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -464,7 +464,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel constexpr bool COMPUTE_ACTIVATIONS = IS_DACT || IS_ACT; constexpr bool NO_ACTIVATIONS = !COMPUTE_ACTIVATIONS; - if constexpr (NO_ACTIVATIONS) { + if constexpr (NO_ACTIVATIONS && !IS_DBIAS) { if (noop != nullptr && noop[0] == 1.0f) { return; } diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 8c57f112d2..4c5c94fb02 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -64,7 +64,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) using transformer_engine::dispatch::mxfp8::swizzle::gemm_swizzled_scale_idx; - if constexpr (NO_ACTIVATIONS) { + if constexpr (NO_ACTIVATIONS && !IS_DBIAS) { if (noop != nullptr && noop[0] == 1.0f) { return; } @@ -646,7 +646,12 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, float *const workspace_ptr = IS_DBIAS ? reinterpret_cast(workspace->data.dptr) : nullptr; float *const amax_ptr = reinterpret_cast(output->amax.dptr); - const float *noop_ptr = reinterpret_cast(noop->data.dptr); + constexpr bool NO_ACTIVATIONS = !(IS_DACT || IS_ACT); + // zero_scales_kernel should ignore the noop tensor whenever the quantization kernel also ignores it, + // which happens when there are no fusions (NO ACT and DBIAS). Since zero_scales_kernel is not templated with these variants + // it doesn't know when to ignore, so we need to override the noop pointer to nullptr before passing it to the kernel. + const float *const noop_ptr = + (NO_ACTIVATIONS && !IS_DBIAS) ? reinterpret_cast(noop->data.dptr) : nullptr; TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input.dtype(), IType,