diff --git a/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py b/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py index a89bcf004..eaf7d7d98 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py @@ -1172,6 +1172,7 @@ def grouped_gemm_quant_wrapper_sm100( row_scale_tensor: Optional[torch.Tensor] = None, acc_dtype: torch.dtype = torch.float32, d_dtype: torch.dtype = torch.bfloat16, + d_tensor: Optional[torch.Tensor] = None, cd_major: str = "n", mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, @@ -1212,6 +1213,9 @@ def grouped_gemm_quant_wrapper_sm100( conversion. acc_dtype: Accumulator data type d_dtype: Output D tensor data type + d_tensor: Optional preallocated output tensor to write into instead of + allocating. Must match the internal layout: shape (valid_m, n_out, 1), + stride (n_out, 1, valid_m * n_out), dtype d_dtype, on a_tensor.device. cd_major: CD major dimension (only "n"-major layout is supported) mma_tiler_mn: MMA tiler shape cluster_shape_mn: Cluster shape @@ -1288,7 +1292,21 @@ def grouped_gemm_quant_wrapper_sm100( _logger.debug("grouped_gemm_quant_wrapper_sm100: Creating output tensors") if cd_major == "n": - d_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=d_dtype, device=a_tensor.device) + expected_shape = (valid_m, n_out, 1) + expected_stride = (n_out, 1, valid_m * n_out) + if d_tensor is None: + d_tensor = torch.empty_strided(expected_shape, expected_stride, dtype=d_dtype, device=a_tensor.device) + elif ( + tuple(d_tensor.shape) != expected_shape + or tuple(d_tensor.stride()) != expected_stride + or d_tensor.dtype != d_dtype + or d_tensor.device != a_tensor.device + ): + raise ValueError( + f"d_tensor must have shape {expected_shape}, stride {expected_stride}, " + f"dtype {d_dtype}, device {a_tensor.device}, but got shape {tuple(d_tensor.shape)}, " + f"stride {tuple(d_tensor.stride())}, dtype {d_tensor.dtype}, device {d_tensor.device}." + ) d_col_tensor = ( torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=d_dtype, device=a_tensor.device) if is_low_precision_output_config diff --git a/test/python/fe_api/test_grouped_gemm_quant.py b/test/python/fe_api/test_grouped_gemm_quant.py index 03fcae24c..88b9092ab 100644 --- a/test/python/fe_api/test_grouped_gemm_quant.py +++ b/test/python/fe_api/test_grouped_gemm_quant.py @@ -246,6 +246,7 @@ def test_grouped_gemm_quant_wrapper_fp8( discrete_col_sfd=discrete_col_sfd, use_dynamic_sched=use_dynamic_sched, request=request, + provide_d_tensor=True, ) @@ -1110,6 +1111,7 @@ def _test_grouped_gemm_quant_wrapper( input_mutator=None, use_dynamic_sched=False, enable_bias=False, + provide_d_tensor=False, ): """Test GroupedGemmQuant API via the wrapper function (with caching).""" try: @@ -1151,6 +1153,16 @@ def _test_grouped_gemm_quant_wrapper( if input_mutator is not None: input_mutator(inputs, cfg) + caller_d_tensor = None + if provide_d_tensor: + valid_m, n_out = inputs["a_tensor"].shape[0], cfg["n"] + caller_d_tensor = torch.empty_strided( + (valid_m, n_out, 1), + (n_out, 1, valid_m * n_out), + dtype=cfg["d_dtype"], + device=inputs["a_tensor"].device, + ) + try: for _ in range(2): # Run twice to test caching path outputs = grouped_gemm_quant_wrapper_sm100( @@ -1166,6 +1178,7 @@ def _test_grouped_gemm_quant_wrapper( row_scale_tensor=inputs.get("row_scale_tensor"), acc_dtype=cfg["acc_dtype"], d_dtype=cfg["d_dtype"], + d_tensor=caller_d_tensor, cd_major=cfg["cd_major"], mma_tiler_mn=cfg["mma_tiler_mn"], cluster_shape_mn=cfg["cluster_shape_mn"], @@ -1179,6 +1192,9 @@ def _test_grouped_gemm_quant_wrapper( except (ValueError, NotImplementedError) as e: pytest.skip(f"Unsupported testcase: {e}") + if provide_d_tensor: + assert outputs["d_tensor"].data_ptr() == caller_d_tensor.data_ptr() + check_ref_grouped_gemm_quant( inputs, outputs,