Add non-TMA MXFP8 cast-only kernels for specialized rowwise-only and row+colwise

transformer_engine/common/CMakeLists.txt

@@ -264,6 +264,8 @@ list(APPEND transformer_engine_cuda_arch_specific_sources
      activation/swiglu_grouped.cu
      activation/swiglu_grouped_dbias.cu
      cast/cast.cu
+     cast/mxfp8/specialized/cast_rowwise.cu
+     cast/mxfp8/specialized/cast_bidim.cu
      cast/cast_dbias.cu
      cast/cast_grouped.cu
      cast/cast_grouped_dbias.cu

transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh

@@ -21,6 +21,8 @@
 #include "../../util/ptx_arch_spec.cuh"
 #include "../../utils.cuh"
 #include "../core/common.cuh"
+#include "specialized/cast_bidim.h"
+#include "specialized/cast_rowwise.h"
 #include "specialized/quantize_mxfp8.cuh"
 #include "swizzle.cuh"
 
@@ -745,6 +747,19 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop,
                   !use_2d_quantization && scaling_type_has_specialized_support) {
                 switch (scaling_type) {
                   case ScalingType::ROWWISE: {
+                    // The register-resident kernel supersedes the staged one below
+                    // wherever it applies: 10-20% faster on the cast-only path.
+                    // hasSpec has already established cast-only, so what is left
+                    // is the input type and the scale layout, neither of which
+                    // that kernel generalizes over yet.
+                    if constexpr (std::is_same_v<IType, bf16> && !WITH_GEMM_SWIZZLED_SCALES) {
+                      specialized::launch_cast_rowwise<OType>(
+                          input.data.dptr, output->data.dptr,
+                          reinterpret_cast<void *>(scales_rowwise_ptr), static_cast<int>(rows),
+                          static_cast<int>(cols), static_cast<int>(scale_stride_rowwise), stream);
+                      break;
+                    }
+
                     using traits = specialized::CastTraits<IType, OType, true, false,
                                                            WITH_GEMM_SWIZZLED_SCALES>;
                     auto kernel = specialized::quantize_mxfp8_kernel_cast_only<traits>;
@@ -765,6 +780,25 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop,
                     break;
                   }
                   case ScalingType::BIDIMENSIONAL: {
+                    // The register-resident kernel supersedes the TMA one below
+                    // wherever it applies: it reads its 32-row tile once and drives
+                    // both the rowwise and colwise passes from registers, roughly
+                    // 1.2x faster on shapes that fit its tiling.  hasSpec has
+                    // already established cast-only, leaving the input type, the
+                    // scale layout, and the tile alignment.
+                    if constexpr (std::is_same_v<IType, bf16> && !WITH_GEMM_SWIZZLED_SCALES) {
+                      if (rows % 32 == 0 && cols % 256 == 0) {
+                        specialized::launch_cast_bidim<OType>(
+                            input.data.dptr, output->data.dptr,
+                            reinterpret_cast<void *>(scales_rowwise_ptr),
+                            output->columnwise_data.dptr,
+                            reinterpret_cast<void *>(scales_colwise_ptr), static_cast<int>(rows),
+                            static_cast<int>(cols), static_cast<int>(scale_stride_rowwise),
+                            static_cast<int>(scale_stride_colwise), stream);
+                        break;
+                      }
+                    }
+
                     using traits =
                         specialized::CastTraitsSwizzle<IType, OType,
                                                        /*NumStages=*/2, /*IterM=*/1, /*IterN=*/4,