[Common] Ensure quantization kernels handle noop properly

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."

transformer_engine/common/cast/fp8/quantize_fp8.cuh

@@ -251,9 +251,14 @@ template <bool IS_ACT, typename ParamOP, float (*OP)(float, const ParamOP &), ty
           typename OType>
 __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)
 
+  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 +358,7 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK)
 }  // namespace quantize_1D_kernel
 
 template <bool IS_ACT, typename ParamOP, float (*OP)(float, const ParamOP &)>
-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 +373,7 @@ void quantize_1D(const Tensor &input, Tensor *output, cudaStream_t stream) {
   float *const amax_ptr = reinterpret_cast<float *>(output->amax.dptr);
   float *const scale_inv_ptr = reinterpret_cast<float *>(output->scale_inv.dptr);
   const float *const scale_ptr = reinterpret_cast<float *>(output->scale.dptr);
+  const float *noop_ptr = reinterpret_cast<const float *>(noop->data.dptr);
 
   const dim3 block(THREADS_PER_BLOCK);
   const dim3 grid(blocks);
@@ -379,9 +385,10 @@ void quantize_1D(const Tensor &input, Tensor *output, cudaStream_t stream) {
           const IType *input_ptr = reinterpret_cast<const IType *>(input.data.dptr);
           OType *output_ptr = reinterpret_cast<OType *>(output->data.dptr);
 
-          cast_fp8_1D_kernel<IS_ACT, ParamOP, OP, IType, OType><<<grid, block, 0, stream>>>(
-              input_ptr, output_ptr, amax_ptr, scale_inv_ptr, scale_ptr, N););  // NOLINT(*)
-  );                                                                            // NOLINT(*)
+          cast_fp8_1D_kernel<IS_ACT, ParamOP, OP, IType, OType>
+          <<<grid, block, 0, stream>>>(input_ptr, output_ptr, amax_ptr, scale_inv_ptr, scale_ptr, N,
+                                       noop_ptr););  // NOLINT(*)
+  );                                                 // NOLINT(*)
   NVTE_CHECK_CUDA(cudaGetLastError());
 }
 
@@ -537,7 +544,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<IS_ACT, ParamOP, OP>(input, output, stream);
+        quantize_1D<IS_ACT, ParamOP, OP>(input, noop, output, stream);
       } else {
         // Unaligned
         CastVectorizedUnaryKernelLauncher<ParamOP, OP>(input, noop, output, stream);

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;

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;
     }

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<float *>(workspace->data.dptr) : nullptr;
   float *const amax_ptr = reinterpret_cast<float *>(output->amax.dptr);
-  const float *noop_ptr = reinterpret_cast<const float *>(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<const float *>(noop->data.dptr) : nullptr;
 
   TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(
       input.dtype(), IType,