From ba50412200ddd7c56d13040161d04b6428e23c5d Mon Sep 17 00:00:00 2001 From: Jiangfei Duan Date: Wed, 16 Sep 2026 20:48:04 -0700 Subject: [PATCH] fix: acquire SM100 grouped-wgrad tensor maps before TMA use Acquire device-written descriptors in the consuming warp when switching experts. Reuse acquires across consecutive tiles to avoid per-tile overhead. Add an eager/graph regression test that changes tensor addresses while reusing a cached plan. Signed-off-by: Jiangfei Duan --- .../moe_blockscaled_grouped_gemm_wgrad.py | 23 ++++++ .../test_wgrad_tma_descriptor_reuse.py | 79 +++++++++++++++++++ 2 files changed, 102 insertions(+) create mode 100644 test/python/fe_api/grouped_gemm/test_wgrad_tma_descriptor_reuse.py diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py index 8e665cc726..f16514eb1b 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py @@ -50,6 +50,7 @@ MoEWeightMode, WGradInputOrder, WgradSfTensormapConstructor, + gmem_ptr_to_generic, ) from ..moe_sched_extension import ( WgradScaledGemmSchedExtension, @@ -1023,10 +1024,23 @@ class SharedStorage: sched_pipeline.consumer_release(sched_consumer_state) sched_consumer_state.advance() + last_acquired_expert = cutlass.Int32(-1) while work_tile_info.is_valid_tile: k_tile_cnt = work_tile_info.k_tile_cnt ext.update_expert_info(offs, work_tile_info.expert_idx) + # The preceding device kernel writes descriptors through the generic + # proxy. Acquire them before TMA reads, including on graph replay. + # Descriptors remain immutable within this kernel, so consecutive + # tiles of the same expert can reuse the acquire. + if work_tile_info.expert_idx != last_acquired_expert: + if cutlass.const_expr(self.input_order == WGradInputOrder.TensorRagged): + cpasync.fence_tma_desc_acquire(gmem_ptr_to_generic(desc_workspace.get_desc_ptr("a", work_tile_info.expert_idx))) + cpasync.fence_tma_desc_acquire(gmem_ptr_to_generic(desc_workspace.get_desc_ptr("b", work_tile_info.expert_idx))) + cpasync.fence_tma_desc_acquire(gmem_ptr_to_generic(desc_workspace.get_desc_ptr("sfa", work_tile_info.expert_idx))) + cpasync.fence_tma_desc_acquire(gmem_ptr_to_generic(desc_workspace.get_desc_ptr("sfb", work_tile_info.expert_idx))) + last_acquired_expert = work_tile_info.expert_idx + real_a, desc_ptr_a = ext.get_gmem_tensor( "a", mA_mkl, @@ -1458,10 +1472,19 @@ class SharedStorage: sched_pipeline.consumer_release(sched_consumer_state) sched_consumer_state.advance() + last_acquired_expert = cutlass.Int32(-1) while work_tile_info.is_valid_tile: k_tile_cnt = work_tile_info.k_tile_cnt ext.update_expert_info(offs, work_tile_info.expert_idx) + # Only the warp issuing TMA stores needs the output descriptor. + # Empty experts still store zeros and must acquire it as well. + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + if warp_idx == self.epilogue_warp_id[0]: + if work_tile_info.expert_idx != last_acquired_expert: + cpasync.fence_tma_desc_acquire(gmem_ptr_to_generic(desc_workspace.get_desc_ptr("c", work_tile_info.expert_idx))) + last_acquired_expert = work_tile_info.expert_idx + real_c, desc_ptr_c = ext.get_gmem_tensor( "c", mC_mnl, diff --git a/test/python/fe_api/grouped_gemm/test_wgrad_tma_descriptor_reuse.py b/test/python/fe_api/grouped_gemm/test_wgrad_tma_descriptor_reuse.py new file mode 100644 index 0000000000..b14b1ac02e --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_wgrad_tma_descriptor_reuse.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A cached wgrad plan must acquire descriptors rewritten between graph nodes.""" + +import pytest +import torch +import cudnn + + +@pytest.mark.L0 +@pytest.mark.parametrize("discrete", [False, True]) +@pytest.mark.parametrize("ragged", [False, True]) +@pytest.mark.parametrize("accumulate", [False, True]) +def test_wgrad_tma_descriptor_reuse(discrete, ragged, accumulate): + if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + pytest.skip("Exercises the SM100 blockscaled wgrad kernel.") + + stream = torch.cuda.Stream() + calls, outputs = [], [] + # Identical signatures reuse one plan/workspace, but the device helper must + # rebind its descriptors to distinct input, scale and output addresses. + for value, scale_byte, tokens in ((1, 127, (128, 128, 0, 256)), (2, 128, (256, 0, 128, 128))): + a = torch.full((512, 256), value, device="cuda").to(torch.float8_e4m3fn).T + b = torch.ones((512, 256), device="cuda").to(torch.float8_e4m3fn) + # Inputs are constant, so their physical layout is also valid in ragged mode. + sa = torch.full((256, 16), scale_byte, device="cuda", dtype=torch.uint8).view(torch.float8_e8m0fnu) + sb = sa.clone() + offsets = torch.tensor(tokens, device="cuda", dtype=torch.int32).cumsum(0, dtype=torch.int32) + if discrete: + dest = [torch.empty((256, 256), device="cuda", dtype=torch.bfloat16) for _ in tokens] + pointers = torch.tensor([out.data_ptr() for out in dest], device="cuda", dtype=torch.int64) + output_args = dict(output_mode="discrete", wgrad_ptrs=pointers) + else: + dense = torch.empty((4, 256, 256), device="cuda", dtype=torch.bfloat16) + dest = list(dense.unbind()) + output_args = dict(output_mode="dense", wgrad_tensor=dense) + calls.append( + dict( + a_tensor=a, + b_tensor=b, + sfa_tensor=sa, + sfb_tensor=sb, + offsets_tensor=offsets, + **output_args, + wgrad_dtype=torch.bfloat16, + acc_dtype=torch.float32, + sf_vec_size=32, + input_order="tensor_ragged" if ragged else "tensor2d", + accumulate_on_output=accumulate, + current_stream=stream.cuda_stream, + ) + ) + scale = 2 ** (scale_byte - 127) + outputs.extend((out, n * value * scale * scale + int(accumulate)) for out, n in zip(dest, tokens)) + + def run(): + for kwargs in calls: + cudnn.grouped_gemm_wgrad_wrapper_sm100(**kwargs) + + def check(fn): + # Reset outside capture on every replay: stale warmup outputs must not + # hide a missing store, and accumulation always starts from a known value. + for out, _ in outputs: + out.fill_(1 if accumulate else -1) + fn() + for out, expected in outputs: + torch.testing.assert_close(out, torch.full_like(out, expected), rtol=0, atol=0) + + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + check(run) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run() + for _ in range(8): + check(graph.replay) + torch.cuda.current_stream().wait_stream(stream)