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 523b63640..d0302e67b 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 @@ -22,6 +22,7 @@ Extension: moe_sched_extension.py (WgradDense / WgradDiscrete) """ +import re from importlib.metadata import PackageNotFoundError, version from typing import Literal, Type, Tuple, Optional @@ -66,7 +67,19 @@ def _using_internal_cutlass_dsl() -> bool: return True -_USING_INTERNAL_CUTLASS_DSL = _using_internal_cutlass_dsl() +def _cutlass_dsl_needs_fp4_layout_workaround() -> bool: + # Public cutlass-dsl wheels before 4.8 interpret packed sub-byte + # from_dlpack layouts in byte units, so the FP4 A/B layouts must be + # recast to element units. + if _using_internal_cutlass_dsl(): + return False + match = re.match(r"(\d+)\.(\d+)", getattr(cutlass, "__version__", "") or "") + if match is None: + return False + return (int(match.group(1)), int(match.group(2))) < (4, 8) + + +_NEEDS_FP4_LAYOUT_WORKAROUND = _cutlass_dsl_needs_fp4_layout_workaround() class BlockScaledMoEGroupedGemmWgradKernel: @@ -329,10 +342,10 @@ def __call__( out_single_expert: Optional[cute.Tensor] = None, ) -> None: - # Public CUTLASS DSL 4.5 needs the packed-FP4 from_dlpack layout - # workaround. Rubin and the internal DSL wheel consume the native - # 4-bit layout directly. - needs_fp4_layout_workaround = self.architecture != "sm_107" and not _USING_INTERNAL_CUTLASS_DSL + # Public CUTLASS DSL < 4.8 needs the packed-FP4 from_dlpack layout + # workaround. Rubin, the internal DSL wheel, and public wheels >= 4.8 + # consume the native 4-bit layout directly. + needs_fp4_layout_workaround = self.architecture != "sm_107" and _NEEDS_FP4_LAYOUT_WORKAROUND if cutlass.const_expr(needs_fp4_layout_workaround and mat_a.iterator.dtype.width < 8): mat_a = cute.make_tensor( mat_a.iterator, diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py index 98a998ae6..d246d0707 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_quant.py @@ -178,13 +178,7 @@ def _blocked_sf_to_flat(sf_tensor: torch.Tensor, rows: int, cols: int) -> torch. row_idx = torch.arange(rows, device=sf_tensor.device, dtype=torch.long).view(rows, 1) col_idx = torch.arange(sf_cols, device=sf_tensor.device, dtype=torch.long).view(1, sf_cols) col_blocks = sf_tensor.shape[1] // 4 - linear = ( - (row_idx // 128) * col_blocks * 512 - + (col_idx // 4) * 512 - + (row_idx % 32) * 16 - + ((row_idx // 32) % 4) * 4 - + (col_idx % 4) - ) + linear = (row_idx // 128) * col_blocks * 512 + (col_idx // 4) * 512 + (row_idx % 32) * 16 + ((row_idx // 32) % 4) * 4 + (col_idx % 4) return sf_tensor.flatten()[linear]