Skip to content

Allow caller-provided output tensor in grouped_gemm_quant_wrapper_sm100 #338

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Anerudhan merged 1 commit into from
Jul 4, 2026
Merged

Allow caller-provided output tensor in grouped_gemm_quant_wrapper_sm100 #338

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 19 additions & 1 deletion python/cudnn/grouped_gemm/grouped_gemm_quant/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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}."
)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
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
Expand Down
16 changes: 16 additions & 0 deletions test/python/fe_api/test_grouped_gemm_quant.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)


Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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(
Expand All @@ -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"],
Expand All @@ -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,
Expand Down