[None][fix] DenseGEMM MoE: keep FC1/FC2 global accesses inside the problem

tensorrt_llm/_torch/kernels/blackwell/dense_blockscaled_gemm_act_fusion.py

@@ -1822,7 +1822,10 @@ def kernel(
                                         << cutlass.Uint32(scale_byte_idx * 8)
                                     )
                         else:
-                            cute.autovec_copy(tCrSFC, tCgSFC)
+                            # Cluster-padded and partial N tiles can reach past the SF tensor.
+                            if sfc_subtile_idx_mn[0] < cute.size(tCgSFC_mn, mode=[3]):
+                                if sfc_subtile_idx_mn[1] < cute.size(tCgSFC_mn, mode=[4]):
+                                    cute.autovec_copy(tCrSFC, tCgSFC)
 
                         #
                         # Compute quantized output values and convert to C type

tensorrt_llm/_torch/moe/kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py

@@ -2620,11 +2620,11 @@ def kernel(
                         tCrSFC.store(tCrSFC_pvscale.load().to(self.sf_dtype))
 
                         #
-                        # Store SFC to global memory
+                        # Store SFC to global memory (guarded for the N boundary)
                         #
-                        # TODO: Need to think about predicate on it
-                        # if cute.elem_less():
-                        cute.autovec_copy(tCrSFC, tCgSFC)
+                        # SF blocks past the last column alias the next 128 rows' SFs.
+                        if sfc_subtile_idx_mn[1] < cute.size(tCgSFC_mn, mode=[4]):
+                            cute.autovec_copy(tCrSFC, tCgSFC)
 
                         #
                         # Compute quantized output values and convert to C type

tensorrt_llm/_torch/moe/kernels/blackwell/moe_as_dense_gemm/fc1.py

@@ -1424,8 +1424,12 @@ def kernel(
             )
 
             ### core code to get alpha_scale (pre-SwiGLU)
-            unroll_scale_idx = work_tile.tile_idx[1] // (
-                self.weight_per_expert // self.cta_tile_shape_mnk[1]
+            # A cluster can run past the last N tile, i.e. past the last expert: clamp the alpha
+            # index there and skip the tile's direct global accesses below.
+            num_experts = alpha_scale.shape[0]
+            n_tiles_per_expert = self.weight_per_expert // self.cta_tile_shape_mnk[1]
+            unroll_scale_idx = cutlass.min(
+                work_tile.tile_idx[1] // n_tiles_per_expert, num_experts - 1
             )
             current_alpha_scale = alpha_scale[unroll_scale_idx, work_tile.tile_idx[2]]
 
@@ -1614,13 +1618,14 @@ def kernel(
                         m_start = cur_tile_coord[0] * self.epi_tile[0]
                         if m_start < m_total:
                             if epi_tidx < (m_total - m_start):
-                                current_alpha_scale_post = tTR_gAlphaPost[
-                                    epi_tidx, cur_tile_coord[0], expert_idx, batch_idx
-                                ]
-                                for i in cutlass.range_constexpr(cute.size(tCompute)):
-                                    tCompute[i] = tCompute[i] * cutlass.Float32(
-                                        current_alpha_scale_post
-                                    )
+                                if expert_idx < num_experts:
+                                    current_alpha_scale_post = tTR_gAlphaPost[
+                                        epi_tidx, cur_tile_coord[0], expert_idx, batch_idx
+                                    ]
+                                    for i in cutlass.range_constexpr(cute.size(tCompute)):
+                                        tCompute[i] = tCompute[i] * cutlass.Float32(
+                                            current_alpha_scale_post
+                                        )
 
                     if cutlass.const_expr(self.generate_sfc):
                         #
@@ -1701,10 +1706,11 @@ def kernel(
                         tCrSFC.store(tCrSFC_pvscale.load().to(self.sf_dtype))
 
                         #
-                        # Store SFC to global memory (guarded for M boundary)
+                        # Store SFC to global memory (guarded for the M and N boundaries)
                         #
                         if m_tile_in_bounds:
-                            cute.autovec_copy(tCrSFC, tCgSFC)
+                            if expert_idx < num_experts:
+                                cute.autovec_copy(tCrSFC, tCgSFC)
 
                         #
                         # Compute quantized output values and convert to C type
@@ -1803,8 +1809,8 @@ def kernel(
                 # then we don't need to update alpha_scale,
                 # otherwise we need to update alpha_scale with ldg
                 if work_tile.is_valid_tile and prev_work_tile_tile_idx_n != work_tile.tile_idx[1]:
-                    unroll_scale_idx = work_tile.tile_idx[1] // (
-                        self.weight_per_expert // self.cta_tile_shape_mnk[1]
+                    unroll_scale_idx = cutlass.min(
+                        work_tile.tile_idx[1] // n_tiles_per_expert, num_experts - 1
                     )
                     current_alpha_scale = alpha_scale[unroll_scale_idx, work_tile.tile_idx[2]]
 

tensorrt_llm/_torch/moe/kernels/blackwell/moe_as_dense_gemm/fc2.py

@@ -1222,6 +1222,8 @@ def kernel(
                 galpha_scale_mnl_current_tile = galpha_scale_mnl[
                     None, alpha_tile_idx, None, batch_idx_w6
                 ]
+                # Bound alpha loads by the global token row, not the row within the tile.
+                alpha_row_base = alpha_tile_idx * self.alpha_scale_dim
 
                 tAgAlphaScale = thr_copy_alpha_scale.partition_S(galpha_scale_mnl_current_tile)
                 tAsAlphaScale = thr_copy_alpha_scale.partition_D(sAlphaScale)
@@ -1251,7 +1253,8 @@ def kernel(
                     for iter_idx in cutlass.range(num_iters, unroll_full=True):
                         iter_coord = (None, iter_idx)
                         pred = cutlass.Boolean(
-                            32 * iter_idx + cute.arch.lane_idx() < malpha_scale_mnl.shape[0]
+                            alpha_row_base + 32 * iter_idx + cute.arch.lane_idx()
+                            < malpha_scale_mnl.shape[0]
                         )
                         if pred:
                             cute.copy(
@@ -1865,7 +1868,8 @@ def kernel(
 
                             m_global = m_base_tile + m_local_start
                             n_global = n_base_subtile + n_local
-                            if m_global + 7 < m_total:
+                            # Bound by C itself: with swap_ab, M is the hidden size and N the tokens.
+                            if m_global + 7 < mC_raw.shape[0] and n_global < mC_raw.shape[1]:
                                 rVec = cute.make_rmem_tensor(
                                     cute.make_layout((4, 2), stride=(2, 1)),
                                     self.c_dtype,
@@ -1909,21 +1913,25 @@ def kernel(
                                 1
                             ] + subtile_idx * cute.size(tTR_rC)
 
+                            # Unlike TMA stores, atomics are not clipped at the last column.
                             for index in cutlass.range(self.epi_loop_size_atomic, unroll_full=True):
                                 coord_n = base_coord_n + index * self.element_offset_atomic
-                                scatter_out = cute.domain_offset((0, coord_n, 0), scatter_base)
-                                if cutlass.const_expr(self.c_dtype == cutlass.Float16):
-                                    vectorized_atomic_add_fp16x8(
-                                        rOut_epi[index, None, None], scatter_out
-                                    )
-                                elif cutlass.const_expr(self.c_dtype == cutlass.BFloat16):
-                                    vectorized_atomic_add_bf16x8(
-                                        rOut_epi[index, None, None], scatter_out
-                                    )
-                                elif cutlass.const_expr(self.c_dtype == cutlass.Float32):
-                                    vectorized_atomic_add_fp32x2(rOut_epi[index, None], scatter_out)
-                                else:
-                                    atomic_add_func(rOut_epi[index], scatter_out)
+                                if coord_n < mC_raw.shape[1]:
+                                    scatter_out = cute.domain_offset((0, coord_n, 0), scatter_base)
+                                    if cutlass.const_expr(self.c_dtype == cutlass.Float16):
+                                        vectorized_atomic_add_fp16x8(
+                                            rOut_epi[index, None, None], scatter_out
+                                        )
+                                    elif cutlass.const_expr(self.c_dtype == cutlass.BFloat16):
+                                        vectorized_atomic_add_bf16x8(
+                                            rOut_epi[index, None, None], scatter_out
+                                        )
+                                    elif cutlass.const_expr(self.c_dtype == cutlass.Float32):
+                                        vectorized_atomic_add_fp32x2(
+                                            rOut_epi[index, None], scatter_out
+                                        )
+                                    else:
+                                        atomic_add_func(rOut_epi[index], scatter_out)
                 else:
                     #
                     # Standard TMA store path (split_k=1)

tests/unittest/_torch/moe/kernels/parallel/test_moe_densegemm.py

@@ -237,6 +237,48 @@ def test_nvfp4_dense_gemm_swiglu_blackwell(
     )
 
 
+@pytest.mark.skipif(
+    get_sm_version() not in (100, 103),
+    reason="This test is only supported on SM 100 and SM 103 GPUs",
+)
+def test_nvfp4_dense_gemm_swiglu_cluster_padded_n_tiles():
+    """N tiles that the cluster shape pads past the last expert leave every output byte as is."""
+    from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLNVFP4DenseGemmSwigluRunner
+
+    sf_vec_size = 16
+    # 7 experts x 1536 columns are 42 N tiles of 256: cluster N 4 pads two of them, cluster N 2
+    # none. 1024 tokens take several persistent waves, so padded tiles run after valid ones.
+    num_expert, weight_per_expert, m, k = 7, 1536, 1024, 256
+    a_bf16 = torch.randint(-1, 2, (m, k), dtype=torch.int32, device="cuda").to(torch.bfloat16)
+    b_bf16 = torch.randint(
+        -1, 2, (num_expert, weight_per_expert, k), dtype=torch.int32, device="cuda"
+    ).to(torch.bfloat16)
+    a_global_sf = a_bf16.abs().max().float() / (448 * 6)
+    b_global_sf = b_bf16.abs().amax(dim=(1, 2)).float() / (448 * 6)
+    a, a_sf = torch.ops.trtllm.fp4_quantize(a_bf16, 1 / a_global_sf, sf_vec_size, False)
+    b, b_sf = torch.ops.trtllm.fp4_quantize(b_bf16, 1 / b_global_sf, sf_vec_size, False)
+    inputs = [
+        a.view(torch.float4_e2m1fn_x2),
+        b.view(torch.float4_e2m1fn_x2),
+        a_sf,
+        b_sf.view(num_expert, weight_per_expert, k // sf_vec_size),
+        a_global_sf * b_global_sf,
+        torch.rand(m, num_expert, dtype=torch.float32, device="cuda") * 0.4 + 0.8,
+        # |up * silu(gate)| <= k * k for entries in [-1, 1].
+        torch.tensor([448 * 6 / k**2], dtype=torch.float32, device="cuda"),
+    ]
+    runner = CuteDSLNVFP4DenseGemmSwigluRunner(
+        num_expert, weight_per_expert, torch.float4_e2m1fn_x2, sf_vec_size
+    )
+    padded_c, padded_sf = runner(inputs, tactic=((128, 256), (1, 4)))
+    c, c_sf = runner(inputs, tactic=((128, 256), (1, 2)))
+
+    assert torch.equal(padded_c.view(torch.uint8), c.view(torch.uint8))
+    assert torch.equal(padded_sf, c_sf), (
+        f"{(padded_sf != c_sf).sum().item()} scale-factor bytes differ"
+    )
+
+
 def nvfp4_dense_gemm_fc2_ref(
     a: torch.Tensor,
     b: torch.Tensor,