Build uploaded using `kernels` (batch 3/10).

.gitattributes

@@ -13,3 +13,4 @@ build/torch29-cxx11-cu129-x86_64-linux/_deep_gemm_cuda_a68a39f.abi3.so filter=lf
 build/torch29-cxx11-cu130-x86_64-linux/_deep_gemm_cuda_a68a39f.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch210-cxx11-cu126-aarch64-linux/_deep_gemm_cuda_a68a39f.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch210-cxx11-cu128-aarch64-linux/_deep_gemm_cuda_a68a39f.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-cu130-aarch64-linux/_deep_gemm_cuda_a68a39f.abi3.so filter=lfs diff=lfs merge=lfs -text

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/interleaved_epilogue.h

@@ -0,0 +1,407 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+  \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
+
+  The epilogue rearranges the result of a matrix product through shared memory to match canonical
+  tensor layouts in global memory. Epilogues support conversion and reduction operations.
+
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/array.h"
+#include "cutlass/layout/vector.h"
+#include "cutlass/layout/tensor.h"
+#include "cutlass/tensor_coord.h"
+#include "cutlass/aligned_buffer.h"
+
+#include "cutlass/gemm/gemm.h"
+
+#include "cutlass/transform/pitch_linear_thread_map.h"
+#include "cutlass/transform/threadblock/regular_tile_iterator.h"
+
+#include "cutlass/epilogue/threadblock/epilogue_base_streamk.h"
+#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace threadblock {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Epilogue operator without splitk
+template <
+    /// Shape of threadblock tile (concept: GemmShape)
+    typename Shape_,
+    /// Warp-level MMA operator (concept: gemm::warp::MmaTensorOp)
+    typename WarpMmaOperator_,
+    /// Number of partitions of the K dimension
+    int PartitionsK,
+    /// Tile iterator reading and writing output tensors
+    typename OutputTileIterator_,
+    /// Fragment iterator selecting accumulators
+    typename AccumulatorFragmentIterator_,
+    /// Output operator
+    typename OutputOp_,
+    /// Number of interleaved k
+    int InterleavedK>
+class InterleavedEpilogue :
+  public EpilogueBaseStreamK<
+    Shape_,
+    PartitionsK,
+    WarpMmaOperator_,
+    AccumulatorFragmentIterator_>
+{
+public:
+
+  using BaseStreamK = EpilogueBaseStreamK<
+    Shape_,
+    PartitionsK,
+    WarpMmaOperator_,
+    AccumulatorFragmentIterator_>;
+
+  using Shape = Shape_;
+  using WarpMmaOperator = WarpMmaOperator_;
+  static int const kPartitionsK = PartitionsK;
+  using AccumulatorFragmentIterator = AccumulatorFragmentIterator_;
+  using OutputTileIterator = OutputTileIterator_;
+  using OutputOp = OutputOp_;
+
+  /// The complete warp-level accumulator tile
+  using AccumulatorTile = typename AccumulatorFragmentIterator::AccumulatorTile;
+
+  /// Fragment type used by the accumulator tile's fragment iterator
+  using AccumulatorFragment = typename AccumulatorFragmentIterator::Fragment;
+
+  /// Accumulator element
+  using ElementAccumulator = typename AccumulatorTile::Element;
+
+  /// Output element
+  using ElementOutput = typename OutputTileIterator::Element;
+
+  /// Output access size
+  static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess;
+
+  /// Tensor reference to destination tensor
+  using TensorRef = typename OutputTileIterator::TensorRef;
+
+  /// Tensor reference to sync tensor
+  using SyncTensorRef =
+      typename cutlass::TensorRef<int, cutlass::layout::PackedVectorLayout>;
+
+  /// Const tensor reference to source tensor
+  using ConstTensorRef = typename OutputTileIterator::ConstTensorRef;
+
+  /// Array type used to output
+  using OutputAccessType = Array<typename OutputTileIterator::Element,
+                                 OutputTileIterator::kElementsPerAccess>;
+
+  /// Array type used by output functor
+  using AccumulatorAccessType =
+      Array<ElementAccumulator, OutputTileIterator::kElementsPerAccess>;
+
+  /// Number of warps
+  using WarpCount =
+      gemm::GemmShape<Shape::kM / WarpMmaOperator::Shape::kM,
+                      Shape::kN / WarpMmaOperator::Shape::kN, kPartitionsK>;
+
+public:
+
+  static_assert(OutputTileIterator::kElementsPerAccess,
+                "This must not be zero.");
+
+  static_assert(!(OutputTileIterator::Fragment::kElements %
+                  OutputTileIterator::kElementsPerAccess),
+                "Divisibility");
+
+public:
+
+  /// Aspect for when epilogue source is not needed
+  struct SourceAspectNotNeeded
+  {
+    /// Constructor
+    CUTLASS_DEVICE
+    SourceAspectNotNeeded()
+    {}
+
+    /// Invoke the output functor over each vector of output
+    CUTLASS_DEVICE
+    void apply_output_operator(
+      typename OutputTileIterator::Fragment &output_fragment,
+      OutputOp const &output_op,
+      typename AccumulatorFragmentIterator::Fragment const &aligned_accum_fragment)
+    {
+      OutputAccessType *output_frag_ptr =
+        reinterpret_cast<OutputAccessType *>(&output_fragment);
+
+      AccumulatorAccessType const *compute_frag_ptr =
+        reinterpret_cast<AccumulatorAccessType const *>(&aligned_accum_fragment);
+
+      int const kOutputOpIterations =
+        OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess;
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kOutputOpIterations; ++i)
+      {
+        // Call the output operator
+        output_frag_ptr[i] = output_op(compute_frag_ptr[i]);
+      }
+    }
+  };
+
+
+  /// Aspect for when epilogue source is needed
+  struct SourceAspectNeeded
+  {
+    OutputTileIterator source_iterator;
+
+    typename OutputTileIterator::Fragment source_fragment;
+
+    /// Invoke the output functor over each vector of output
+    CUTLASS_DEVICE
+    static void apply_output_operator(
+      typename OutputTileIterator::Fragment &output_fragment,
+      OutputOp const &output_op,
+      typename AccumulatorFragmentIterator::Fragment const &aligned_accum_fragment,
+      typename OutputTileIterator::Fragment const &source_fragment)
+    {
+      OutputAccessType *output_frag_ptr =
+        reinterpret_cast<OutputAccessType *>(&output_fragment);
+
+      AccumulatorAccessType const *compute_frag_ptr =
+        reinterpret_cast<AccumulatorAccessType const *>(&aligned_accum_fragment);
+
+      OutputAccessType const *source_frag_ptr =
+        reinterpret_cast<OutputAccessType const *>(&source_fragment);
+
+      int const kOutputOpIterations =
+        OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess;
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kOutputOpIterations; ++i)
+      {
+        // Call the output operator
+        output_frag_ptr[i] = output_op(compute_frag_ptr[i], source_frag_ptr[i]);
+      }
+    }
+
+    /// Constructor
+    CUTLASS_DEVICE
+    SourceAspectNeeded(OutputTileIterator source_iterator) :
+      source_iterator(source_iterator)
+    {
+      source_fragment.clear();
+    }
+
+    /// Invoke the output functor over each vector of output
+    CUTLASS_DEVICE
+    void apply_output_operator(
+      typename OutputTileIterator::Fragment &output_fragment,
+      OutputOp const &output_op,
+      typename AccumulatorFragmentIterator::Fragment const &aligned_accum_fragment)
+    {
+      // Load addend source fragment from global memory
+      source_iterator.load(source_fragment);
+      ++source_iterator;
+
+      apply_output_operator(output_fragment, output_op, aligned_accum_fragment, source_fragment);
+    }
+  };
+
+
+  /// Shared storage allocation needed by the epilogue
+  struct SharedStorage {};
+
+
+public:
+
+  /// Constructor
+  CUTLASS_DEVICE
+  InterleavedEpilogue(
+      SharedStorage &shared_storage,  ///< Shared storage object
+      int thread_idx,                 ///< ID of a thread within the threadblock
+      int warp_idx,                   ///< ID of warp within threadblock
+      int lane_idx)                   ///< Id of thread within warp
+  :
+      BaseStreamK(thread_idx)
+  {}
+
+
+  /// Aggregates the accumulator sets shared by peer blocks in the global workspace,
+  /// performing epilogue computations, writing to output
+  CUTLASS_DEVICE
+  void reduce(
+      int peer_idx_begin,
+      int peer_idx_end,
+      int reduce_fragment_idx,
+      void *element_workspace,
+      OutputOp const &output_op,                      ///< Output operator
+      OutputTileIterator destination_iterator,        ///< Tile iterator for destination
+      OutputTileIterator source_iterator)             ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
+  {
+    // Redcuce peer accumulator fragments into one fragment
+    AccumulatorFragment accum_fragment;
+    BaseStreamK::reduce(accum_fragment, peer_idx_begin, peer_idx_end, reduce_fragment_idx, element_workspace);
+
+    // Source-fragment data (zero-initialized for scenarios where the
+    // output operator allows us to skip loading it from global input)
+    typename OutputTileIterator::Fragment source_fragment;
+    source_fragment.clear();
+
+    if (output_op.is_source_needed())
+    {
+      source_iterator += reduce_fragment_idx;
+      source_iterator.load(source_fragment);
+    }
+
+    // Compute the output result
+    typename OutputTileIterator::Fragment output_fragment;
+
+    // Apply the output operator
+    SourceAspectNeeded::apply_output_operator(output_fragment, output_op, accum_fragment, source_fragment);
+
+    // Store the final result
+    destination_iterator += reduce_fragment_idx;
+    destination_iterator.store(output_fragment);
+  }
+
+
+  /// Perform the epilogue computations and stream the result to global memory.
+  CUTLASS_DEVICE
+  void operator()(
+    OutputOp const &output_op,                      ///< Output operator
+    OutputTileIterator destination_iterator,        ///< Tile iterator for destination
+    AccumulatorTile const &accumulators)            ///< Complete warp-level accumulator tile
+  {
+    operator()(output_op, destination_iterator, accumulators, SourceAspectNotNeeded());
+  }
+
+
+  /// Perform the epilogue computations and stream the result to global memory.  Implements
+  /// two alternative codepaths, depending on whether the output op requires addend data to be loaded.
+  CUTLASS_DEVICE
+  void operator()(
+    OutputOp const &output_op,                      ///< Output operator
+    OutputTileIterator destination_iterator,        ///< Tile iterator for destination
+    AccumulatorTile const &accumulators,            ///< Complete warp-level accumulator tile
+    OutputTileIterator source_iterator )            ///< Tile iterator for addend source
+  {
+    if (output_op.is_source_needed())
+    {
+      operator()(output_op, destination_iterator, accumulators, SourceAspectNeeded(source_iterator));
+    }
+    else
+    {
+      operator()(output_op, destination_iterator, accumulators, SourceAspectNotNeeded());
+    }
+  }
+
+
+  /// Perform the epilogue computations and stream the result to global memory.  Implements a
+  /// single codepath, regardless of whether the output op requires addend data to be loaded
+  CUTLASS_DEVICE
+  void unified(
+    OutputOp const &output_op,                      ///< Output operator
+    OutputTileIterator destination_iterator,        ///< Tile iterator for destination
+    AccumulatorTile const &accumulators,            ///< Complete warp-level accumulator tile
+    OutputTileIterator source_iterator )            ///< Tile iterator for addend source
+  {
+    if (!output_op.is_source_needed())
+    {
+      source_iterator.clear_mask();
+      __syncthreads();  // Dummy (CUDA 11.0)
+    }
+
+    operator()(output_op, destination_iterator, accumulators, SourceAspectNeeded(source_iterator));
+  }
+
+
+  /// Streams the result to global memory
+  template <typename SourceAspect>
+  CUTLASS_DEVICE
+  void operator()(
+    OutputOp const &output_op,                      ///< Output operator
+    OutputTileIterator destination_iterator,        ///< Tile iterator for destination
+    AccumulatorTile const &accumulators,            ///< Complete warp-level accumulator tile
+    SourceAspect source)
+  {
+    //
+    // Iterator over warp-level accumulator fragment
+    //
+
+    AccumulatorFragmentIterator accum_fragment_iterator(accumulators);
+
+    //
+    // Iterate over accumulator tile
+    //
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) {
+
+      //
+      // Convert fragment
+      //
+
+      typename AccumulatorFragmentIterator::Fragment accum_fragment;
+
+      accum_fragment_iterator.load(accum_fragment);
+      ++accum_fragment_iterator;
+
+      //
+      // Compute the output result
+      //
+
+      typename OutputTileIterator::Fragment output_fragment;
+      source.apply_output_operator(output_fragment, output_op, accum_fragment);
+
+      //
+      // Store the final result
+      //
+
+      destination_iterator.set_iteration_index(iter);
+      destination_iterator.store(output_fragment);
+      ++destination_iterator;
+    }
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/output_iterator_parameter.h

@@ -0,0 +1,223 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+
+#include "cutlass/conv/convolution.h"
+#include "cutlass/conv/conv2d_problem_size.h"
+#include "cutlass/conv/conv3d_problem_size.h"
+#include "cutlass/layout/tensor.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/tensor_ref.h"
+
+namespace cutlass {
+namespace epilogue {
+namespace threadblock {
+
+template<
+  typename TensorLayout_,                             ///! The original output tensor layout
+  typename OutputIteratorLayout_,                     ///! Layout used by epilogue output iterator
+  typename TensorRef_,                                ///! Input tensor to epilogue output iterator
+  conv::Operator ConvOperator,                        ///! Convolutional operator (Fprop, Dgrad, Wgrad)
+  typename ConvProblemSize_                          ///! Convolutional operator on 2D or 3D problem
+>
+struct ConvOutputIteratorParameter {
+
+  using TensorLayout = TensorLayout_;
+  using OutputIteratorLayout = OutputIteratorLayout_;
+  using OutputTensorCoord = typename OutputIteratorLayout::TensorCoord;
+  using TensorRef = TensorRef_;
+  static conv::Operator const kConvolutionalOperator = ConvOperator;
+  using ConvProblemSize = ConvProblemSize_;
+
+  /// Wgrad stride idx for implicit gemm algorithm 
+  // Conv2d row-major matrix (KxRSC) 
+  // Conv3d row-major matrix (KxTRSC)
+  static int const kWgradStrideIdx = 
+    platform::is_same<TensorLayout, layout::TensorNHWC>::value ? 2 : 3;
+
+  /// This chooses the appropriate stride element of the C tensor.
+  static int const kTensorStrideIdx = 
+    (kConvolutionalOperator == conv::Operator::kWgrad ? kWgradStrideIdx : 0);
+
+  CUTLASS_HOST_DEVICE
+  static OutputIteratorLayout layout(const TensorRef & ref) {
+    return ref.stride(kTensorStrideIdx);
+  }
+
+  CUTLASS_HOST_DEVICE
+  static OutputTensorCoord extent(ConvProblemSize problem_size) {
+    return conv::implicit_gemm_problem_size(kConvolutionalOperator, problem_size).mn();
+  }
+};
+
+template<
+  typename TensorRef_,                                ///! Input tensor to epilogue output iterator
+  typename ConvProblemSize_                          ///! Convolutional operator on 2D or 3D problem
+>
+struct ConvOutputIteratorParameter<layout::TensorNHWC, layout::TensorNHWC, TensorRef_, conv::Operator::kFprop, ConvProblemSize_> {
+
+  using TensorLayout = layout::TensorNHWC;
+  using OutputIteratorLayout = layout::TensorNHWC;
+  using MappedLayout = layout::RowMajor;
+  using OutputTensorCoord = typename OutputIteratorLayout::TensorCoord;
+  using MappedTensorCoord = typename MappedLayout::TensorCoord;
+  using TensorRef = TensorRef_;
+  static conv::Operator const kConvolutionalOperator = conv::Operator::kFprop;
+  using ConvProblemSize = ConvProblemSize_;
+
+  CUTLASS_HOST_DEVICE
+  static OutputIteratorLayout layout(const TensorRef & ref) {
+    return ref.stride();
+  }
+
+  CUTLASS_HOST_DEVICE
+  static MappedTensorCoord extent(ConvProblemSize problem_size) {
+    return conv::implicit_gemm_problem_size(kConvolutionalOperator, problem_size).mn();
+  }
+};
+
+template<
+  typename TensorRef_,                                ///! Input tensor to epilogue output iterator
+  typename ConvProblemSize_                          ///! Convolutional operator on 2D or 3D problem
+>
+struct ConvOutputIteratorParameter<layout::TensorNHWC, layout::TensorNHWC, TensorRef_, conv::Operator::kDeconv, ConvProblemSize_> {
+
+  using TensorLayout = layout::TensorNHWC;
+  using OutputIteratorLayout = layout::TensorNHWC;
+  using MappedLayout = layout::RowMajor;
+  using OutputTensorCoord = typename OutputIteratorLayout::TensorCoord;
+  using MappedTensorCoord = typename MappedLayout::TensorCoord;
+  using TensorRef = TensorRef_;
+  static conv::Operator const kConvolutionalOperator = conv::Operator::kDeconv;
+  using ConvProblemSize = ConvProblemSize_;
+
+  CUTLASS_HOST_DEVICE
+  static OutputIteratorLayout layout(const TensorRef & ref) {
+    return ref.stride();
+  }
+
+  CUTLASS_HOST_DEVICE
+  static MappedTensorCoord extent(ConvProblemSize problem_size) {
+    return conv::implicit_gemm_problem_size(kConvolutionalOperator, problem_size).mn();
+  }
+};
+
+template<
+  typename TensorRef_,                                ///! Input tensor to epilogue output iterator
+  typename ConvProblemSize_                          ///! Convolutional operator on 2D or 3D problem
+>
+struct ConvOutputIteratorParameter<layout::TensorNDHWC, layout::TensorNDHWC, TensorRef_, conv::Operator::kFprop, ConvProblemSize_> {
+
+  using TensorLayout = layout::TensorNDHWC;
+  using OutputIteratorLayout = layout::TensorNDHWC;
+  using MappedLayout = layout::RowMajor;
+  using OutputTensorCoord = typename OutputIteratorLayout::TensorCoord;
+  using MappedTensorCoord = typename MappedLayout::TensorCoord;
+  using TensorRef = TensorRef_;
+  static conv::Operator const kConvolutionalOperator = conv::Operator::kFprop;
+  using ConvProblemSize = ConvProblemSize_;
+
+  CUTLASS_HOST_DEVICE
+  static OutputIteratorLayout layout(const TensorRef & ref) {
+    return ref.stride();
+  }
+
+  CUTLASS_HOST_DEVICE
+  static MappedTensorCoord extent(ConvProblemSize problem_size) {
+    return conv::implicit_gemm_problem_size(kConvolutionalOperator, problem_size).mn();
+  }
+};
+
+template<
+  typename TensorRef_,                                ///! Input tensor to epilogue output iterator
+  typename ConvProblemSize_                          ///! Convolutional operator on 2D or 3D problem
+>
+struct ConvOutputIteratorParameter<layout::TensorNDHWC, layout::TensorNDHWC, TensorRef_, conv::Operator::kDeconv, ConvProblemSize_> {
+
+  using TensorLayout = layout::TensorNDHWC;
+  using OutputIteratorLayout = layout::TensorNDHWC;
+  using MappedLayout = layout::RowMajor;
+  using OutputTensorCoord = typename OutputIteratorLayout::TensorCoord;
+  using MappedTensorCoord = typename MappedLayout::TensorCoord;
+  using TensorRef = TensorRef_;
+  static conv::Operator const kConvolutionalOperator = conv::Operator::kDeconv;
+  using ConvProblemSize = ConvProblemSize_;
+
+  CUTLASS_HOST_DEVICE
+  static OutputIteratorLayout layout(const TensorRef & ref) {
+    return ref.stride();
+  }
+
+  CUTLASS_HOST_DEVICE
+  static MappedTensorCoord extent(ConvProblemSize problem_size) {
+    return conv::implicit_gemm_problem_size(kConvolutionalOperator, problem_size).mn();
+  }
+};
+
+template <
+  int InterleavedK,
+  typename TensorRef_,
+  conv::Operator ConvOperator,
+  typename ConvProblemSize_
+>
+struct ConvOutputIteratorParameter<
+  layout::TensorNCxHWx<InterleavedK>, 
+  layout::TensorNCxHWx<InterleavedK>,
+  TensorRef_,
+  ConvOperator,
+  ConvProblemSize_>
+{ 
+
+  using TensorLayout = typename layout::TensorNCxHWx<InterleavedK>;
+  using OutputIteratorLayout = typename layout::TensorNCxHWx<InterleavedK>;
+  using OutputTensorCoord = typename OutputIteratorLayout::TensorCoord;
+  using TensorRef = TensorRef_;
+  static conv::Operator const kConvolutionalOperator = ConvOperator;
+  using ConvProblemSize = ConvProblemSize_;
+
+  CUTLASS_HOST_DEVICE
+  static OutputIteratorLayout layout(const TensorRef & ref) {
+    return ref.stride();
+  }
+
+  CUTLASS_HOST_DEVICE
+  static OutputTensorCoord extent(ConvProblemSize problem_size) {
+    return problem_size.output_extent();
+  }
+
+};
+
+} // namespace threadblock
+} // namespace epilogue
+} // namespace cutlass

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/output_tile_thread_map.h

@@ -0,0 +1,628 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+  \brief Metaprogram for determining the mapping of output elements to threads for epilogue tiles.
+
+  
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/matrix_shape.h"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/fast_math.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace threadblock {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Tuple defining point in output tile
+template <
+  int Column,
+  int Row,
+  int Group,
+  int Cluster,
+  int Tile
+>
+struct OutputTileShape {
+  static int const kColumn = Column;
+  static int const kRow = Row;
+  static int const kGroup = Group;
+  static int const kCluster = Cluster;
+  static int const kTile = Tile;
+
+  static int const kCount = kColumn * kRow * kGroup * kCluster * kTile;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <typename Iterations, typename Delta>
+struct OutputTileThreadMapHelpers {
+
+  /// Determines the iteration index of a vector access according to the thread map
+  CUTLASS_HOST_DEVICE
+  static void iteration_index(
+    int &column_idx,
+    int &row_idx,
+    int &group_idx,
+    int &cluster_idx,
+    int &tile_idx,
+    int iter_idx) {
+
+    column_idx = iter_idx % Iterations::kColumn;
+    int residual   = iter_idx / Iterations::kColumn;
+
+    row_idx    = residual % Iterations::kRow;
+    residual       = residual / Iterations::kRow;
+
+    group_idx  = residual % Iterations::kGroup;
+    residual       = residual / Iterations::kGroup;
+
+    cluster_idx = residual % Iterations::kCluster;
+    tile_idx    = residual / Iterations::kCluster;
+  }
+
+  /// Computes the offset of a given vector access
+  CUTLASS_HOST_DEVICE
+  static MatrixCoord iteration_offset(int iter_idx) {
+
+    int column_idx;
+    int row_idx;
+    int group_idx;
+    int cluster_idx;
+    int tile_idx;
+
+    iteration_index(column_idx, row_idx, group_idx, cluster_idx, tile_idx, iter_idx);
+
+    return
+      MatrixCoord(
+        row_idx     * Delta::kRow     +
+        group_idx   * Delta::kGroup   +
+        cluster_idx * Delta::kCluster +
+        tile_idx    * Delta::kTile,
+
+        column_idx  * Delta::kColumn);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+template <
+  typename ThreadMap_,
+  typename Shape_,
+  typename Iterations_,
+  typename Delta_,
+  typename Count_
+>
+struct OutputTileThreadMap : public OutputTileThreadMapHelpers<Iterations_, Delta_> {
+
+  /// Conventional thread map (concept: ThreadMap)
+  using ThreadMap = ThreadMap_;
+
+  /// Number of threads participating in the operation
+  static int const kThreads = ThreadMap::kThreads;
+
+  /// Number of scalar elements per access
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+
+  /// Shape of the tile
+  using Shape = Shape_;
+
+  /// Iterations performed by each thread
+  using Iterations = Iterations_;
+
+  /// Delta between accesses
+  using Delta = Delta_;
+
+  /// Number of iterator iterations 
+  using Count = Count_;
+
+  /// Initial offset function
+  CUTLASS_HOST_DEVICE
+  static MatrixCoord initial_offset(int thread_idx) {
+
+    using Index = typename layout::PitchLinearCoord::Index;
+    
+    layout::PitchLinearCoord coord = ThreadMap::initial_offset(thread_idx);
+
+    Index cluster = coord.strided() / (Shape::kGroup * Shape::kRow);
+    Index cluster_residual = coord.strided() % (Shape::kGroup * Shape::kRow);
+
+    Index group = cluster_residual / (Shape::kRow);
+    Index row = cluster_residual % (Shape::kRow);
+
+    return MatrixCoord{
+      row + group * Shape::kRow * Count::kRow 
+        + cluster * Shape::kGroup * Count::kGroup * Shape::kRow * Count::kRow,
+      coord.contiguous()
+    };
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+/// RowArrangement determines how one or more warps cover a region of consecutive rows.
+template <
+  typename Shape,
+  int WarpsRemaining,
+  int ElementsPerAccess,
+  int ElementSize,
+  bool Is2dTile
+>
+struct RowArrangement;
+
+/// RowArrangement in which each warp's access is a 1D tiled arrangement.
+template <
+  typename Shape,
+  int WarpsRemaining,
+  int ElementsPerAccess,
+  int ElementSize
+>
+struct RowArrangement<Shape, WarpsRemaining, ElementsPerAccess, ElementSize, false> {
+  static int const kWarpSize = 32;
+  static int const kElementsPerAccess = ElementsPerAccess;
+  static int const kElementSize = ElementSize;
+
+  static int const kIterationsRow = 1;
+  static int const kDeltaRow = 1;
+  static int const kIterationsColumn = Shape::kColumn / kElementsPerAccess / kWarpSize;
+  static int const kDeltaColumn = kWarpSize * kElementsPerAccess;
+
+  static int const kAccessWidth = kWarpSize;
+  static int const kAccessRows = 1;
+  static int const kWarpPartitionsRow = 1;
+  static int const kWarpPartitionsColumn = WarpsRemaining;
+};
+
+/// RowArrangement in which each warp's access is a 2D tiled arrangement.
+template <
+  typename Shape,
+  int WarpsRemaining,
+  int ElementsPerAccess,
+  int ElementSize
+>
+struct RowArrangement<Shape, WarpsRemaining, ElementsPerAccess, ElementSize, true> {
+
+  static int const kMemoryAccessSize = 256; // Preferred access size
+  static int const kWarpSize = 32;
+
+  static int const kElementsPerAccess = ElementsPerAccess;
+  static int const kElementSize = ElementSize;
+
+  struct Detail {
+    static int const kShapeRow = Shape::kRow / WarpsRemaining;
+    static int const kShapeWidth = Shape::kColumn / kElementsPerAccess;
+
+    static int const kTargetMemoryAccessWidth = 
+      kMemoryAccessSize / (kElementsPerAccess * kElementSize / 8);
+
+    static int const kTargetAccessRows = kWarpSize / kTargetMemoryAccessWidth;
+  };
+
+  static int const kAccessWidth = 
+    (Detail::kTargetAccessRows > Detail::kShapeRow ?
+      kWarpSize / Detail::kShapeRow
+      : const_min(
+          Detail::kShapeWidth,
+        const_min(kWarpSize, kMemoryAccessSize / (kElementsPerAccess * kElementSize / 8))
+        ));
+
+  static int const kAccessRows =
+    (Detail::kTargetAccessRows > Detail::kShapeRow ?
+      Detail::kShapeRow
+      : const_min(Shape::kRow, kWarpSize / kAccessWidth));
+
+  static int const kIterationsRow = Detail::kShapeRow / kAccessRows;
+  static int const kDeltaRow = kAccessRows;
+
+  static int const kIterationsColumn = Detail::kShapeWidth / kAccessWidth;
+  static int const kDeltaColumn = kAccessWidth * kElementsPerAccess;
+
+  static_assert( kAccessWidth * kElementsPerAccess <= Shape::kColumn, "Accessing too many elements per access");
+  static_assert( kIterationsColumn > 0, "Iteration Count Column must be > 0" );
+  static_assert( kIterationsRow > 0, "Iteration Count Row must be > 0" );
+
+  static int const kWarpPartitionsRow = 1;
+  static int const kWarpPartitionsColumn = 1;
+};
+
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Template metaprogram for partitioning a 4D space across warps to achieve several performance
+/// objectives:
+///
+///   - coalesced memory accesses in units of 128 Byte lines
+///   - minimal address arithmetic
+///   - minimal predicate calculations
+///
+template <
+  typename Shape_,
+  typename Count_,
+  int Threads,
+  int ElementsPerAccess,
+  int ElementSize
+>
+struct OutputTileOptimalThreadMap {
+
+  using Shape = Shape_;
+  using Count = Count_;
+
+  static int const kWarpSize = 32;
+  static int const kThreads = Threads;
+  static int const kWarpCount = kThreads / kWarpSize;
+
+  static int const kElementsPerAccess = ElementsPerAccess;
+  static int const kElementSize = ElementSize;
+
+  //
+  // Metaprogram computation
+  //
+
+  struct Detail {
+
+    // Clusters
+    static int const kIterationsCluster = 
+      ((Shape::kCluster > kWarpCount) ?
+        Shape::kCluster / kWarpCount
+        : 1);
+
+    static int const kDeltaCluster =
+      ((Shape::kCluster > kWarpCount) ?
+        Shape::kRow * Count::kRow * Shape::kGroup * Count::kGroup * Shape::kCluster / kIterationsCluster
+        : 1);
+
+    static int const kCompactedDeltaCluster =
+      ((Shape::kCluster > kWarpCount) ?
+        Shape::kRow * Shape::kGroup * Shape::kCluster / kIterationsCluster
+        : 1);
+
+    static int const kWarpPartitionsCluster =
+      ((Shape::kCluster > kWarpCount) ?
+        kWarpCount
+        : kWarpCount / Shape::kCluster);
+
+    static int const kWarpsRemainingForGroups =
+      ((Shape::kCluster > kWarpCount) ? 1 : kWarpCount / Shape::kCluster);
+
+    // Groups
+    static int const kIterationsGroup =
+      ((Shape::kGroup > kWarpsRemainingForGroups) ?
+        Shape::kGroup / kWarpsRemainingForGroups
+        : 1);
+
+    static int const kDeltaGroup =
+      ((Shape::kGroup > kWarpsRemainingForGroups) ?
+        Shape::kRow * Count::kRow * Shape::kGroup / kIterationsGroup
+        : 1);
+
+    static int const kCompactedDeltaGroup =
+      ((Shape::kGroup > kWarpsRemainingForGroups) ?
+        Shape::kRow * Shape::kGroup / kIterationsGroup
+        : 1);
+
+    static int const kWarpPartitionsGroup =
+      ((Shape::kGroup > kWarpsRemainingForGroups) ?
+        1
+        : kWarpsRemainingForGroups / Shape::kGroup);
+
+    static int const kWarpsRemainingForRows =
+      ((Shape::kGroup > kWarpsRemainingForGroups) ?
+        1
+        : kWarpsRemainingForGroups / Shape::kGroup);
+    
+    // Rows
+    using RowArrangement = detail::RowArrangement<
+      Shape,
+      kWarpsRemainingForRows,
+      kElementsPerAccess,
+      kElementSize,
+      (Shape::kRow > kWarpsRemainingForRows)
+    >;
+
+    // Warp partitions
+    using WarpPartitions = OutputTileShape<
+      RowArrangement::kWarpPartitionsColumn,
+      RowArrangement::kWarpPartitionsRow,
+      kWarpPartitionsGroup,
+      kWarpPartitionsCluster,
+      1>;
+
+    static int const kAccessWidth = RowArrangement::kAccessWidth;
+    static int const kAccessRows = RowArrangement::kAccessRows;
+  };
+
+  //
+  // Output
+  //
+
+  using Iterations = OutputTileShape<
+    Detail::RowArrangement::kIterationsColumn, 
+    Detail::RowArrangement::kIterationsRow, 
+    Detail::kIterationsGroup, 
+    Detail::kIterationsCluster, 
+    1>;
+
+  using Delta = OutputTileShape<
+    Detail::RowArrangement::kDeltaColumn,
+    Detail::RowArrangement::kDeltaRow,
+    Detail::kDeltaGroup,
+    Detail::kDeltaCluster,
+    1>;
+
+  /// Initial offset function
+  CUTLASS_HOST_DEVICE
+  static MatrixCoord initial_offset(int thread_idx) {
+
+//    int warp_idx = __shfl_sync(0xffffffff, thread_idx / kWarpSize, 0);
+    int warp_idx = thread_idx / kWarpSize;
+    int lane_idx = thread_idx % kWarpSize;
+
+    // Compute warp location
+    int cluster_idx = warp_idx / Detail::WarpPartitions::kCluster;
+    int residual_cluster = warp_idx % Detail::WarpPartitions::kCluster;
+
+    int group_idx = residual_cluster / Detail::WarpPartitions::kGroup;
+    int residual_group = residual_cluster % Detail::WarpPartitions::kGroup;
+
+    int row_idx = residual_group / Detail::WarpPartitions::kRow;
+    int col_idx = residual_group % Detail::WarpPartitions::kRow;
+
+    // Compute per-lane offset
+    int lane_row_offset = lane_idx / Detail::kAccessWidth;
+    int lane_col_offset = lane_idx % Detail::kAccessWidth;
+
+    // Compute coordinate in output space
+    int cluster_offset = cluster_idx * Shape::kRow * Count::kRow * Shape::kGroup * Count::kGroup;
+    int group_offset = group_idx * Shape::kRow * Count::kRow;
+    int row_offset = row_idx * Iterations::kRow * Detail::kAccessRows;
+    int column_offset = col_idx * Iterations::kColumn * Detail::kAccessWidth * kElementsPerAccess;
+
+    return MatrixCoord(
+      cluster_offset + group_offset + row_offset + lane_row_offset,
+      column_offset + lane_col_offset * kElementsPerAccess
+    );
+  }
+
+  /// Computes the offset of a given vector access
+  CUTLASS_HOST_DEVICE
+  static MatrixCoord iteration_offset(int iter_idx) {
+    return OutputTileThreadMapHelpers<Iterations, Delta>::iteration_offset(iter_idx);
+  }
+
+  /// Compacted thread map in which the 4D region is contiguous
+  struct CompactedThreadMap {
+
+
+    using Shape = Shape_;
+
+    using TileShape = MatrixShape<
+      Shape::kTile * Shape::kCluster * Shape::kGroup * Shape::kRow,
+      Shape::kColumn
+    >;
+
+    using Iterations = OutputTileShape<
+      Detail::RowArrangement::kIterationsColumn,
+      Detail::RowArrangement::kIterationsRow,
+      Detail::kIterationsGroup,
+      Detail::kIterationsCluster,
+      1>;
+
+    using Delta = OutputTileShape<
+      Detail::RowArrangement::kDeltaColumn,
+      Detail::RowArrangement::kDeltaRow,
+      Detail::kCompactedDeltaGroup,
+      Detail::kCompactedDeltaCluster,
+      1>;
+
+    /// Number of elements within each vector access
+    static int const kElementsPerAccess = ElementsPerAccess;
+
+    /// Number  of threads
+    static int const kThreads = Threads;
+
+    /// Function to compute each thread's initial offset
+    CUTLASS_HOST_DEVICE
+    static MatrixCoord initial_offset(int thread_idx) {
+
+//      int warp_idx = __shfl_sync(0xffffffff, thread_idx / kWarpSize, 0);
+      int warp_idx = thread_idx / kWarpSize;
+      int lane_idx = thread_idx % kWarpSize;
+
+      // Compute warp location
+      int cluster_idx = warp_idx / Detail::WarpPartitions::kCluster;
+      int residual_cluster = warp_idx % Detail::WarpPartitions::kCluster;
+
+      int group_idx = residual_cluster / Detail::WarpPartitions::kGroup;
+      int residual_group = residual_cluster % Detail::WarpPartitions::kGroup;
+
+      int row_idx = residual_group / Detail::WarpPartitions::kRow;
+      int col_idx = residual_group % Detail::WarpPartitions::kRow;
+
+      // Compute per-lane offset
+      int lane_row_offset = lane_idx / Detail::kAccessWidth;
+      int lane_col_offset = lane_idx % Detail::kAccessWidth;
+
+      // Compute coordinate in output space
+      int cluster_offset = cluster_idx * Shape::kRow * Shape::kGroup;
+      int group_offset = group_idx * Shape::kRow;
+      int row_offset = row_idx * Iterations::kRow * Detail::kAccessRows;
+      int column_offset = col_idx * Iterations::kColumn * Detail::kAccessWidth * kElementsPerAccess;
+
+      MatrixCoord coord(
+        cluster_offset + group_offset + row_offset + lane_row_offset,
+        column_offset + lane_col_offset * kElementsPerAccess
+      );
+
+      return coord;
+    }
+  };
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Template metaprogram for partitioning a 3D interleaved layout across warps
+/// to achieve several performance objectives:
+///
+///   - coalesced memory accesses in units of 64 Byte lines
+///   - minimal address arithmetic
+///   - minimal predicate calculations
+///
+template <typename WarpCount_, typename Iterations_, int Threads,
+          int ElementsPerAccess, int ElementSize>
+struct InterleavedOutputTileThreadMap {
+  using WarpCount = WarpCount_;
+
+  static int const kWarpSize = 32;
+  static int const kThreads = Threads;
+  static int const kWarpCount = kThreads / kWarpSize;
+
+  static int const kElementsPerAccess = ElementsPerAccess;
+  static int const kElementSize = ElementSize;
+
+  //
+  // Metaprogram computation
+  //
+
+  struct Detail {};
+
+  //
+  // Output
+  //
+
+  using Iterations = Iterations_;
+
+  using Delta = layout::PitchLinearShape<kWarpSize * kElementsPerAccess, 1>;
+
+  /// Initial offset function
+  CUTLASS_HOST_DEVICE
+  static layout::PitchLinearCoord initial_offset(int thread_idx) {
+    int warp_idx = thread_idx / kWarpSize;
+    int lane_idx = thread_idx % kWarpSize;
+
+    // Compute warp location
+    layout::PitchLinearCoord warp_footprint{
+        Delta::kContiguous * Iterations::kContiguous,
+        Delta::kStrided * Iterations::kStrided};
+
+    layout::PitchLinearCoord warp_offset{warp_idx % WarpCount::kContiguous,
+                                         warp_idx / WarpCount::kContiguous};
+
+    // Compute per-lane offset
+    layout::PitchLinearCoord thread_offset_in_warp{
+        lane_idx * kElementsPerAccess, 0};
+
+    layout::PitchLinearCoord thread_offset_in_threadblock_tile =
+        warp_footprint * warp_offset + thread_offset_in_warp;
+
+    return thread_offset_in_threadblock_tile;
+  }
+};
+
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Template metaprogram for partitioning a 4D interleaved layout across warps
+/// to achieve several performance objectives:
+///
+///   - coalesced memory accesses in units of 64 Byte lines
+///   - minimal address arithmetic
+///   - minimal predicate calculations
+///
+template <typename WarpCount_, typename Iterations_, int Threads,
+          int ElementsPerAccess, int ElementSize>
+struct InterleavedConvOutputTileThreadMap {
+  using WarpCount = WarpCount_;
+
+  static int const kWarpSize = 32;
+  static int const kThreads = Threads;
+  static int const kWarpCount = kThreads / kWarpSize;
+
+  static int const kElementsPerAccess = ElementsPerAccess;
+  static int const kElementSize = ElementSize;
+
+  //
+  // Metaprogram computation
+  //
+
+  struct Detail {};
+
+  //
+  // Output
+  //
+
+  using Iterations = Iterations_;
+
+  using Delta = MatrixShape<kWarpSize / 4, 4 * kElementsPerAccess>;
+
+  /// Initial offset function
+  CUTLASS_HOST_DEVICE
+  static MatrixCoord initial_offset(int thread_idx) {
+    int warp_idx = thread_idx / kWarpSize;
+    int lane_idx = thread_idx % kWarpSize;
+
+    // Compute warp location
+    MatrixCoord warp_footprint{
+        Delta::kRow * Iterations::kRow,
+        Delta::kColumn * Iterations::kColumn,
+    };
+
+    MatrixCoord warp_offset{warp_idx % WarpCount::kRow,
+                            warp_idx / WarpCount::kRow};
+
+    // Compute per-lane offset
+    MatrixCoord thread_offset_in_warp{lane_idx / 4,
+                                      (lane_idx % 4) * kElementsPerAccess};
+
+    MatrixCoord thread_offset_in_threadblock_tile =
+        warp_footprint * warp_offset + thread_offset_in_warp;
+
+    return thread_offset_in_threadblock_tile;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace epilogue
+} // namespace cutlass

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator.h

@@ -0,0 +1,1387 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+  \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
+
+  The epilogue rearranges the result of a matrix product through shared memory to match canonical
+  tensor layouts in global memory. Epilogues support conversion and reduction operations.
+
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/layout/tensor.h"
+#include "cutlass/layout/permute.h"
+#include "cutlass/matrix_shape.h"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/transform/pitch_linear_thread_map.h"
+#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
+#include "cutlass/arch/arch.h"
+#include "cutlass/arch/memory.h"
+#include "cutlass/epilogue/threadblock/predicated_tile_iterator_params.h"
+#include "cutlass/conv/conv2d_problem_size.h"
+#include "cutlass/conv/conv3d_problem_size.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace epilogue {
+namespace threadblock {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Tile iterator used to load and store output tile from global memory in epilogue.
+///
+/// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator
+///
+template <
+  typename ThreadMap_,       ///< Thread map (conept: OutputTileThreadMap)
+  typename Element_,         ///< Element data type
+  bool ScatterD = false,     ///< Scatter D operand or not
+  typename PermuteDLayout = layout::NoPermute, ///< Permute D operand or not
+  bool UseCUDAStore = false
+>
+class PredicatedTileIterator {
+public:
+  using ThreadMap = ThreadMap_;
+  using Shape = typename ThreadMap::Shape;
+
+  using Element = Element_;
+
+  using Layout = layout::RowMajor;
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = MatrixCoord;
+
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+  static int const kThreads = ThreadMap::kThreads;
+  static int const kIterations = ThreadMap::Count::kTile;
+
+  static bool constexpr PermuteD = !layout::is_trivial_permute<PermuteDLayout>;
+
+  static_assert( ThreadMap::Iterations::kRow > 0,"ThreadMap::Iterations::kRow must be > 0");
+  static_assert( ThreadMap::Iterations::kGroup > 0,"ThreadMap::Iterations::kGroup must be > 0");
+  static_assert( ThreadMap::Iterations::kCluster > 0,"ThreadMap::Iterations::kCluster must be > 0");
+  static_assert( ThreadMap::Iterations::kColumn > 0,"ThreadMap::Iterations::kColumn must be > 0");
+
+  /// Fragment object
+  using Fragment = Array<
+    Element,
+    ThreadMap::Iterations::kColumn *
+    ThreadMap::Iterations::kRow *
+    ThreadMap::Iterations::kGroup *
+    ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>;
+
+  /// Memory access size
+  using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess>;
+
+  //
+  // Parameters struct
+  //
+
+  /// Uses a non-template class
+  struct Params : PredicatedTileIteratorParams {
+    using Base = PredicatedTileIteratorParams;
+
+    CUTLASS_HOST_DEVICE
+    Params() { }
+
+    CUTLASS_HOST_DEVICE
+    Params(Layout const &layout):
+      PredicatedTileIteratorParams(
+        layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess,
+        make_OutputTileThreadMapDesc<ThreadMap>()
+      ) 
+    { }
+
+    CUTLASS_HOST_DEVICE
+    Params(Layout const &layout,
+           // Not needed.  Added to be compatible with strided conv epilogue.
+           cutlass::Tensor4DCoord const &tensor_extent):
+      Params(layout)
+    { }
+
+    CUTLASS_HOST_DEVICE
+    Params(Layout const &layout,
+           // Not needed.  Added to be compatible with strided conv epilogue.
+           cutlass::Tensor5DCoord const &tensor_extent):
+      Params(layout)
+    { }
+
+    CUTLASS_HOST_DEVICE
+    Params(Base const &base) : 
+      Base(base) { }
+  };
+
+  /// Mask object
+  struct Mask {
+
+    static int const kCount = ThreadMap::Iterations::kColumn;
+
+    /// Predicate state
+    bool predicates[kCount];
+
+    //
+    // Mask
+    //
+    CUTLASS_HOST_DEVICE
+    Mask() {
+      enable();
+    }
+
+    ///< Efficiently disables all accesses guarded by mask
+    CUTLASS_HOST_DEVICE void clear() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = false;
+      }
+    }
+
+    ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask
+    CUTLASS_DEVICE void enable() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = true;
+      }
+    }
+  };
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Parameters structure containing reference and precomputed state.
+  PredicatedTileIteratorParams params_;
+
+  /// Byte-level pointer. This pointer is usually for both load() and store(), unless PermuteD is performed. When having PermuteD, byte_pointer_ is only for load().
+  uint8_t *byte_pointer_;
+
+  /// Byte-level pointer for store(). Due to PermuteD Op, store_byte_pointer_ may be with different address computation compared to byte_pointer_.
+  uint8_t *store_byte_pointer_;
+
+  /// Array of boolean values to contain steady-state predicates
+  Mask mask_;
+
+  /// Extent of the matrix tile in rows
+  Index extent_row_;
+
+  /// Extent of the matrix tile in rows
+  Index extent_column_;
+
+  /// A thread's starting row position (assuming steady-state predicates have been computed)
+  Index thread_start_row_;
+
+  /// A thread's starting column
+  Index thread_start_column_;
+
+  /// Internal state counter
+  int state_[3];
+
+  /// Scatter indices
+  int const *indices_;
+
+  /// PermuteDLayout
+  PermuteDLayout permute_layout_;
+
+  //
+  // Static asserts about internal strides
+  //
+
+  static_assert(sizeof(extent_row_) == 4, "Expected 32b extents");
+  static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents");
+  static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides");
+
+private:
+
+  //
+  // Methods
+  //
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  CUTLASS_DEVICE
+  PredicatedTileIterator(
+    PredicatedTileIteratorParams const & params,
+    Element *pointer,
+    TensorCoord extent,
+    int thread_idx,
+    TensorCoord threadblock_offset = TensorCoord(),
+    int const *indices = nullptr
+  ): 
+    params_(params), indices_(indices),
+    permute_layout_(PitchLinearCoord(extent.column(), extent.row()), params_.stride * kElementsPerAccess / sizeof(AccessType))
+  {
+
+    TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset;
+
+    extent_row_ = extent.row();
+    extent_column_ = extent.column();
+
+    thread_start_row_ = thread_offset.row();
+    thread_start_column_ = thread_offset.column();
+
+    // Initialize predicates
+    CUTLASS_PRAGMA_UNROLL
+    for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) {
+
+      mask_.predicates[c] = ((thread_offset.column()
+        + ThreadMap::Delta::kColumn * c) < extent.column());
+    }
+
+    // Null pointer performs no accesses
+    if (!pointer) {
+      mask_.clear();
+    }
+
+    if (ScatterD && !indices) {
+      mask_.clear();
+    }
+
+    // Initialize byte_pointer_
+    byte_pointer_ = reinterpret_cast<uint8_t *>(pointer) +
+      LongIndex(thread_offset.row()) * LongIndex(params_.stride) +
+      LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess;
+
+    if (ScatterD) {
+      byte_pointer_ = reinterpret_cast<uint8_t *>(pointer) +
+        LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess;
+    }
+
+    // store_byte_pointer_ is set to be the same with byte_pointer_ unless PermuteD is used.
+    store_byte_pointer_ = PermuteD ? reinterpret_cast<uint8_t *>(pointer) : byte_pointer_;
+
+    // Initialize internal state counter
+    state_[0] = state_[1] = state_[2] = 0;
+  }
+
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    store_byte_pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
+    byte_pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load_with_byte_offset(Fragment &frag, int64_t byte_offset) const {
+
+    uint8_t *byte_pointer = byte_pointer_;
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int frag_row_idx =
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          int row_offset = row * ThreadMap::Delta::kRow 
+            + group * ThreadMap::Delta::kGroup 
+            + cluster * ThreadMap::Delta::kCluster;
+
+          bool row_guard = ((row_offset + thread_start_row_) < extent_row_);
+
+          AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset);
+
+          if (ScatterD && row_guard) {
+            assert(indices_);
+
+            memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset +
+              LongIndex(indices_[row_offset + thread_start_row_]) * LongIndex(params_.stride));
+          }
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+
+            bool guard = row_guard && mask_.predicates[column];
+
+            cutlass::arch::global_load<
+              AccessType,
+              sizeof(AccessType)
+            >(
+                frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn +
+                         column],
+                (void *)&memory_pointer[column * ThreadMap::Delta::kColumn /
+                                        kElementsPerAccess],
+                guard);
+          }
+
+          if (row + 1 < ThreadMap::Iterations::kRow) {
+            if (!ScatterD) {
+              byte_pointer += params_.increment_row;
+            }
+          }
+        }
+
+        if (group + 1 < ThreadMap::Iterations::kGroup) {
+          byte_pointer += params_.increment_group;
+        }
+      }
+
+      if (cluster + 1 < ThreadMap::Iterations::kCluster) {
+        byte_pointer += params_.increment_cluster;
+      }
+    }
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load(Fragment &frag) const {
+
+    load_with_byte_offset(frag, 0);
+  }
+
+  /// Stores a fragment to memory
+  CUTLASS_DEVICE
+  void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) const {
+    uint8_t *byte_pointer = store_byte_pointer_;
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int frag_row_idx =
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          int row_offset = row * ThreadMap::Delta::kRow
+            + group * ThreadMap::Delta::kGroup
+            + cluster * ThreadMap::Delta::kCluster;
+
+          bool row_guard = ((row_offset + thread_start_row_) < extent_row_);
+
+          AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset);
+
+          if (ScatterD && row_guard) {
+            assert(indices_);
+
+            memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset +
+              LongIndex(indices_[row_offset + thread_start_row_]) * LongIndex(params_.stride));
+          }
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+
+            bool guard = row_guard && mask_.predicates[column];
+            
+            if (PermuteD) {
+
+              int col_offset = column * ThreadMap::Delta::kColumn;
+
+              int col = col_offset + thread_start_column_;
+              int row = row_offset + thread_start_row_;
+
+              // Locate memory_pointer
+              memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset
+                 + permute_layout_(PitchLinearCoord(col, row)) * sizeof(AccessType) / kElementsPerAccess);
+            }
+
+            if (UseCUDAStore) {
+              if (guard) {
+                memory_pointer[0] =
+                    frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column];
+              }
+            } else {
+              cutlass::arch::global_store<AccessType, sizeof(AccessType)>(
+                  frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column],
+                  (void *)&memory_pointer[0],
+                  guard);
+            }
+
+            if (!PermuteD) {
+              memory_pointer += (ThreadMap::Delta::kColumn / kElementsPerAccess);
+            }
+          }
+
+          if (row + 1 < ThreadMap::Iterations::kRow) {
+            if (!ScatterD && !PermuteD) {
+              byte_pointer += params_.increment_row;
+            }
+          }
+        }
+
+        if (group + 1 < ThreadMap::Iterations::kGroup) {
+          if (!ScatterD && !PermuteD) {
+            byte_pointer += params_.increment_group;
+          }
+        }
+      }
+
+      if (cluster + 1 < ThreadMap::Iterations::kCluster) {
+        if (!ScatterD && !PermuteD) {
+          byte_pointer += params_.increment_cluster;
+        }
+      }
+    }
+  }
+
+  /// Stores a fragment to memory
+  CUTLASS_DEVICE
+  void store(Fragment const &frag) const {
+
+    store_with_byte_offset(frag, 0);
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void downsample_load_with_byte_offset(Fragment &frag, int64_t byte_offset, int convolution_P, int convolution_Q, int add_P, int add_Q, int problem_N) const {
+
+    uint8_t *byte_pointer = byte_pointer_;
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int frag_row_idx = 
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          int row_offset = row * ThreadMap::Delta::kRow 
+            + group * ThreadMap::Delta::kGroup 
+            + cluster * ThreadMap::Delta::kCluster;
+
+          bool row_guard = ((row_offset + thread_start_row_) < extent_row_);
+
+          int output_row = row_offset + thread_start_row_;
+          int output_N = output_row / (convolution_P * convolution_Q);
+          int output_PQ = output_row % (convolution_P * convolution_Q);
+          int output_P = output_PQ / convolution_Q;
+          int output_Q = output_PQ % convolution_Q;
+
+          int input_row = output_N * 2 * convolution_P * 2 * convolution_Q +
+            (2 * output_P + add_P) * 2 * convolution_Q + 2 * output_Q + add_Q;
+
+          int64_t byte_offset = (input_row-output_row)*problem_N*sizeof(float);
+
+          AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+
+            bool guard = row_guard && mask_.predicates[column];
+
+            cutlass::arch::global_load<
+              AccessType, 
+              sizeof(AccessType)
+            >(
+                frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn +
+                         column],
+                (void *)&memory_pointer[column * ThreadMap::Delta::kColumn /
+                                        kElementsPerAccess],
+                guard);
+          }
+
+          if (row + 1 < ThreadMap::Iterations::kRow) {
+            byte_pointer += params_.increment_row;
+          }
+        }
+
+        if (group + 1 < ThreadMap::Iterations::kGroup) {
+          byte_pointer += params_.increment_group;
+        }
+      }
+
+      if (cluster + 1 < ThreadMap::Iterations::kCluster) {
+        byte_pointer += params_.increment_cluster;
+      }
+    }
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void upsample_load_with_byte_offset(Fragment &frag, int64_t byte_offset, int convolution_P, int convolution_Q, int add_P, int add_Q, int problem_N) const {
+
+    uint8_t *byte_pointer = byte_pointer_;
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int frag_row_idx = 
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          int row_offset = row * ThreadMap::Delta::kRow 
+            + group * ThreadMap::Delta::kGroup 
+            + cluster * ThreadMap::Delta::kCluster;
+
+          bool row_guard = ((row_offset + thread_start_row_) < extent_row_);
+
+          int output_row = row_offset + thread_start_row_;
+          int output_N = output_row / (convolution_P * convolution_Q);
+          int output_PQ = output_row % (convolution_P * convolution_Q);
+          int output_P = output_PQ / convolution_Q;
+          int output_Q = output_PQ % convolution_Q;
+          int row_add_P = add_P;
+          int row_add_Q = add_Q;
+	  if (output_P > convolution_P - 2) row_add_P = 0;
+	  if (output_Q > convolution_Q - 2) row_add_Q = 0;
+
+          int input_row = output_N * (convolution_P/2) * (convolution_Q/2) +
+            ((output_P + row_add_P)/2) * (convolution_Q/2) + (output_Q + row_add_Q)/2;
+
+          int64_t byte_offset = (input_row-output_row)*problem_N*sizeof(float);
+
+          AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+
+            bool guard = row_guard && mask_.predicates[column];
+
+            cutlass::arch::global_load<
+              AccessType, 
+              sizeof(AccessType)
+            >(
+                frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn +
+                         column],
+                (void *)&memory_pointer[column * ThreadMap::Delta::kColumn /
+                                        kElementsPerAccess],
+                guard);
+          }
+
+          if (row + 1 < ThreadMap::Iterations::kRow) {
+            byte_pointer += params_.increment_row;
+          }
+        }
+
+        if (group + 1 < ThreadMap::Iterations::kGroup) {
+          byte_pointer += params_.increment_group;
+        }
+      }
+
+      if (cluster + 1 < ThreadMap::Iterations::kCluster) {
+        byte_pointer += params_.increment_cluster;
+      }
+    }
+  }
+
+  CUTLASS_DEVICE
+  MatrixCoord thread_start() const {
+    return MatrixCoord(thread_start_row_, thread_start_column_);
+  }
+
+  /// Need to get the thread start row from the tile iterator
+  CUTLASS_DEVICE
+  int32_t thread_start_row() const {
+    return thread_start_row_;
+  }
+
+  /// Need to get the thread start row from the tile iterator
+  CUTLASS_DEVICE
+  int32_t thread_start_column() const {
+    return thread_start_column_;
+  }
+
+  /// Extent of the matrix in rows
+  CUTLASS_DEVICE
+  Index extent_row() const {
+    return extent_row_;
+  }
+
+  /// Extent of the matrix in columns
+  CUTLASS_DEVICE
+  Index extent_column() const {
+    return extent_column_;
+  }
+
+  /// Advances to the next position to load or store
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIterator &operator++() {
+
+    ++state_[0];
+
+    if (!ScatterD) {
+      byte_pointer_ += params_.advance_row;
+    }
+
+    if (!ScatterD && !PermuteD) {
+      store_byte_pointer_ += params_.advance_row;
+    }
+
+    thread_start_row_ += ThreadMap::Shape::kRow;
+
+    if (state_[0] == ThreadMap::Count::kRow) {
+
+      state_[0] = 0;
+      ++state_[1];
+
+      if (!ScatterD) {
+        byte_pointer_ += params_.advance_group;
+      }
+
+      if (!ScatterD && !PermuteD) {
+        store_byte_pointer_ += params_.advance_group;
+      }
+
+      thread_start_row_ += (ThreadMap::Shape::kGroup - 1) *
+        ThreadMap::Shape::kRow * ThreadMap::Count::kRow;
+
+      if (state_[1] == ThreadMap::Count::kGroup) {
+
+        state_[1] = 0;
+        ++state_[2];
+
+        if (!ScatterD) {
+          byte_pointer_ += params_.advance_cluster;
+        }
+
+        if (!ScatterD && !PermuteD) {
+          store_byte_pointer_ += params_.advance_cluster;
+        }
+
+        thread_start_row_ += ThreadMap::Count::kGroup *
+          ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow;
+
+        if (state_[2] == ThreadMap::Count::kCluster) {
+          state_[2] = 0;
+
+          if (!ScatterD) {
+            byte_pointer_ += params_.advance_tile;
+          }
+
+          if (!ScatterD && !PermuteD) {
+            store_byte_pointer_ += params_.advance_tile;
+          }
+
+          thread_start_row_ += ThreadMap::Shape::kGroup * ThreadMap::Shape::kRow
+            * ThreadMap::Shape::kCluster * ThreadMap::Shape::kTile;
+        }
+      }
+    }
+
+    return *this;
+  }
+
+  /// Advances a number of positions to load or store
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIterator &operator+=(int increment)
+  {
+    // Row
+    state_[0] += increment;
+    int increment_row = state_[0] / ThreadMap::Count::kRow;
+    state_[0] = state_[0] % ThreadMap::Count::kRow;
+
+    byte_pointer_ += (params_.advance_row * increment);
+    store_byte_pointer_ += (params_.advance_row * increment);
+    thread_start_row_ += (ThreadMap::Shape::kRow * increment);
+
+    // Group
+    state_[1] += increment_row;
+    int increment_group = state_[1] / ThreadMap::Count::kGroup;
+    state_[1] = state_[1] % ThreadMap::Count::kGroup;
+
+    byte_pointer_ += (params_.advance_group * increment_row);
+    store_byte_pointer_ += (params_.advance_group * increment_row);
+    thread_start_row_ +=
+        (ThreadMap::Shape::kGroup - 1) *
+        ThreadMap::Shape::kRow *
+        ThreadMap::Count::kRow *
+        increment_row;
+
+
+    // Cluster
+    state_[2] += increment_group;
+    int increment_cluster = state_[2] / ThreadMap::Count::kCluster;
+    state_[2] = state_[2] % ThreadMap::Count::kCluster;
+
+    byte_pointer_ += (params_.advance_cluster * increment_group);
+    store_byte_pointer_ += (params_.advance_cluster * increment_group);
+    thread_start_row_ +=
+        ThreadMap::Count::kGroup *
+        ThreadMap::Shape::kGroup *
+        ThreadMap::Count::kRow *
+        ThreadMap::Shape::kRow *
+        increment_group;
+
+    // Tile
+    byte_pointer_ += (params_.advance_tile * increment_cluster);
+    store_byte_pointer_ += (params_.advance_tile * increment_cluster);
+    thread_start_row_ +=
+        ThreadMap::Shape::kGroup *
+        ThreadMap::Shape::kRow *
+        ThreadMap::Shape::kCluster *
+        ThreadMap::Shape::kTile *
+        increment_cluster;
+
+    return *this;
+  }
+
+  ///< Efficiently disables all accesses guarded by mask
+  CUTLASS_DEVICE void clear_mask() {
+    mask_.clear();
+  }
+
+  ///< Efficiently enables all accesses guarded by mask
+  CUTLASS_DEVICE void enable_mask() {
+    mask_.enable();
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void get_mask(Mask &mask) const {
+    mask = mask_;
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void set_mask(Mask const &mask) {
+    mask_ = mask;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Tile iterator used to load output tile from global memory in epilogue.
+///
+/// Satisfies: ReadableTileIterator | InterleavedPredicatedTileIterator | ForwardTileIterator
+///
+template <
+  typename ThreadMap_,       ///< Thread map (conept: OutputTileThreadMap)
+  typename Element_,         ///< Element data type
+  int InterleavedN           ///< Number of Interleaved N 
+>
+class InterleavedPredicatedTileIterator {
+public:
+  using ThreadMap = ThreadMap_;
+
+  using Element = Element_;
+
+  using Layout = layout::ColumnMajorInterleaved<InterleavedN>;
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = layout::PitchLinearCoord;
+
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+  static int const kThreads = ThreadMap::kThreads;
+  static int const kIterations = ThreadMap::Iterations::kCount;
+
+  /// Fragment object
+  using Fragment = Array<Element, ThreadMap::kElementsPerAccess>;
+
+  /// Memory access size
+  using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess>;
+
+  /// Uses a non-template class
+  struct Params : InterleavedPredicatedTileIteratorParams {
+    using Base = InterleavedPredicatedTileIteratorParams;
+
+    CUTLASS_HOST_DEVICE
+    Params() { }
+
+    CUTLASS_HOST_DEVICE
+    Params(Layout const &layout): 
+      Base(
+        layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess,
+        make_InterleavedPredicatedTileIteratorDesc<Element, ThreadMap>()
+      ) { }
+
+    CUTLASS_HOST_DEVICE
+    Params(Base const &base) : 
+      Base(base) { }
+  };
+
+  /// Mask object
+  struct Mask {
+    static int const kCount = (ThreadMap::Iterations::kContiguous < 8)
+                                  ? 8
+                                  : ThreadMap::Iterations::kContiguous;
+
+    /// Predicate state
+    bool predicates[kCount];
+
+    //
+    // Mask
+    //
+    CUTLASS_HOST_DEVICE
+    Mask() {
+      enable();
+    }
+
+    ///< Efficiently disables all accesses guarded by mask
+    CUTLASS_HOST_DEVICE void clear() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = false;
+      }
+    }
+
+    ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask
+    CUTLASS_DEVICE void enable() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = true;
+      }
+    }
+  };
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Parameters structure containing reference and precomputed state.
+  Params params_;
+
+  /// Byte-level pointer
+  uint8_t *byte_pointer_;
+
+  /// Array of boolean values to contain steady-state predicates
+  Mask mask_;
+
+  /// Extent of the matrix tile in columns
+  Index extent_col_;
+
+  /// A thread's starting column position (assuming steady-state predicates have
+  /// been computed)
+  Index thread_start_col_;
+
+  /// Internal iteration counter
+  int iteration_contiguous_;
+
+  int iteration_strided_;
+
+private:
+
+  //
+  // Methods
+  //
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  CUTLASS_DEVICE
+  InterleavedPredicatedTileIterator(
+    Params const & params,
+    Element *pointer,
+    TensorCoord extent,
+    int thread_idx,
+    TensorCoord threadblock_offset,
+    int const *indices = nullptr     ///< gather/scatter indices, note no support for gather/scatter at this specialization
+  ):
+    params_(params) {
+    TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) +
+                                TensorCoord(threadblock_offset.contiguous() * InterleavedN,
+                                 threadblock_offset.strided() / InterleavedN);
+
+    extent_col_ = extent.strided() / InterleavedN;
+    thread_start_col_ = thread_offset.strided();
+
+    // Initialize predicates
+    CUTLASS_PRAGMA_UNROLL
+    for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) {
+      mask_.predicates[c] =
+          ((thread_offset.contiguous() + ThreadMap::Delta::kContiguous * c) <
+           (extent.contiguous() * InterleavedN));
+    }
+
+    // Initialize pointer
+    byte_pointer_ = reinterpret_cast<uint8_t *>(pointer) + 
+      LongIndex(thread_offset.strided()) * LongIndex(params_.stride) + 
+      LongIndex(thread_offset.contiguous()) * sizeof(AccessType) / kElementsPerAccess;
+
+    // Initialize internal state counter
+    iteration_contiguous_ = iteration_strided_ = 0;
+  }
+
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    byte_pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load(Fragment &frag) {
+
+    uint8_t *byte_pointer = byte_pointer_;
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+    AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer);
+
+    int col_offset = iteration_strided_ * ThreadMap::Delta::kStrided;
+
+    bool col_guard = ((thread_start_col_ + col_offset) < extent_col_);
+
+    bool guard = col_guard && mask_.predicates[iteration_contiguous_];
+
+    cutlass::arch::global_load<
+      AccessType, 
+      sizeof(AccessType)
+    >(
+        *frag_ptr,
+        (void *)memory_pointer,
+        guard);
+  }
+
+  /// Stores a fragment to memory
+  CUTLASS_DEVICE
+  void store(Fragment const &frag) {
+    uint8_t *byte_pointer = byte_pointer_;
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+    AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer);
+
+    int col_offset = iteration_strided_ * ThreadMap::Delta::kStrided;
+
+    bool col_guard = ((thread_start_col_ + col_offset) < extent_col_);
+
+    bool guard = col_guard && mask_.predicates[iteration_contiguous_];
+
+    cutlass::arch::global_store<AccessType, sizeof(AccessType)>(
+        *frag_ptr, (void *)memory_pointer, guard);
+  }
+
+  /// Overrides the internal iteration index
+  CUTLASS_HOST_DEVICE
+  void set_iteration_index(int iteration) {
+    iteration_contiguous_ = iteration % ThreadMap::Iterations::kContiguous;
+    iteration_strided_ = iteration / ThreadMap::Iterations::kContiguous;
+  }
+
+  /// Advances to the next position to load or store
+  CUTLASS_HOST_DEVICE
+  InterleavedPredicatedTileIterator &operator++() {
+
+    ++iteration_contiguous_;
+    byte_pointer_ += params_.advance_row;
+
+    if (iteration_contiguous_ == ThreadMap::Iterations::kContiguous) {
+
+      iteration_contiguous_ = 0;
+      ++iteration_strided_;
+      byte_pointer_ += params_.advance_column;
+
+      if (iteration_strided_ == ThreadMap::Iterations::kStrided) {
+        iteration_strided_ = 0;
+      }
+    }
+
+    return *this;
+  }
+
+  /// Advances a number of positions to load or store
+  CUTLASS_HOST_DEVICE
+  InterleavedPredicatedTileIterator &operator+=(int increment)
+  {
+    // Contiguous
+    iteration_contiguous_ += increment;
+    int increment_strided = iteration_contiguous_ / ThreadMap::Iterations::kContiguous;
+    iteration_contiguous_ = iteration_contiguous_ % ThreadMap::Iterations::kContiguous;
+    byte_pointer_ += (params_.advance_row * increment);
+
+    // Strided
+    iteration_strided_ += increment_strided;
+    byte_pointer_ += (params_.advance_column * increment_strided);
+
+    return *this;
+  }
+
+  ///< Efficiently disables all accesses guarded by mask
+  CUTLASS_DEVICE void clear_mask() {
+    mask_.clear();
+  }
+
+  ///< Efficiently enables all accesses guarded by mask
+  CUTLASS_DEVICE void enable_mask() {
+    mask_.enable();
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void get_mask(Mask &mask) {
+    mask = mask_;
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void set_mask(Mask const &mask) {
+    mask_ = mask;
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+
+/// Tile iterator used to load output tile from global memory in epilogue.
+///
+/// Satisfies: ReadableTileIterator | InterleavedMaskedTileIterator | ForwardTileIterator
+///
+template <
+  typename ThreadMap_,       ///< Thread map (conept: OutputTileThreadMap)
+  typename Element_,         ///< Element data type
+  int InterleavedN           ///< Number of Interleaved N
+>
+class InterleavedConvPredicatedTileIterator {
+public:
+  using ThreadMap = ThreadMap_;
+
+  using Element = Element_;
+
+  using Layout = layout::TensorNCxHWx<InterleavedN>;
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = Tensor4DCoord;
+
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+  static int const kThreads = ThreadMap::kThreads;
+  static int const kIterations = ThreadMap::Iterations::kCount;
+
+  /// Fragment object
+  using Fragment = Array<Element, ThreadMap::kElementsPerAccess>;
+
+  /// Memory access size
+  using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess>;
+
+  //
+  // Parameters struct
+  //
+
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    LongIndex stride_col;           ///< stride in bytes between columns
+    LongIndex stride_row;           ///< stride in bytes between rows
+
+    //
+    // Methods
+    //
+
+    CUTLASS_HOST_DEVICE
+    Status initialize(typename Layout::Stride stride_) {
+      stride_col = stride_[1];
+      stride_row = stride_[2];
+
+      return Status::kSuccess;
+    }
+
+    CUTLASS_HOST_DEVICE
+    Params() {
+      initialize(cutlass::make_Coord(0, 0, 0));
+    }
+
+    CUTLASS_HOST_DEVICE
+    Params(Layout const &layout) {
+
+      initialize(layout.stride());
+    }
+
+    CUTLASS_HOST_DEVICE
+    Params(Layout const &layout,
+           // Not needed.  Added to be compatible with strided conv epilogue.
+           cutlass::Tensor4DCoord const &tensor_extent):
+      Params(layout)
+    { }
+
+  };
+
+  /// Mask object
+  struct Mask {
+    static int const kCount =
+        (ThreadMap::Iterations::kRow < 8) ? 8 : ThreadMap::Iterations::kRow;
+
+    /// Predicate state
+    bool predicates[kCount];
+
+    //
+    // Mask
+    //
+    CUTLASS_HOST_DEVICE
+    Mask() {
+      enable();
+    }
+
+    ///< Efficiently disables all accesses guarded by mask
+    CUTLASS_HOST_DEVICE void clear() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = false;
+      }
+    }
+
+    ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask
+    CUTLASS_DEVICE void enable() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = true;
+      }
+    }
+  };
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Parameters structure containing reference and precomputed state.
+  Params params_;
+
+  /// Byte-level pointer
+  uint8_t *byte_pointer_;
+
+  /// Array of boolean values to contain steady-state predicates
+  Mask mask_;
+
+  /// Extent of the matrix tile in columns
+  Index extent_col_;
+
+  /// Extent of the matrix tile in rows
+  Index extent_row_;
+
+  /// Extent of the matrix tile in pq 
+  Index extent_pq_;
+
+  /// A thread's starting row position (assuming steady-state predicates have
+  /// been computed)
+  Index thread_start_row_;
+
+  /// A thread's starting column position (assuming steady-state predicates have
+  /// been computed)
+  Index thread_start_col_;
+
+  /// Internal iteration counter
+  LongIndex iteration_row_;
+  LongIndex iteration_col_;
+
+  uint32_t pq_mul_;
+
+  uint32_t pq_shr_;
+
+private:
+
+  //
+  // Methods
+  //
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  CUTLASS_DEVICE
+  InterleavedConvPredicatedTileIterator(
+    Params const & params,
+    Element *pointer,
+    TensorCoord extent,
+    int thread_idx,
+    MatrixCoord threadblock_offset
+  ):
+    params_(params) {
+    MatrixCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset;
+                                
+    extent_col_ = extent.c();
+    extent_pq_ = extent.h() * extent.w();
+    extent_row_ = extent.n() * extent_pq_;
+
+    find_divisor(pq_mul_, pq_shr_, extent_pq_);
+
+    thread_start_row_ = thread_offset.row();
+    thread_start_col_ = thread_offset.column();
+
+    // Initialize predicates
+    CUTLASS_PRAGMA_UNROLL
+    for (int r = 0; r < ThreadMap::Iterations::kRow; ++r) {
+      mask_.predicates[r] =
+          ((thread_offset.row() + ThreadMap::Delta::kRow * r) < extent_row_);
+    }
+
+    // Initialize pointer
+    byte_pointer_ = reinterpret_cast<uint8_t *>(pointer) +
+                    ((thread_start_col_ / InterleavedN) * params_.stride_col +
+                     (thread_start_col_ % InterleavedN)) *
+                        sizeof_bits<Element>::value / 8;
+
+    // Initialize internal state counter
+    iteration_row_ = iteration_col_ = 0;
+  }
+
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    byte_pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load(Fragment &frag) {
+
+    int col_offset = iteration_col_ * ThreadMap::Delta::kColumn;
+    bool col_guard = ((thread_start_col_ + col_offset) < extent_col_);
+    bool guard = col_guard && mask_.predicates[iteration_row_];
+
+    int n, pq_rem;
+
+    fast_divmod(n, pq_rem,
+                thread_start_row_ + iteration_row_ * ThreadMap::Delta::kRow,
+                extent_pq_, pq_mul_, pq_shr_);
+
+    uint8_t *byte_pointer =
+        byte_pointer_ + (n * params_.stride_row + pq_rem * InterleavedN) *
+                            sizeof_bits<Element>::value / 8;
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+    AccessType const *memory_pointer =
+        reinterpret_cast<AccessType const *>(byte_pointer);
+
+    cutlass::arch::global_load<
+      AccessType, 
+      sizeof(AccessType)
+    >(
+        *frag_ptr,
+        (void *)memory_pointer,
+        guard);
+  }
+
+  /// Stores a fragment to memory
+  CUTLASS_DEVICE
+  void store(Fragment const &frag) {
+
+    int col_offset = iteration_col_ * ThreadMap::Delta::kColumn;
+    bool col_guard = ((thread_start_col_ + col_offset) < extent_col_);
+    bool guard = col_guard && mask_.predicates[iteration_row_];
+
+    int n, pq_rem;
+
+    fast_divmod(n, pq_rem,
+                thread_start_row_ + iteration_row_ * ThreadMap::Delta::kRow,
+                extent_pq_, pq_mul_, pq_shr_);
+
+    uint8_t *byte_pointer =
+        byte_pointer_ + (n * params_.stride_row + pq_rem * InterleavedN) *
+                            sizeof_bits<Element>::value / 8;
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+    AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer);
+
+    cutlass::arch::global_store<AccessType, sizeof(AccessType)>(
+        *frag_ptr, (void *)memory_pointer, guard);
+  }
+
+  /// Overrides the internal iteration index
+  CUTLASS_HOST_DEVICE
+  void set_iteration_index(int iteration) {
+    iteration_row_ = iteration % ThreadMap::Iterations::kRow;
+    iteration_col_ = iteration / ThreadMap::Iterations::kRow;
+  }
+
+  /// Advances to the next position to load or store
+  CUTLASS_HOST_DEVICE
+  InterleavedConvPredicatedTileIterator &operator++() {
+
+    ++iteration_row_;
+
+    if (iteration_row_ == ThreadMap::Iterations::kRow) {
+
+      iteration_row_ = 0;
+      ++iteration_col_;
+      byte_pointer_ += params_.stride_col;
+
+      if (iteration_col_ == ThreadMap::Iterations::kColumn) {
+        iteration_col_ = 0;
+      }
+    }
+
+    return *this;
+  }
+
+  ///< Efficiently disables all accesses guarded by mask
+  CUTLASS_DEVICE void clear_mask() {
+    mask_.clear();
+  }
+
+  ///< Efficiently enables all accesses guarded by mask
+  CUTLASS_DEVICE void enable_mask() {
+    mask_.enable();
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void get_mask(Mask &mask) {
+    mask = mask_;
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void set_mask(Mask const &mask) {
+    mask_ = mask;
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_affine.h

@@ -0,0 +1,615 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+  
+  \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
+
+  The epilogue rearranges the result of a matrix product through shared memory to match canonical
+  tensor layouts in global memory. Epilogues support conversion and reduction operations.
+
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/layout/tensor.h"
+#include "cutlass/matrix_shape.h"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/transform/pitch_linear_thread_map.h"
+#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
+#include "cutlass/arch/arch.h"
+#include "cutlass/arch/memory.h"
+#include "cutlass/epilogue/threadblock/predicated_tile_iterator_params.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace epilogue {
+namespace threadblock {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Tile iterator used to load and store output tile from global memory in epilogue.
+///
+/// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator
+///
+/// It provides a fast path for the case Rank = 2 which does not need div/rem to 
+/// calculate modes.
+
+template <
+  typename ThreadMap_,       ///< Thread map (conept: OutputTileThreadMap)
+  typename Element_,         ///< Element data type
+  int Rank
+>
+class PredicatedTileIteratorAffineRankN {
+public:
+  using ThreadMap = ThreadMap_;
+  using Shape = typename ThreadMap::Shape;
+
+  using Element = Element_;
+
+  using Layout = layout::AffineRankN<Rank>;
+  using TensorRef = TensorRef<Element, Layout>;
+  using TensorView = TensorView<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = typename Layout::TensorCoord;
+
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+  static int const kThreads = ThreadMap::kThreads;
+  static int const kIterations = ThreadMap::Count::kTile;
+
+  static_assert( ThreadMap::Iterations::kRow > 0,"ThreadMap::Iterations::kRow must be > 0");
+  static_assert( ThreadMap::Iterations::kGroup > 0,"ThreadMap::Iterations::kGroup must be > 0");
+  static_assert( ThreadMap::Iterations::kCluster > 0,"ThreadMap::Iterations::kCluster must be > 0");
+  static_assert( ThreadMap::Iterations::kColumn > 0,"ThreadMap::Iterations::kColumn must be > 0");
+  static_assert( !(Layout::kRank % 2), 
+    "Layout rank must be even. This assumes the first half of the modes correspond to the 'row' "
+    "and the second half of the modes correspond to the 'column'");
+
+  static bool const kBigEndian = false;
+
+  /// Fragment object
+  using Fragment = Array<
+    Element, 
+    ThreadMap::Iterations::kColumn * 
+    ThreadMap::Iterations::kRow * 
+    ThreadMap::Iterations::kGroup * 
+    ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>;
+
+  /// Memory access size
+  using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess>;
+
+  //
+  // Parameters struct
+  //
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    Layout layout;
+
+    /// Stride in units of bytes along M modes
+    Coord<Layout::kRank/2, typename Layout::LongIndex> stride_m;
+
+    /// Stride in units of bytes along N modes
+    Coord<Layout::kRank/2, typename Layout::LongIndex> stride_n;
+
+    /// Fast divmod objects divided by tensor extents
+    FastDivmod divmod_m[(Layout::kRank == 2) ? 1 : (Layout::kRank/2 - 1)];
+
+    /// Fast divmod objects divided by tensor extents
+    FastDivmod divmod_n[(Layout::kRank == 2) ? 1 : (Layout::kRank/2 - 1)];
+
+    int64_t rank2_inc_col;
+    int64_t rank2_inc_row;
+
+    //
+    // Methods
+    //
+    CUTLASS_HOST_DEVICE
+    Params() { }
+
+    CUTLASS_HOST_DEVICE
+    Params(TensorCoord const &extent, Layout const &layout_): layout(layout_) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < Layout::kRank / 2; ++i) {
+        stride_m[i] = OffsetBytes<Element>(layout_.stride()[i]);
+        stride_n[i] = OffsetBytes<Element>(layout_.stride()[i + Layout::kRank / 2]);
+      }
+
+      if (kBigEndian) {
+        // "Big Endian" scheme
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < Layout::kRank / 2 - 1; ++i) {
+          divmod_m[i] = FastDivmod(extent[i + 1]);
+          divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2 + 1]);
+        }
+      }
+      else {
+        // "Little Endian" scheme
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < Layout::kRank / 2 - 1; ++i) {
+          divmod_m[i] = FastDivmod(extent[i]);
+          divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2]);
+        }
+      }
+
+      #if 0
+      //
+      // Debug print statements to verify extents and strides are passed correctly.
+      //
+      printf("PredicatedTileIteratorAffine::Params() entered\n");
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < Layout::kRank; ++i) {
+        printf("  extent[%d]: %d\n", i, extent[i]);
+      }
+      for (int i = 0; i < Layout::kRank; ++i) {
+        printf("  stride[%d]: %ld\n", i, layout_.stride()[i]);
+      }
+      printf("PredicatedTileIteratorAffine::Params() returning\n");
+      #endif
+    }
+
+    CUTLASS_HOST_DEVICE
+    Params(Layout const &layout_): layout(layout_) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < Layout::kRank / 2; ++i) {
+        stride_m[i] = OffsetBytes<Element>(layout_.stride()[i]);
+        stride_n[i] = OffsetBytes<Element>(layout_.stride()[i + Layout::kRank / 2]);
+      }
+
+      rank2_inc_col = ThreadMap::Delta::kColumn * stride_n[0];
+      rank2_inc_row = ThreadMap::Delta::kRow * stride_m[0];
+    }
+  };
+
+  /// Mask object
+  struct Mask {
+
+    static int const kCount = ThreadMap::Iterations::kColumn;
+
+    /// Predicate state
+    bool predicates[kCount];
+
+    //
+    // Mask
+    //
+    CUTLASS_HOST_DEVICE
+    Mask() {
+      enable();
+    }
+
+    ///< Efficiently disables all accesses guarded by mask
+    CUTLASS_HOST_DEVICE void clear() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = false;
+      }
+    }
+
+    ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask
+    CUTLASS_DEVICE void enable() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = true;
+      }
+    }
+  };
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Parameters structure containing reference and precomputed state.
+  Params params_;
+
+  /// Byte-level pointer
+  uint8_t *byte_pointer_;
+
+  /// Array of boolean values to contain steady-state predicates
+  Mask mask_;
+
+  /// Extent of the matrix tile in rows
+  Index extent_row_;
+
+  /// Extent of the matrix tile in columns
+  Index extent_col_;
+
+  /// A thread's starting row position (assuming steady-state predicates have been computed)
+  Index thread_start_row_;
+
+  /// A thread's starting column position (assuming steady-state predicates have been computed)
+  Index thread_start_column_;
+
+  /// Internal state counter
+  int state_[3];
+
+  /// Offsets in columns, cached for performance
+  int64_t offset_modes_n_[ThreadMap::Iterations::kColumn];
+ 
+  //
+  // Static asserts about internal strides
+  //
+
+  static_assert(sizeof(extent_row_) == 4, "Expected 32b extents");
+  static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents");
+
+private:
+
+  //
+  // Methods
+  //
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  CUTLASS_DEVICE
+  PredicatedTileIteratorAffineRankN(
+    Params const & params,
+    Element *pointer,
+    MatrixCoord extent,
+    int thread_idx,
+    MatrixCoord threadblock_offset = MatrixCoord(),
+    int const *indices = nullptr     ///< gather/scatter indices, note no support for gather/scatter at this specialization
+  ): 
+    params_(params)
+  {
+
+    MatrixCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset;
+
+    extent_row_ = extent.row();
+    extent_col_ = extent.column();
+
+    thread_start_row_ = thread_offset.row();
+    thread_start_column_ = thread_offset.column();
+
+    if (Layout::kRank > 2) {
+      // Initialize predicates
+      CUTLASS_PRAGMA_UNROLL
+      for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) {
+
+        // 
+        // Compute coordinate and decompose into N modes
+        //
+
+        int coord_n = thread_start_column_ + c * ThreadMap::Delta::kColumn;
+
+        mask_.predicates[c] = coord_n < extent.column();
+        
+        Coord<Layout::kRank / 2, Index> modes_n;
+
+        int64_t offset_modes_n = 0;
+
+        if (kBigEndian) {
+          modes_n = CoordinateDecomposition<Layout::kRank / 2>(coord_n, params_.divmod_n);
+
+          offset_modes_n = dot(modes_n, params_.stride_n);
+        }
+        else {
+          modes_n = CoordinateDecompositionLittleEndian<Layout::kRank / 2>(coord_n, params_.divmod_n);
+
+          offset_modes_n = dot(modes_n, params_.stride_n);
+        }
+
+        offset_modes_n_[c] = offset_modes_n;
+
+      }
+
+      if (!pointer) {
+        mask_.clear();
+      }
+    }
+
+    // Initialize pointer
+    byte_pointer_ = reinterpret_cast<uint8_t *>(pointer);
+
+    // Initialize internal state counter
+    state_[0] = state_[1] = state_[2] = 0;
+  }
+
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    byte_pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load_with_byte_offset(Fragment &frag, int64_t byte_offset) {
+    uint8_t const *byte_pointer = byte_pointer_;
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        int row_begin = thread_start_row_ + group * ThreadMap::Delta::kGroup + cluster * ThreadMap::Delta::kCluster;
+        int64_t offset_modes_m = row_begin * params_.stride_m[0];
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int frag_row_idx = 
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          // 
+          // Compute coordinate and decompose into M modes
+          //
+
+          int coord_m = row * ThreadMap::Delta::kRow + row_begin;
+
+          Coord<Layout::kRank / 2, Index> modes_m;
+
+          if (Layout::kRank > 2) {
+            if (kBigEndian) {
+              modes_m = CoordinateDecomposition<Layout::kRank / 2>(coord_m, params_.divmod_m);
+            } else {
+              modes_m = CoordinateDecompositionLittleEndian<Layout::kRank / 2>(coord_m, params_.divmod_m);
+            }
+
+            offset_modes_m = dot(modes_m, params_.stride_m);
+          }
+
+          //
+          // Compute the offset due to modes M
+          //
+
+          bool row_guard = (coord_m < extent_row_);
+          int64_t offset_modes_n = thread_start_column_ * params_.stride_n[0];
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+
+            // 
+            // Compute coordinate and decompose into N modes
+            //
+            
+            if (Layout::kRank > 2) {
+              offset_modes_n = offset_modes_n_[column];
+            }
+
+            //
+            // Compute the pointer and access
+            //
+            bool guard;
+
+            if (Layout::kRank > 2) {
+              guard = row_guard && mask_.predicates[column];
+            } else {
+              guard = (coord_m < extent_row_) && 
+              ((thread_start_column_ + ThreadMap::Delta::kColumn * column) < extent_col_);
+            }
+
+            cutlass::arch::global_load<
+              AccessType, 
+              sizeof(AccessType)
+            >(
+              frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column],
+              (void *)(byte_pointer + offset_modes_m + offset_modes_n + byte_offset),
+              guard
+            );
+
+            if (Layout::kRank == 2) {
+              offset_modes_n += params_.rank2_inc_col;
+            }
+          }
+
+          if (Layout::kRank == 2) {
+            offset_modes_m += params_.rank2_inc_row;
+          }
+        }
+      }
+    }
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load(Fragment &frag) {
+
+    load_with_byte_offset(frag, 0);
+  }
+
+  /// Stores a fragment to memory
+  CUTLASS_DEVICE
+  void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) {
+    uint8_t *byte_pointer = byte_pointer_;
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        int row_begin = thread_start_row_ + group * ThreadMap::Delta::kGroup + cluster * ThreadMap::Delta::kCluster;
+        int64_t offset_modes_m = row_begin * params_.stride_m[0];
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int frag_row_idx = 
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          // 
+          // Compute coordinate and decompose into M modes
+          //
+
+          int coord_m = row * ThreadMap::Delta::kRow + row_begin;
+
+          Coord<Layout::kRank / 2, Index> modes_m;
+
+          if (Layout::kRank > 2) {
+            if (kBigEndian) {
+              modes_m = CoordinateDecomposition<Layout::kRank / 2>(coord_m, params_.divmod_m);
+            } else {
+              modes_m = CoordinateDecompositionLittleEndian<Layout::kRank / 2>(coord_m, params_.divmod_m);
+            }
+
+            offset_modes_m = dot(modes_m, params_.stride_m);
+          }
+
+          //
+          // Compute the offset due to modes M
+          //
+
+          bool row_guard = (coord_m < extent_row_);
+          int64_t offset_modes_n = thread_start_column_ * params_.stride_n[0];
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+
+            // 
+            // Compute coordinate and decompose into N modes
+            //
+            
+            if (Layout::kRank > 2) {
+              offset_modes_n = offset_modes_n_[column];
+            } 
+
+            //
+            // Compute the pointer and access
+            //
+            bool guard;
+            if (Layout::kRank > 2) {            
+              guard = row_guard && mask_.predicates[column];
+            } else {
+              guard = (coord_m < extent_row_) && ((thread_start_column_ + ThreadMap::Delta::kColumn * column) < extent_col_);
+            }
+
+            cutlass::arch::global_store<AccessType, sizeof(AccessType)>(
+                frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column],
+                (void *)(byte_pointer + offset_modes_m + offset_modes_n + byte_offset),
+                guard);
+
+            if (Layout::kRank == 2) {
+              offset_modes_n += params_.rank2_inc_col;
+            }
+          }
+
+          if (Layout::kRank == 2) {
+            offset_modes_m += params_.rank2_inc_row;
+          }
+        }
+      }
+    }
+  }
+
+  /// Stores a fragment to memory
+  CUTLASS_DEVICE
+  void store(Fragment const &frag) {
+
+    store_with_byte_offset(frag, 0);
+  }
+
+  /// Advances to the next position to load or store
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorAffineRankN &operator++() {
+
+    ++state_[0];
+    thread_start_row_ += ThreadMap::Shape::kRow;
+    
+    if (state_[0] == ThreadMap::Count::kRow) {
+
+      state_[0] = 0;
+      ++state_[1];
+
+      thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * 
+        ThreadMap::Shape::kRow * ThreadMap::Count::kRow;
+
+      if (state_[1] == ThreadMap::Count::kGroup) {
+
+        state_[1] = 0;
+        ++state_[2];
+
+        thread_start_row_ += ThreadMap::Count::kGroup * 
+          ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow;
+
+        if (state_[2] == ThreadMap::Count::kCluster) {
+          state_[2] = 0;
+        }
+      }
+    }
+
+    return *this;
+  }
+
+  ///< Efficiently disables all accesses guarded by mask
+  CUTLASS_DEVICE void clear_mask() {
+    mask_.clear();
+  }
+
+  ///< Efficiently enables all accesses guarded by mask
+  CUTLASS_DEVICE void enable_mask() {
+    mask_.enable();
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void get_mask(Mask &mask) {
+    mask = mask_;
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void set_mask(Mask const &mask) {
+    mask_ = mask;
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_affine_layout_params.h

@@ -0,0 +1,156 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+  \brief 
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/fast_math.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace threadblock {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  int Rank
+>
+struct PredicatedTileIteratorAffineLayoutRankNParams {
+  using Layout = layout::AffineRankN<Rank>;
+  using TensorCoord = typename Layout::TensorCoord;
+
+  static bool const kBigEndian = false;
+  
+  //
+  // Data members
+  //
+
+  Layout layout;
+
+  /// Stride in units of bytes along M modes
+  Coord<Layout::kRank/2, typename Layout::LongIndex> stride_m;
+
+  /// Stride in units of bytes along N modes
+  Coord<Layout::kRank/2, typename Layout::LongIndex> stride_n;
+
+  /// Fast divmod objects divided by tensor extents
+  FastDivmod divmod_m[(Layout::kRank == 2) ? 1 : (Layout::kRank/2 - 1)];
+
+  /// Fast divmod objects divided by tensor extents
+  FastDivmod divmod_n[(Layout::kRank == 2) ? 1 : (Layout::kRank/2 - 1)];
+
+  int64_t rank2_inc_col;
+  int64_t rank2_inc_row;
+
+  //
+  // Methods
+  //
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorAffineLayoutRankNParams() { }
+
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorAffineLayoutRankNParams(TensorCoord const &extent, 
+                                                Layout const &layout_,
+                                                int64_t element_sizeof_bits)
+  : layout(layout_) 
+  {
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Layout::kRank / 2; ++i) {
+      stride_m[i] = OffsetBytes(layout_.stride()[i], element_sizeof_bits);
+      stride_n[i] = OffsetBytes(layout_.stride()[i + Layout::kRank / 2], element_sizeof_bits);
+    }
+
+    if (kBigEndian) {
+      // "Big Endian" scheme
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < Layout::kRank / 2 - 1; ++i) {
+        divmod_m[i] = FastDivmod(extent[i + 1]);
+        divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2 + 1]);
+      }
+    }
+    else {
+      // "Little Endian" scheme
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < Layout::kRank / 2 - 1; ++i) {
+        divmod_m[i] = FastDivmod(extent[i]);
+        divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2]);
+      }
+    }
+
+    #if 0
+    //
+    // Debug print statements to verify extents and strides are passed correctly.
+    //
+    printf("PredicatedTileIteratorAffine::Params() entered\n");
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Layout::kRank; ++i) {
+      printf("  extent[%d]: %d\n", i, extent[i]);
+    }
+    for (int i = 0; i < Layout::kRank; ++i) {
+      printf("  stride[%d]: %ld\n", i, layout_.stride()[i]);
+    }
+    printf("PredicatedTileIteratorAffine::Params() returning\n");
+    #endif
+  }
+
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorAffineLayoutRankNParams(Layout const &layout_,
+                                                int32_t threadmap_delta_kColumn,
+                                                int32_t threadmap_delta_kRow,
+                                                int64_t element_sizeof_bits)
+  : layout(layout_) 
+  {
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Layout::kRank / 2; ++i) {
+      stride_m[i] = OffsetBytes(layout_.stride()[i], element_sizeof_bits);
+      stride_n[i] = OffsetBytes(layout_.stride()[i + Layout::kRank / 2], element_sizeof_bits);
+    }
+
+    rank2_inc_col = threadmap_delta_kColumn * stride_n[0];
+    rank2_inc_row = threadmap_delta_kRow * stride_m[0];
+  }
+};
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_blas3.h

@@ -0,0 +1,633 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+  \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
+
+  The epilogue rearranges the result of a matrix product through shared memory to match canonical
+  tensor layouts in global memory. Epilogues support conversion and reduction operations.
+
+  
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/layout/tensor.h"
+#include "cutlass/matrix_shape.h"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/transform/pitch_linear_thread_map.h"
+#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
+#include "cutlass/arch/arch.h"
+#include "cutlass/arch/memory.h"
+#include "cutlass/epilogue/threadblock/predicated_tile_iterator_params.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace epilogue {
+namespace threadblock {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Tile iterator used to load and store output tile from global memory in epilogue.
+///
+/// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator
+///
+template <
+  typename ThreadMap_,                     ///< Thread map (conept: OutputTileThreadMap)
+  typename Element_,                        ///< Element data type
+  BlasMode BlasMode_ = BlasMode::kGemm   ///< Tile Iterator for a Symmetric or Hermitian Kernel
+>
+class PredicatedTileIteratorBlas3 {
+public:
+  using ThreadMap = ThreadMap_;
+  using Shape = typename ThreadMap::Shape;
+
+  using Element = Element_;
+
+  using Layout = layout::RowMajor;
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = MatrixCoord;
+
+  static BlasMode const kBlasMode = BlasMode_;
+
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+  static int const kThreads = ThreadMap::kThreads;
+  static int const kIterations = ThreadMap::Count::kTile;
+
+  static_assert( ThreadMap::Iterations::kRow > 0,"ThreadMap::Iterations::kRow must be > 0");
+  static_assert( ThreadMap::Iterations::kGroup > 0,"ThreadMap::Iterations::kGroup must be > 0");
+  static_assert( ThreadMap::Iterations::kCluster > 0,"ThreadMap::Iterations::kCluster must be > 0");
+  static_assert( ThreadMap::Iterations::kColumn > 0,"ThreadMap::Iterations::kColumn must be > 0");
+
+  /// Fragment object
+  using Fragment = Array<
+    Element, 
+    ThreadMap::Iterations::kColumn * 
+    ThreadMap::Iterations::kRow * 
+    ThreadMap::Iterations::kGroup * 
+    ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>;
+
+  /// Memory access size
+  using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess>;
+  static_assert( AccessType::kElements == 1, "BLAS3 Epilogue must use AccessType::kElements as 1");
+
+  //
+  // Parameters struct
+  //
+
+  /// Uses a non-template class
+  struct Params : PredicatedTileIteratorParams {
+
+    CUTLASS_HOST_DEVICE
+    Params() { }
+
+    CUTLASS_HOST_DEVICE
+    Params(Layout const &layout): 
+      PredicatedTileIteratorParams(
+        layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess,
+        make_OutputTileThreadMapDesc<ThreadMap>()
+      ) 
+    {
+        
+    }
+  };
+
+  /// Mask object
+  struct Mask {
+
+    static int const kCount = ThreadMap::Iterations::kColumn;
+
+    /// Predicate state
+    bool predicates[kCount];
+
+    //
+    // Mask
+    //
+    CUTLASS_HOST_DEVICE
+    Mask() {
+      enable();
+    }
+
+    ///< Efficiently disables all accesses guarded by mask
+    CUTLASS_HOST_DEVICE void clear() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = false;
+      }
+    }
+
+    ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask
+    CUTLASS_DEVICE void enable() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = true;
+      }
+    }
+  };
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Parameters structure containing reference and precomputed state.
+  PredicatedTileIteratorParams params_;
+
+  /// Byte-level pointer
+  uint8_t *byte_pointer_;
+
+  /// Fill Mode for a tile on diagonal of a symmetric kernel
+  cutlass::FillMode fill_mode;
+
+  /// Array of boolean values to contain steady-state predicates
+  Mask mask_;
+
+  /// Extent of the matrix tile in rows
+  Index extent_row_;
+
+  /// A thread's starting row position (assuming steady-state predicates have been computed)
+  Index thread_start_row_;
+
+  /// Internal state counter
+  int state_[3];
+
+  /// Starting address of the matrix  
+  size_t matrix_start_addr; 
+ 
+  static_assert((kBlasMode == BlasMode::kSymmetric || kBlasMode == BlasMode::kHermitian), 
+    "Unsupported blas3 mode.");
+
+private:
+
+  //
+  // Methods
+  //
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  CUTLASS_DEVICE
+  PredicatedTileIteratorBlas3(
+    PredicatedTileIteratorParams const & params,
+    Element *pointer,
+    TensorCoord extent,
+    int thread_idx,
+    TensorCoord threadblock_offset
+    , cutlass::FillMode fill_mode
+  ): 
+    params_(params), fill_mode(fill_mode)
+  {
+
+    TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset;
+
+    extent_row_ = extent.row();
+    thread_start_row_ = thread_offset.row();
+
+    // Initialize predicates
+    CUTLASS_PRAGMA_UNROLL
+    for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) {
+
+      mask_.predicates[c] = ((thread_offset.column() 
+        + ThreadMap::Delta::kColumn * c) < extent.column());
+    }
+
+    // Check Symmetric kernel modes (Lower and Upper - for diagonal CTAs, None for rest CTAs)
+    if ((kBlasMode == BlasMode::kSymmetric || kBlasMode == BlasMode::kHermitian) && 
+        fill_mode == cutlass::FillMode::kInvalid) {
+      arch::device_breakpoint();
+    }
+
+    // Starting address of the matrix
+    matrix_start_addr =  reinterpret_cast<size_t>(pointer); 
+
+    // Initialize pointer
+    byte_pointer_ = reinterpret_cast<uint8_t *>(pointer) + 
+      LongIndex(thread_offset.row()) * LongIndex(params_.stride) + 
+      LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess;
+
+    // Initialize internal state counter
+    state_[0] = state_[1] = state_[2] = 0;
+  }
+
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    byte_pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load_with_byte_offset(Fragment &frag, int64_t byte_offset) {
+
+    uint8_t *byte_pointer = byte_pointer_;
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int frag_row_idx = 
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          int row_offset = row * ThreadMap::Delta::kRow 
+            + group * ThreadMap::Delta::kGroup 
+            + cluster * ThreadMap::Delta::kCluster;
+
+          bool row_guard = ((row_offset + thread_start_row_) < extent_row_);
+
+          AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+
+            bool guard = row_guard && mask_.predicates[column];
+
+            cutlass::arch::global_load<
+              AccessType, 
+              sizeof(AccessType)
+            >(
+                frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn +
+                         column],
+                (void *)&memory_pointer[column * ThreadMap::Delta::kColumn /
+                                        kElementsPerAccess],
+                guard);
+          }
+
+          if (row + 1 < ThreadMap::Iterations::kRow) {
+            byte_pointer += params_.increment_row;
+          }
+        }
+
+        if (group + 1 < ThreadMap::Iterations::kGroup) {
+          byte_pointer += params_.increment_group;
+        }
+      }
+
+      if (cluster + 1 < ThreadMap::Iterations::kCluster) {
+        byte_pointer += params_.increment_cluster;
+      }
+    }
+  }
+
+  /// Loads a fragment on the diagonal of a symmetric kernel to memory 
+  CUTLASS_DEVICE
+  void load_symmetric_with_byte_offset(Fragment &frag, int64_t byte_offset) {
+
+    uint8_t *byte_pointer = byte_pointer_;
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    bool isLowerMode = (fill_mode == cutlass::FillMode::kLower) ? true : false;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int frag_row_idx = 
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          int row_offset = row * ThreadMap::Delta::kRow 
+            + group * ThreadMap::Delta::kGroup 
+            + cluster * ThreadMap::Delta::kCluster;
+
+          bool row_guard = ((row_offset + thread_start_row_) < extent_row_);
+
+          AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset);
+
+          // Offset of row from beginning of the matrix per thread
+          size_t row_start_offset = (size_t)memory_pointer - matrix_start_addr;
+
+          // Absolute row index
+          int row_index = int(row_start_offset/params_.stride);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+
+            bool guard = row_guard && mask_.predicates[column];
+
+            // Offset of column from beginning of row per thread     
+            size_t col_start_offset = row_start_offset + 
+                        (column * ThreadMap::Delta::kColumn / kElementsPerAccess) * sizeof(AccessType);
+
+            // Absolute column index
+            size_t col_index = (col_start_offset%params_.stride)/sizeof(AccessType);
+            guard = guard && ( (isLowerMode && row_index >= col_index) ||
+                               (!isLowerMode && row_index <= col_index) );
+
+            cutlass::arch::global_load<
+              AccessType, 
+              sizeof(AccessType)
+            >(
+                frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn +
+                         column],
+                (void *)&memory_pointer[column * ThreadMap::Delta::kColumn /
+                                        kElementsPerAccess],
+                guard);
+
+            // The imaginary parts of the diagonal elements of a complex element are assumed and set to zero
+            if (guard && kBlasMode == BlasMode::kHermitian && cutlass::is_complex<Element>::value) {
+              Element *scalar_ptr = reinterpret_cast<Element *>(frag_ptr);
+
+              if (row_index == col_index) {
+                scalar_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column] = 
+                  real(scalar_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column]);
+              }
+            }
+          }
+
+          if (row + 1 < ThreadMap::Iterations::kRow) {
+            byte_pointer += params_.increment_row;
+          }
+        }
+
+        if (group + 1 < ThreadMap::Iterations::kGroup) {
+          byte_pointer += params_.increment_group;
+        }
+      }
+
+      if (cluster + 1 < ThreadMap::Iterations::kCluster) {
+        byte_pointer += params_.increment_cluster;
+      }
+    }
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load(Fragment &frag) {
+    
+    if (fill_mode == cutlass::FillMode::kNone) {
+      load_with_byte_offset(frag, 0);
+    }
+    else {
+      load_symmetric_with_byte_offset(frag, 0);
+    }
+  }
+
+  /// Stores a fragment to memory
+  CUTLASS_DEVICE
+  void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) {
+    uint8_t *byte_pointer = byte_pointer_;
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int frag_row_idx = 
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          int row_offset = row * ThreadMap::Delta::kRow 
+            + group * ThreadMap::Delta::kGroup 
+            + cluster * ThreadMap::Delta::kCluster;
+
+          bool row_guard = ((row_offset + thread_start_row_) < extent_row_);
+
+          AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+
+            bool guard = row_guard && mask_.predicates[column];
+
+            cutlass::arch::global_store<AccessType, sizeof(AccessType)>(
+                frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column],
+                (void *)&memory_pointer[column * ThreadMap::Delta::kColumn / kElementsPerAccess],
+                guard);
+          }
+
+          if (row + 1 < ThreadMap::Iterations::kRow) {
+            byte_pointer += params_.increment_row;
+          }
+        }
+
+        if (group + 1 < ThreadMap::Iterations::kGroup) {
+          byte_pointer += params_.increment_group;
+        }
+      }
+
+      if (cluster + 1 < ThreadMap::Iterations::kCluster) {
+        byte_pointer += params_.increment_cluster;
+      }
+    }
+  }
+
+  /// Stores a fragment on the diagonal of a symmetric kernel to memory 
+  CUTLASS_DEVICE
+  void store_symmetric_with_byte_offset(Fragment const &frag, int64_t byte_offset) {
+    uint8_t *byte_pointer = byte_pointer_;
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    bool isLowerMode = (fill_mode == cutlass::FillMode::kLower) ? true : false;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int frag_row_idx = 
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          int row_offset = row * ThreadMap::Delta::kRow 
+            + group * ThreadMap::Delta::kGroup 
+            + cluster * ThreadMap::Delta::kCluster;
+
+          bool row_guard = ((row_offset + thread_start_row_) < extent_row_);
+
+          AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset);
+
+          // Offset of row from beginning of the matrix per thread
+          size_t row_start_offset = (size_t)memory_pointer - matrix_start_addr;
+
+          // Absolute row index
+          int row_index = int(row_start_offset/params_.stride);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+
+            bool guard = row_guard && mask_.predicates[column];
+
+            // Offset of column from beginning of row per thread     
+            size_t col_start_offset = row_start_offset + 
+                        (column * ThreadMap::Delta::kColumn / kElementsPerAccess) * sizeof(AccessType);
+
+            // Absolute column index
+            size_t col_index = (col_start_offset%params_.stride)/sizeof(AccessType);
+
+            guard = guard && ( (isLowerMode && row_index >= col_index) ||
+                               (!isLowerMode && row_index <= col_index) );
+
+            // The imaginary parts of the diagonal elements of a complex element are assumed and set to zero
+            if (guard && kBlasMode == BlasMode::kHermitian && cutlass::is_complex<Element>::value) {
+
+              AccessType *frag_ptr_modify = const_cast<AccessType *>(frag_ptr);
+              Element *scalar_ptr = reinterpret_cast<Element *>(frag_ptr_modify);
+
+              if (row_index == col_index) {
+                scalar_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column] = 
+                  real(scalar_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column]);
+              }
+            }
+
+            cutlass::arch::global_store<AccessType, sizeof(AccessType)>(
+                frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn +
+                         column],
+                (void *)&memory_pointer[column * ThreadMap::Delta::kColumn /
+                                        kElementsPerAccess],
+                guard);
+          }
+
+          if (row + 1 < ThreadMap::Iterations::kRow) {
+            byte_pointer += params_.increment_row;
+          }
+        }
+
+        if (group + 1 < ThreadMap::Iterations::kGroup) {
+          byte_pointer += params_.increment_group;
+        }
+      }
+
+      if (cluster + 1 < ThreadMap::Iterations::kCluster) {
+        byte_pointer += params_.increment_cluster;
+      }
+    }
+  }
+
+  /// Stores a fragment to memory
+  CUTLASS_DEVICE
+  void store(Fragment const &frag) {
+    
+    if (fill_mode == cutlass::FillMode::kNone) {
+      store_with_byte_offset(frag, 0);
+    }
+    else {
+      store_symmetric_with_byte_offset(frag, 0); 
+    }
+
+  }
+
+  /// Advances to the next position to load or store
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorBlas3 &operator++() {
+
+    ++state_[0];
+    byte_pointer_ += params_.advance_row;
+    thread_start_row_ += ThreadMap::Shape::kRow;
+    
+    if (state_[0] == ThreadMap::Count::kRow) {
+
+      state_[0] = 0;
+      ++state_[1];
+      byte_pointer_ += params_.advance_group;
+
+      thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * 
+        ThreadMap::Shape::kRow * ThreadMap::Count::kRow;
+
+      if (state_[1] == ThreadMap::Count::kGroup) {
+
+        state_[1] = 0;
+        ++state_[2];
+        byte_pointer_ += params_.advance_cluster;
+
+        thread_start_row_ += ThreadMap::Count::kGroup * 
+          ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow;
+
+        if (state_[2] == ThreadMap::Count::kCluster) {
+          state_[2] = 0;
+          byte_pointer_ += params_.advance_tile;
+        }
+      }
+    }
+
+    return *this;
+  }
+
+  ///< Efficiently disables all accesses guarded by mask
+  CUTLASS_DEVICE void clear_mask() {
+    mask_.clear();
+  }
+
+  ///< Efficiently enables all accesses guarded by mask
+  CUTLASS_DEVICE void enable_mask() {
+    mask_.enable();
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void get_mask(Mask &mask) {
+    mask = mask_;
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void set_mask(Mask const &mask) {
+    mask_ = mask;
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_conv.h

@@ -0,0 +1,562 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+  \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
+
+  The epilogue rearranges the result of a matrix product through shared memory to match canonical
+  tensor layouts in global memory. Epilogues support conversion and reduction operations.
+
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/layout/tensor.h"
+#include "cutlass/layout/permute.h"
+#include "cutlass/matrix_shape.h"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/transform/pitch_linear_thread_map.h"
+#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
+#include "cutlass/arch/arch.h"
+#include "cutlass/arch/memory.h"
+#include "cutlass/epilogue/threadblock/predicated_tile_iterator_params.h"
+#include "cutlass/conv/conv2d_problem_size.h"
+#include "cutlass/conv/conv3d_problem_size.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace epilogue {
+namespace threadblock {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Tile iterator used to load and store output tile from global memory in epilogue.
+///
+/// Satisfies: ReadableTileIterator | PredicatedTileIteratorConv | ForwardTileIterator
+///
+template <
+  typename ThreadMap_,       ///< Thread map (conept: OutputTileThreadMap)
+  typename Element_,         ///< Element data type
+  bool ScatterD = false,     ///< Scatter D operand or not
+  typename PermuteDLayout = layout::NoPermute, ///< Permute D operand or not
+  bool UseCUDAStore = false,
+  int Rank = 4
+>
+class PredicatedTileIteratorConv {
+public:
+  using ThreadMap = ThreadMap_;
+  using Shape = typename ThreadMap::Shape;
+
+  using Element = Element_;
+
+  static int const kRank = Rank;
+  using Layout = typename platform::conditional<kRank == 4,
+                                       layout::TensorNHWC,
+                                       layout::TensorNDHWC>::type;
+
+  using Stride = typename Layout::Stride;
+  static int const kStrideRank = Layout::kStrideRank;
+
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  using MappedLayout = layout::RowMajor;
+  using Index = typename MappedLayout::Index;
+  using LongIndex = typename MappedLayout::LongIndex;
+  using TensorCoord = typename MappedLayout::TensorCoord;
+
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+  static int const kThreads = ThreadMap::kThreads;
+  static int const kIterations = ThreadMap::Count::kTile;
+
+  static bool constexpr PermuteD = !layout::is_trivial_permute<PermuteDLayout>;
+
+  static_assert( ThreadMap::Iterations::kRow > 0,"ThreadMap::Iterations::kRow must be > 0");
+  static_assert( ThreadMap::Iterations::kGroup > 0,"ThreadMap::Iterations::kGroup must be > 0");
+  static_assert( ThreadMap::Iterations::kCluster > 0,"ThreadMap::Iterations::kCluster must be > 0");
+  static_assert( ThreadMap::Iterations::kColumn > 0,"ThreadMap::Iterations::kColumn must be > 0");
+
+  /// Fragment object
+  using Fragment = Array<
+    Element,
+    ThreadMap::Iterations::kColumn *
+    ThreadMap::Iterations::kRow *
+    ThreadMap::Iterations::kGroup *
+    ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>;
+
+  /// Memory access size
+  using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess>;
+
+  //
+  // Parameters struct
+  //
+
+  /// Uses a non-template class
+  struct Params : PredicatedTileIteratorParams {
+    using Base = PredicatedTileIteratorParams;
+
+    /// Fast divmod objects divided by tensor extents
+    FastDivmod divmod[kStrideRank - 1];
+    Stride tensor_stride;
+
+    CUTLASS_HOST_DEVICE
+    Params() { }
+
+    CUTLASS_HOST_DEVICE
+    Params(Layout const &layout, cutlass::Tensor4DCoord const &tensor_extent):
+      PredicatedTileIteratorParams(
+        layout.stride()[0] * int(sizeof(AccessType)) / kElementsPerAccess,
+        make_OutputTileThreadMapDesc<ThreadMap>()
+      ) {
+      divmod[0] = FastDivmod(tensor_extent[2] /* Q for Fprop & W for Deconv*/);
+      divmod[1] = FastDivmod(tensor_extent[1] /* P for Fprop & H for Deconv*/);
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kStrideRank; ++i) {
+        tensor_stride[i] = layout.stride()[i];
+      }
+    }
+
+    CUTLASS_HOST_DEVICE
+    Params(Layout const &layout, cutlass::Tensor5DCoord const &tensor_extent):
+      PredicatedTileIteratorParams(
+        layout.stride()[0] * int(sizeof(AccessType)) / kElementsPerAccess,
+        make_OutputTileThreadMapDesc<ThreadMap>()
+      ) {
+      divmod[0] = FastDivmod(tensor_extent[3] /* Q for Fprop & W for Deconv*/);
+      divmod[1] = FastDivmod(tensor_extent[2] /* P for Fprop & H for Deconv*/);
+      divmod[2] = FastDivmod(tensor_extent[1] /* Z for Fprop & D for Deconv*/);
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kStrideRank; ++i) {
+        tensor_stride[i] = layout.stride()[i];
+      }
+    }
+
+    CUTLASS_HOST_DEVICE
+    Params(Base const &base) :
+      Base(base) { }
+  };
+
+  /// Mask object
+  struct Mask {
+
+    static int const kCount = ThreadMap::Iterations::kColumn;
+
+    /// Predicate state
+    bool predicates[kCount];
+
+    //
+    // Mask
+    //
+    CUTLASS_HOST_DEVICE
+    Mask() {
+      enable();
+    }
+
+    ///< Efficiently disables all accesses guarded by mask
+    CUTLASS_HOST_DEVICE void clear() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = false;
+      }
+    }
+
+    ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask
+    CUTLASS_DEVICE void enable() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = true;
+      }
+    }
+  };
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Parameters structure containing reference and precomputed state.
+  Params params_;
+
+  /// Byte-level pointer. This pointer is usually for both load() and store(), unless PermuteD is performed. When having PermuteD, byte_pointer_ is only for load().
+  uint8_t *byte_pointer_;
+
+  /// Array of boolean values to contain steady-state predicates
+  Mask mask_;
+
+  /// Extent of the matrix tile in rows
+  Index extent_row_;
+
+  /// Extent of the matrix tile in rows
+  Index extent_column_;
+
+  /// A thread's starting row position (assuming steady-state predicates have been computed)
+  Index thread_start_row_;
+
+  /// A thread's starting column
+  Index thread_start_column_;
+
+  /// Internal state counter
+  int state_[3];
+
+  //
+  // Static asserts about internal strides
+  //
+
+  static_assert(sizeof(extent_row_) == 4, "Expected 32b extents");
+  static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents");
+  static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides");
+
+private:
+
+  //
+  // Methods
+  //
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  CUTLASS_DEVICE
+  PredicatedTileIteratorConv(
+    Params const & params,
+    Element *pointer,
+    TensorCoord extent,
+    int thread_idx,
+    TensorCoord threadblock_offset = TensorCoord()
+  ):
+    params_(params)
+  {
+
+    TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset;
+
+    extent_row_ = extent.row();
+    extent_column_ = extent.column();
+
+    thread_start_row_ = thread_offset.row();
+    thread_start_column_ = thread_offset.column();
+
+    // Initialize predicates
+    CUTLASS_PRAGMA_UNROLL
+    for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) {
+
+      mask_.predicates[c] = ((thread_offset.column()
+        + ThreadMap::Delta::kColumn * c) < extent.column());
+    }
+
+    // Null pointer performs no accesses
+    if (!pointer) {
+      mask_.clear();
+    }
+
+    // Initialize byte_pointer_
+    byte_pointer_ = reinterpret_cast<uint8_t *>(pointer) +
+      LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess;
+
+    // Initialize internal state counter
+    state_[0] = state_[1] = state_[2] = 0;
+  }
+
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    byte_pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load_with_byte_offset(Fragment &frag, int64_t byte_offset) const {
+
+    uint8_t *byte_pointer = byte_pointer_;
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int frag_row_idx =
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          int row_offset = row * ThreadMap::Delta::kRow
+            + group * ThreadMap::Delta::kGroup
+            + cluster * ThreadMap::Delta::kCluster;
+
+          bool row_guard = ((row_offset + thread_start_row_) < extent_row_);
+
+          AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset);
+
+          Stride tensor_coord = CoordinateDecompositionLittleEndian<kStrideRank>(row_offset + thread_start_row_, params_.divmod);
+
+          LongIndex tensor_offset = dot(tensor_coord, params_.tensor_stride);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+
+            bool guard = row_guard && mask_.predicates[column];
+
+            cutlass::arch::global_load<
+              AccessType,
+              sizeof(AccessType)
+            >(
+                frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn +
+                         column],
+                (void *)&memory_pointer[column * ThreadMap::Delta::kColumn /
+                                        kElementsPerAccess + tensor_offset / kElementsPerAccess],
+                guard);
+          }
+        }
+      }
+    }
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load(Fragment &frag) const {
+
+    load_with_byte_offset(frag, 0);
+  }
+
+  /// Stores a fragment to memory
+  CUTLASS_DEVICE
+  void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) const {
+    uint8_t *byte_pointer = byte_pointer_;
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int frag_row_idx =
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          int row_offset = row * ThreadMap::Delta::kRow
+            + group * ThreadMap::Delta::kGroup
+            + cluster * ThreadMap::Delta::kCluster;
+
+          bool row_guard = ((row_offset + thread_start_row_) < extent_row_);
+
+          Stride tensor_coord = CoordinateDecompositionLittleEndian<kStrideRank>((row_offset + thread_start_row_), params_.divmod);
+
+          LongIndex tensor_offset = dot(tensor_coord, params_.tensor_stride);
+
+          AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+
+            bool guard = row_guard && mask_.predicates[column];
+
+            if (UseCUDAStore) {
+              if (guard) {
+                memory_pointer[tensor_offset / kElementsPerAccess] =
+                    frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column];
+              }
+            } else {
+              cutlass::arch::global_store<AccessType, sizeof(AccessType)>(
+                  frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column],
+                  (void *)&memory_pointer[tensor_offset / kElementsPerAccess],
+                  guard);
+            }
+
+            memory_pointer += (ThreadMap::Delta::kColumn / kElementsPerAccess);
+          }
+        }
+      }
+    }
+  }
+
+  /// Stores a fragment to memory
+  CUTLASS_DEVICE
+  void store(Fragment const &frag) const {
+
+    store_with_byte_offset(frag, 0);
+  }
+
+  CUTLASS_DEVICE
+  MatrixCoord thread_start() const {
+    return MatrixCoord(thread_start_row_, thread_start_column_);
+  }
+
+  /// Need to get the thread start row from the tile iterator
+  CUTLASS_DEVICE
+  int32_t thread_start_row() const {
+    return thread_start_row_;
+  }
+
+  /// Need to get the thread start row from the tile iterator
+  CUTLASS_DEVICE
+  int32_t thread_start_column() const {
+    return thread_start_column_;
+  }
+
+  /// Extent of the matrix in rows
+  CUTLASS_DEVICE
+  Index extent_row() const {
+    return extent_row_;
+  }
+
+  /// Extent of the matrix in columns
+  CUTLASS_DEVICE
+  Index extent_column() const {
+    return extent_column_;
+  }
+
+  /// Advances to the next position to load or store
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorConv &operator++() {
+
+    ++state_[0];
+
+    thread_start_row_ += ThreadMap::Shape::kRow;
+
+    if (state_[0] == ThreadMap::Count::kRow) {
+
+      state_[0] = 0;
+      ++state_[1];
+
+      thread_start_row_ += (ThreadMap::Shape::kGroup - 1) *
+        ThreadMap::Shape::kRow * ThreadMap::Count::kRow;
+
+      if (state_[1] == ThreadMap::Count::kGroup) {
+
+        state_[1] = 0;
+        ++state_[2];
+
+        thread_start_row_ += ThreadMap::Count::kGroup *
+          ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow;
+
+        if (state_[2] == ThreadMap::Count::kCluster) {
+          state_[2] = 0;
+
+          thread_start_row_ += ThreadMap::Shape::kGroup * ThreadMap::Shape::kRow
+            * ThreadMap::Shape::kCluster * ThreadMap::Shape::kTile;
+        }
+      }
+    }
+
+    return *this;
+  }
+
+  /// Advances a number of positions to load or store
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorConv &operator+=(int increment)
+  {
+    // Row
+    state_[0] += increment;
+    int increment_row = state_[0] / ThreadMap::Count::kRow;
+    state_[0] = state_[0] % ThreadMap::Count::kRow;
+
+    thread_start_row_ += (ThreadMap::Shape::kRow * increment);
+
+    // Group
+    state_[1] += increment_row;
+    int increment_group = state_[1] / ThreadMap::Count::kGroup;
+    state_[1] = state_[1] % ThreadMap::Count::kGroup;
+
+    thread_start_row_ +=
+        (ThreadMap::Shape::kGroup - 1) *
+        ThreadMap::Shape::kRow *
+        ThreadMap::Count::kRow *
+        increment_row;
+
+    // Cluster
+    state_[2] += increment_group;
+    int increment_cluster = state_[2] / ThreadMap::Count::kCluster;
+    state_[2] = state_[2] % ThreadMap::Count::kCluster;
+
+    thread_start_row_ +=
+        ThreadMap::Count::kGroup *
+        ThreadMap::Shape::kGroup *
+        ThreadMap::Count::kRow *
+        ThreadMap::Shape::kRow *
+        increment_group;
+
+    // Tile
+    thread_start_row_ +=
+        ThreadMap::Shape::kGroup *
+        ThreadMap::Shape::kRow *
+        ThreadMap::Shape::kCluster *
+        ThreadMap::Shape::kTile *
+        increment_cluster;
+
+    return *this;
+  }
+
+  ///< Efficiently disables all accesses guarded by mask
+  CUTLASS_DEVICE void clear_mask() {
+    mask_.clear();
+  }
+
+  ///< Efficiently enables all accesses guarded by mask
+  CUTLASS_DEVICE void enable_mask() {
+    mask_.enable();
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void get_mask(Mask &mask) const {
+    mask = mask_;
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void set_mask(Mask const &mask) {
+    mask_ = mask;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_direct_conv.h

@@ -0,0 +1,445 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+  \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
+
+  The epilogue rearranges the result of a matrix product through shared memory to match canonical
+  tensor layouts in global memory. Epilogues support conversion and reduction operations.
+
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/layout/tensor.h"
+#include "cutlass/layout/permute.h"
+#include "cutlass/matrix_shape.h"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/transform/pitch_linear_thread_map.h"
+#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
+#include "cutlass/arch/arch.h"
+#include "cutlass/arch/memory.h"
+#include "cutlass/epilogue/threadblock/predicated_tile_iterator_params.h"
+#include "cutlass/conv/conv2d_problem_size.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace epilogue {
+namespace threadblock {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Tile iterator used to load and store output tile from global memory in epilogue.
+///
+/// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator
+///
+template <
+  typename ThreadMap_,       ///< Thread map (conept: PitchLinearThreadMap)
+  typename Element_,         ///< Element data type
+  typename ThreadOutputShape_ = cutlass::conv::TensorNHWCShape<1, 1, 1, 1>,
+  typename ThreadBlockOutputShape_ = cutlass::conv::TensorNHWCShape<1, 1, 1, 1>
+>
+class PredicatedTileIteratorDirectConv {
+public:
+  using ThreadMap = ThreadMap_;
+  using Shape = typename ThreadMap::Shape;
+  using ThreadOutputShape = ThreadOutputShape_;
+  using ThreadBlockOutputShape = ThreadBlockOutputShape_;
+
+  using Element = Element_;
+
+  using Layout = layout::RowMajor;
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = MatrixCoord;
+
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+  static int const kThreads = ThreadMap::kThreads;
+
+  using ConvProblemSize = typename cutlass::conv::Conv2dProblemSize;
+
+  /// Fragment object
+  using Fragment = Array<Element, ThreadMap::Iterations::kCount * kElementsPerAccess>;
+
+  /// Memory access size
+  using AccessType = AlignedArray<Element, kElementsPerAccess>;
+
+  static int const kLoadsPerAccess = AccessType::kElements / AccessType::kElements;
+
+  using ThreadTileCount = MatrixShape<
+    ThreadBlockOutputShape::kH / ThreadOutputShape::kH,
+    ThreadBlockOutputShape::kW / ThreadOutputShape::kW
+  >;
+
+  //
+  // Parameters struct
+  //
+
+  /// Uses a non-template class
+  struct Params : PredicatedTileIteratorDirect2dConvParams {
+    using Base = PredicatedTileIteratorDirect2dConvParams;
+
+    CUTLASS_HOST_DEVICE
+    Params() { }
+
+    CUTLASS_HOST_DEVICE
+    Params(Layout const &layout, cutlass::conv::Conv2dProblemSize const &problem_size): 
+      PredicatedTileIteratorDirect2dConvParams(
+        layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess,
+        problem_size,
+        {ThreadBlockOutputShape::kH, ThreadBlockOutputShape::kW}
+      ) 
+    { }
+
+    CUTLASS_HOST_DEVICE
+    Params(Base const &base) : 
+      Base(base) { }
+  };
+
+  /// Mask object
+  struct Mask {
+
+    static int const kCount = ThreadMap::Iterations::kContiguous;
+
+    /// Predicate state
+    bool predicates[kCount];
+
+    //
+    // Mask
+    //
+    CUTLASS_HOST_DEVICE
+    Mask() {
+      enable();
+    }
+
+    ///< Efficiently disables all accesses guarded by mask
+    CUTLASS_HOST_DEVICE void clear() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = false;
+      }
+    }
+
+    ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask
+    CUTLASS_DEVICE void enable() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = true;
+      }
+    }
+  };
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Parameters structure containing reference and precomputed state.
+  PredicatedTileIteratorDirect2dConvParams params_;
+
+  /// Byte-level pointer
+  uint8_t *byte_pointer_;
+
+  ///     
+  Element *pointer_;
+
+
+  /// Array of boolean values to contain steady-state predicates
+  Mask mask_;
+
+  /// Extent of the matrix tile in rows
+  Index extent_row_;
+
+  /// Extent of the matrix tile in rows
+  Index extent_column_;
+
+  /// A thread's starting row position (assuming steady-state predicates have been computed)
+  Index thread_start_row_;
+
+  /// A thread's starting column
+  Index thread_start_column_;
+
+  /// Initial thread output location
+  int thread_start_n_, thread_start_p_, thread_start_q_;
+
+  /// Current threadblock tile index
+  int tile_index_;
+
+  //
+  // Static asserts about internal strides
+  //
+
+  static_assert(sizeof(extent_row_) == 4, "Expected 32b extents");
+  static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents");
+  static_assert(sizeof(PredicatedTileIteratorDirect2dConvParams::stride) == 8, "Expected 64b strides");
+
+private:
+
+  //
+  // Methods
+  //
+
+
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  CUTLASS_DEVICE
+  PredicatedTileIteratorDirectConv(
+    PredicatedTileIteratorDirect2dConvParams const & params,
+    Element *pointer,
+    TensorCoord extent,
+    int thread_idx,
+    TensorCoord threadblock_offset = TensorCoord()
+  ): 
+    params_(params), pointer_(pointer)
+  {
+
+    TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx);
+
+    extent_row_ = extent.row();
+    extent_column_ = extent.column();
+
+    // stride dim (PQ)
+    thread_start_row_ = thread_offset.column();
+    // contiguous dim (Channels)
+    thread_start_column_ = threadblock_offset.column() + thread_offset.row();
+
+    tile_index_ = threadblock_offset.row();
+
+    set_tile_index(0);
+  }
+
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void set_tile_index(const int index) { 
+   
+    int residual;
+    params_.pq_divmod(thread_start_n_, residual, tile_index_ + index);
+    params_.q_divmod(thread_start_p_, thread_start_q_, residual);
+
+    // Compute the base output coord of ThreadBlock
+    thread_start_p_ *= ThreadBlockOutputShape::kH;
+    thread_start_q_ *= ThreadBlockOutputShape::kW;
+
+    // Initialize predicates
+    CUTLASS_PRAGMA_UNROLL
+    for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) {
+      mask_.predicates[c] = ((thread_start_column_ 
+        + c * ThreadMap::Delta::kContiguous) < extent_column_);
+    }
+
+    // Null pointer performs no accesses
+    if (!pointer_) {
+      mask_.clear();
+    }
+
+  }
+
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    byte_pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load_with_byte_offset(Fragment &frag, int64_t byte_offset) const {
+    CUTLASS_PRAGMA_UNROLL
+    for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) {
+        int frag_base_idx = s * ThreadMap::Iterations::kContiguous + c;
+
+        int current_row = thread_start_row_ + s * ThreadMap::Delta::kStrided;
+        int p = current_row / ThreadBlockOutputShape::kW;
+        int q = current_row % ThreadBlockOutputShape::kW;
+
+        int current_p = thread_start_p_ + p;
+        int current_q = thread_start_q_ + q;
+
+        bool row_guard = (current_p) < params_.P && (current_q) < params_.Q &&
+                         (thread_start_n_ < params_.N) && current_row < ThreadMap::Shape::kStrided;
+
+        int output_row_offset =
+            thread_start_n_ * params_.stride_n + current_p * params_.stride_p + current_q;
+
+        uint8_t *byte_pointer =
+            reinterpret_cast<uint8_t *>(pointer_) +
+            LongIndex(output_row_offset) * LongIndex(params_.stride) +
+            LongIndex(thread_start_column_ + c * ThreadMap::Delta::kContiguous) *
+                sizeof(AccessType) / kElementsPerAccess;
+
+        AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+        AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset);
+
+        bool guard = row_guard && mask_.predicates[c];
+
+        cutlass::arch::global_load<AccessType, sizeof(AccessType)>(
+            frag_ptr[frag_base_idx], (void *)&memory_pointer[0], guard);
+      }
+    }
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load(Fragment &frag) const {
+    load_with_byte_offset(frag, 0);
+  }
+
+  /// Stores a fragment to memory
+  CUTLASS_DEVICE
+  void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) const {
+    CUTLASS_PRAGMA_UNROLL
+    for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) {
+        int frag_base_idx = s * ThreadMap::Iterations::kContiguous + c;
+
+        int current_row = thread_start_row_ + s * ThreadMap::Delta::kStrided;
+        int p = current_row / ThreadBlockOutputShape::kW;
+        int q = current_row % ThreadBlockOutputShape::kW;
+
+        int current_p = thread_start_p_ + p;
+        int current_q = thread_start_q_ + q;
+
+        bool row_guard = (current_p) < params_.P && (current_q) < params_.Q &&
+                         (thread_start_n_ < params_.N) && current_row < ThreadMap::Shape::kStrided;
+
+        int output_row_offset =
+            thread_start_n_ * params_.stride_n + current_p * params_.stride_p + current_q;
+
+        uint8_t *byte_pointer =
+            reinterpret_cast<uint8_t *>(pointer_) +
+            LongIndex(output_row_offset) * LongIndex(params_.stride) +
+            LongIndex(thread_start_column_ + c * ThreadMap::Delta::kContiguous) *
+                sizeof(AccessType) / kElementsPerAccess;
+
+        AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+        AccessType *memory_pointer = reinterpret_cast<AccessType *>(byte_pointer + byte_offset);
+
+        bool guard = row_guard && mask_.predicates[c];
+
+        cutlass::arch::global_store<AccessType, sizeof(AccessType)>(
+            frag_ptr[frag_base_idx], (void *)&memory_pointer[0], guard);
+      }
+    }
+  }
+
+  /// Stores a fragment to memory
+  CUTLASS_DEVICE
+  void store(Fragment const &frag) const {
+
+    store_with_byte_offset(frag, 0);
+  }
+
+  CUTLASS_DEVICE
+  MatrixCoord thread_start() const {
+    return MatrixCoord(thread_start_row_, thread_start_column_);
+  }
+
+  /// Need to get the thread start row from the tile iterator
+  CUTLASS_DEVICE
+  int32_t thread_start_row() const {
+    return thread_start_row_;
+  }
+
+  /// Need to get the thread start row from the tile iterator
+  CUTLASS_DEVICE
+  int32_t thread_start_column() const {
+    return thread_start_column_;
+  }
+
+  /// Extent of the matrix in rows
+  CUTLASS_DEVICE
+  Index extent_row() const {
+    return extent_row_;
+  }
+
+  /// Extent of the matrix in columns
+  CUTLASS_DEVICE
+  Index extent_column() const {
+    return extent_column_;
+  }
+
+  /// Advances to the next position to load or store
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorDirectConv &operator++() {
+    // do nothing
+
+    return *this;
+  }
+
+  ///< Efficiently disables all accesses guarded by mask
+  CUTLASS_DEVICE void clear_mask() {
+    mask_.clear();
+  }
+
+  ///< Efficiently enables all accesses guarded by mask
+  CUTLASS_DEVICE void enable_mask() {
+    mask_.enable();
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void get_mask(Mask &mask) const {
+    mask = mask_;
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void set_mask(Mask const &mask) {
+    mask_ = mask;
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_params.h

@@ -0,0 +1,483 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+  \brief 
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+
+#include "cutlass/layout/pitch_linear.h"
+#include "cutlass/layout/matrix.h"
+
+#include "cutlass/conv/conv2d_problem_size.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace threadblock {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct OutputTileShapeDesc {
+
+  int column;
+  int row;
+  int group;
+  int cluster;
+  int tile;
+
+  //
+  // Methods
+  //
+
+  /// Default ctor
+  CUTLASS_HOST_DEVICE
+  OutputTileShapeDesc(): column(0), row(0), group(0), cluster(0), tile(0) { }
+
+  /// Ctor
+  CUTLASS_HOST_DEVICE
+  OutputTileShapeDesc(
+    int column_,
+    int row_,
+    int group_,
+    int cluster_,
+    int tile_
+  ):
+    column(column_),
+    row(row_),
+    group(group_),
+    cluster(cluster_),
+    tile(tile_) { }
+
+  /// Total number of points in the 5D space
+  CUTLASS_HOST_DEVICE
+  int count() const {
+    return column * row * group * cluster * tile;
+  }
+
+  #if 0
+  CUTLASS_HOST_DEVICE
+  void print() const {
+    printf("{%d, %d, %d, %d, %d}", column, row, group, cluster, tile);
+  }
+  #endif
+};
+
+/// Helper template to construct an OutputTileShapeDesc from a OutputTileShape template.
+template <typename Shape>
+CUTLASS_HOST_DEVICE
+OutputTileShapeDesc make_OutputTileShapeDesc() {
+  return OutputTileShapeDesc(
+    Shape::kColumn,
+    Shape::kRow,
+    Shape::kGroup,
+    Shape::kCluster,
+    Shape::kTile
+  );
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Thread map description
+struct OutputTileThreadMapDesc {
+
+  int threads;
+  int elements_per_access;
+  OutputTileShapeDesc shape;
+  OutputTileShapeDesc iterations;
+  OutputTileShapeDesc delta;
+  OutputTileShapeDesc count;
+
+  //
+  // Methods
+  //
+
+  CUTLASS_HOST_DEVICE
+  OutputTileThreadMapDesc() { }
+
+  CUTLASS_HOST_DEVICE
+  OutputTileThreadMapDesc(
+    int threads_,
+    int elements_per_access_,
+    OutputTileShapeDesc shape_,
+    OutputTileShapeDesc iterations_,
+    OutputTileShapeDesc delta_,
+    OutputTileShapeDesc count_
+  ):
+    threads(threads_), 
+    elements_per_access(elements_per_access_),
+    shape(shape_),
+    iterations(iterations_),
+    delta(delta_),
+    count(count_) 
+  {
+    
+  }
+};
+
+/// Helper template to construct an OutputTileShapeDesc from a OutputTileThreadMap template.
+template <typename ThreadMap>
+CUTLASS_HOST_DEVICE
+OutputTileThreadMapDesc make_OutputTileThreadMapDesc() {
+  return OutputTileThreadMapDesc(
+    ThreadMap::kThreads,
+    ThreadMap::kElementsPerAccess,
+    make_OutputTileShapeDesc<typename ThreadMap::Shape>(),
+    make_OutputTileShapeDesc<typename ThreadMap::Iterations>(),
+    make_OutputTileShapeDesc<typename ThreadMap::Delta>(),
+    make_OutputTileShapeDesc<typename ThreadMap::Count>()
+  );
+}
+///////////////////////////////////////////////////////////////////////////////
+
+//
+// Parameters struct for PredicatedTileIterator
+//
+
+struct PredicatedTileIteratorParams {
+
+  using Index = int32_t;
+  using LongIndex = int64_t;
+
+  //
+  // Data members
+  //
+
+  LongIndex stride;               ///< stride in bytes between rows
+
+  LongIndex increment_row;        ///< increment quantity (in bytes) to advance when moving between rows
+  LongIndex increment_group;      ///< increment quantity (in bytes) to advance when moving to the next group
+  LongIndex increment_cluster;    ///< increment quantity (in bytes) to advance when moving to the next cluster
+
+  LongIndex advance_row;          ///< amount to add to move to the next 'row' position
+  LongIndex advance_group;        ///< amount to add to move to the next 'group' position
+  LongIndex advance_cluster;      ///< amount to add to move to the next 'cluster' position
+  LongIndex advance_tile;         ///< amount to add to move to the next 'tile'
+
+  //
+  // Methods
+  //
+
+  CUTLASS_HOST_DEVICE
+  Status initialize(LongIndex stride_, OutputTileThreadMapDesc thread_map) {
+    
+    stride = stride_;
+
+    increment_row = stride * thread_map.delta.row;
+
+    increment_group = stride * thread_map.delta.group
+      - stride * thread_map.delta.row * (thread_map.iterations.row - 1);
+
+    increment_cluster = stride * thread_map.delta.cluster
+      - stride * thread_map.delta.group * (thread_map.iterations.group - 1)
+      - stride * thread_map.delta.row * (thread_map.iterations.row - 1);
+
+    advance_row = stride * thread_map.shape.row;
+
+    advance_group = 
+      stride * 
+      (thread_map.shape.group - 1) * thread_map.shape.row * thread_map.count.row;
+    
+    advance_cluster = 
+      stride * 
+      thread_map.count.group * 
+      thread_map.shape.group * 
+      thread_map.count.row * 
+      thread_map.shape.row;
+    
+    advance_tile =
+      stride * 
+      thread_map.shape.group * 
+      thread_map.shape.row * 
+      thread_map.shape.cluster * 
+      thread_map.shape.tile;
+
+    return Status::kSuccess;
+  }
+
+  CUTLASS_HOST_DEVICE
+  Status initialize(Index stride_, OutputTileThreadMapDesc thread_map) {
+    return initialize(LongIndex(stride_), thread_map); 
+  }
+
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorParams() {
+    initialize(LongIndex(0), OutputTileThreadMapDesc());
+  }
+
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorParams(Index stride, OutputTileThreadMapDesc thread_map) {
+    initialize(stride, thread_map);
+  }
+
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorParams(LongIndex stride, OutputTileThreadMapDesc thread_map) {
+    initialize(stride, thread_map);
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+
+//
+// Parameters struct for PredicatedTileIteratorDirect2dConv
+//
+
+struct PredicatedTileIteratorDirect2dConvParams{
+  using Index = int32_t;
+  using LongIndex = int64_t;
+
+  //
+  // Data members
+  //
+  FastDivmod pq_divmod;
+  FastDivmod q_divmod;
+
+  LongIndex stride;
+  LongIndex stride_n;
+  LongIndex stride_p;
+
+  int N;
+  int P;
+  int Q;
+
+  //
+  // Methods
+  //
+
+  CUTLASS_HOST_DEVICE
+  Status initialize(LongIndex stride_,
+                    cutlass::conv::Conv2dProblemSize const &problem_size,
+                    MatrixCoord threadblock_output_shape) {
+    stride = stride_; // The stride per row of output tensor (bytes)
+    stride_n = problem_size.P * problem_size.Q;
+    stride_p = problem_size.Q ;
+
+    N = problem_size.N;
+    P = problem_size.P;
+    Q = problem_size.Q;
+
+    // Fastdivmod for output O, P, Q
+    if(threadblock_output_shape.row() != 0 && threadblock_output_shape.column() !=0 ){
+      // MSVC emits a "potential divide by 0" warning as error
+      // if the code just divides without a check and substitution.
+
+      CUTLASS_ASSERT(threadblock_output_shape.row() != 0);
+      const auto row_denom = threadblock_output_shape.row() != 0 ?
+        threadblock_output_shape.row() : cutlass::MatrixCoord::Index(1);
+      int tiles_p =
+          (problem_size.P + (threadblock_output_shape.row() - 1)) / row_denom;
+
+      CUTLASS_ASSERT(threadblock_output_shape.column() != 0);
+      const auto col_denom = threadblock_output_shape.column() != 0 ?
+        threadblock_output_shape.column() : cutlass::MatrixCoord::Index(1);
+      int tiles_q = (problem_size.Q + (threadblock_output_shape.column() - 1)) /
+                    col_denom;
+
+      pq_divmod = FastDivmod(tiles_p * tiles_q);
+      q_divmod = FastDivmod(tiles_q);
+    }
+
+    return Status::kSuccess;
+  }
+
+  CUTLASS_HOST_DEVICE
+  Status initialize(
+      Index stride_,
+      cutlass::conv::Conv2dProblemSize const &problem_size = cutlass::conv::Conv2dProblemSize(),
+      MatrixCoord threadblock_output_shape = MatrixCoord()) {
+    return initialize(LongIndex(stride_), problem_size, threadblock_output_shape);
+  }
+
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorDirect2dConvParams() { initialize(LongIndex(0)); }
+
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorDirect2dConvParams(Index stride,
+                               cutlass::conv::Conv2dProblemSize const &problem_size,
+                               MatrixCoord threadblock_output_shape) {
+    initialize(stride, problem_size, threadblock_output_shape);
+  }
+
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorDirect2dConvParams(LongIndex stride,
+                               cutlass::conv::Conv2dProblemSize const &problem_size,
+                               MatrixCoord threadblock_output_shape) {
+    initialize(stride, problem_size, threadblock_output_shape);
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+//  InterleavedPredicatedTileIterator
+///////////////////////////////////////////////////////////////////////////////
+
+
+/// Predicated tile access iterator descriptor object containing template dependent state
+struct InterleavedPredicatedTileIteratorDesc {
+
+  int element_size_bits;
+  int elements_per_access;
+  int threadmap_warp_size;
+  layout::PitchLinearCoord threadmap_iterations;
+  layout::PitchLinearCoord threadmap_delta;
+
+  //
+  // Methods
+  //
+
+  CUTLASS_HOST_DEVICE
+  InterleavedPredicatedTileIteratorDesc() { }
+
+  CUTLASS_HOST_DEVICE
+  InterleavedPredicatedTileIteratorDesc(
+    int element_size_bits_,
+    int elements_per_access_,
+    int threadmap_warp_size_,
+    layout::PitchLinearCoord threadmap_iterations_,
+    layout::PitchLinearCoord threadmap_delta_
+  ):
+    element_size_bits(element_size_bits_),
+    elements_per_access(elements_per_access_),
+    threadmap_warp_size(threadmap_warp_size_),
+    threadmap_iterations(threadmap_iterations_),
+    threadmap_delta(threadmap_delta_) { }
+};
+
+//
+// Parameters struct InterleavedPredicatedTileIterator
+//
+
+struct InterleavedPredicatedTileIteratorParams {
+
+  using Index = int32_t;
+  using LongIndex = int64_t;
+
+  //
+  // Data members
+  //
+
+  LongIndex stride;               ///< stride in bytes between rows
+  LongIndex advance_row;          ///< amount to add to move to the next 'row' position
+  LongIndex advance_column;       ///< amount to add to move to the next 'column' position
+
+  //
+  // Methods
+  //
+
+  CUTLASS_HOST_DEVICE
+  Status initialize(LongIndex stride_, InterleavedPredicatedTileIteratorDesc desc) {
+    
+    stride = stride_;
+
+    advance_row = desc.threadmap_delta.contiguous() * desc.element_size_bits / 8;
+
+    advance_column = stride_ - desc.threadmap_iterations.contiguous() *
+                               desc.elements_per_access *
+                               desc.element_size_bits *
+                               desc.threadmap_warp_size / 8;
+
+    return Status::kSuccess;
+  }
+
+  CUTLASS_HOST_DEVICE
+  InterleavedPredicatedTileIteratorParams() {
+    initialize(LongIndex(0), InterleavedPredicatedTileIteratorDesc());
+  }
+
+  CUTLASS_HOST_DEVICE
+  InterleavedPredicatedTileIteratorParams(Index stride, InterleavedPredicatedTileIteratorDesc desc) {
+    initialize(stride, desc);
+  }
+
+  CUTLASS_HOST_DEVICE
+  InterleavedPredicatedTileIteratorParams(LongIndex stride, InterleavedPredicatedTileIteratorDesc desc) {
+    initialize(stride, desc);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+/// Helper template to construct an OutputTileShapeDesc from a OutputTileThreadMap template.
+template <typename Element, typename ThreadMap>
+CUTLASS_HOST_DEVICE
+InterleavedPredicatedTileIteratorDesc make_InterleavedPredicatedTileIteratorDesc() {
+  return InterleavedPredicatedTileIteratorDesc(
+    sizeof_bits<Element>::value,
+    ThreadMap::kElementsPerAccess,
+    ThreadMap::kWarpSize,
+    {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided},
+    {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided}
+  );
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+/// Helper template to construct an MakePredicatedTileIteratorDesc from a template 
+// dependent state
+template <typename Element, typename Layout,
+   typename ThreadMap>
+  struct MakePredicatedTileIteratorDesc;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Specialization of PredicatedTileAccessIterator for layout::RowMajor output data.
+template <typename Element, typename ThreadMap>
+struct MakePredicatedTileIteratorDesc <
+    Element, layout::RowMajor, ThreadMap> {
+
+  CUTLASS_HOST_DEVICE
+  OutputTileThreadMapDesc operator()() {
+
+    return make_OutputTileThreadMapDesc<ThreadMap>();
+  }
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Specialization of PredicatedTileAccessIterator for layout::ColumnMajorInterleaved<InterleavedN> output data.
+template <typename Element, typename ThreadMap, int InterleavedN>
+struct MakePredicatedTileIteratorDesc <
+    Element, layout::ColumnMajorInterleaved<InterleavedN>, ThreadMap> {
+
+  CUTLASS_HOST_DEVICE
+  InterleavedPredicatedTileIteratorDesc operator()() {
+
+    return make_InterleavedPredicatedTileIteratorDesc<Element, ThreadMap>();
+  }
+};
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_predicates.h

@@ -0,0 +1,309 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+  \brief PredicatedTileIteratorPredicates.
+
+  PredicatedTileIteratorPredicates enables both upper and lower bounds for predicates.
+
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/layout/tensor.h"
+#include "cutlass/matrix_shape.h"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/transform/pitch_linear_thread_map.h"
+#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
+#include "cutlass/arch/arch.h"
+#include "cutlass/arch/memory.h"
+#include "cutlass/epilogue/threadblock/predicated_tile_iterator_params.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace epilogue {
+namespace threadblock {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Tile iterator predicates used to bound computations in epilogue.
+///
+/// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator
+///
+template <
+  typename ThreadMap_,       ///< Thread map (conept: OutputTileThreadMap)
+  typename Element_          ///< Element data type
+>
+class PredicatedTileIteratorPredicates {
+public:
+  using ThreadMap = ThreadMap_;
+  using Shape = typename ThreadMap::Shape;
+
+  using Element = Element_;
+
+  using Layout = layout::RowMajor;
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = MatrixCoord;
+
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+  static int const kThreads = ThreadMap::kThreads;
+  static int const kIterations = ThreadMap::Count::kTile;
+
+  static_assert( ThreadMap::Iterations::kRow > 0,"ThreadMap::Iterations::kRow must be > 0");
+  static_assert( ThreadMap::Iterations::kGroup > 0,"ThreadMap::Iterations::kGroup must be > 0");
+  static_assert( ThreadMap::Iterations::kCluster > 0,"ThreadMap::Iterations::kCluster must be > 0");
+  static_assert( ThreadMap::Iterations::kColumn > 0,"ThreadMap::Iterations::kColumn must be > 0");
+
+  /// Fragment object
+  using Fragment = Array<
+    Element, 
+    ThreadMap::Iterations::kColumn * 
+    ThreadMap::Iterations::kRow * 
+    ThreadMap::Iterations::kGroup * 
+    ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>;
+
+  /// Memory access size
+  using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess>;
+
+  //
+  // Parameters struct
+  //
+
+  /// Uses a non-template class
+  struct Params : PredicatedTileIteratorParams {
+
+    CUTLASS_HOST_DEVICE
+    Params() { }
+
+    CUTLASS_HOST_DEVICE
+    Params(Layout const &layout): 
+      PredicatedTileIteratorParams(
+        layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess,
+        make_OutputTileThreadMapDesc<ThreadMap>()
+      ) 
+    {
+        
+    }
+  };
+
+  /// Mask object
+  struct Mask {
+
+    static int const kCount = ThreadMap::Iterations::kColumn;
+
+    /// Predicate state
+    bool predicates[kCount];
+
+    //
+    // Mask
+    //
+    CUTLASS_HOST_DEVICE
+    Mask() {
+      enable();
+    }
+
+    ///< Efficiently disables all accesses guarded by mask
+    CUTLASS_HOST_DEVICE void clear() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = false;
+      }
+    }
+
+    ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask
+    CUTLASS_DEVICE void enable() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = true;
+      }
+    }
+  };
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Parameters structure containing reference and precomputed state.
+  PredicatedTileIteratorParams params_;
+
+  /// Array of boolean values to contain steady-state predicates
+  Mask mask_;
+
+  /// Extent of the matrix tile in rows
+  Index lower_extent_row_;
+  Index upper_extent_row_;
+
+  /// A thread's starting row position (assuming steady-state predicates have been computed)
+  Index thread_start_row_;
+
+  /// Internal state counter
+  int state_[3];
+ 
+  //
+  // Static asserts about internal strides
+  //
+
+  static_assert(sizeof(lower_extent_row_) == 4, "Expected 32b extents");
+  static_assert(sizeof(upper_extent_row_) == 4, "Expected 32b extents");
+  static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents");
+  static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides");
+
+private:
+
+  //
+  // Methods
+  //
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  CUTLASS_DEVICE
+  PredicatedTileIteratorPredicates(
+    PredicatedTileIteratorParams const & params,
+    TensorCoord lower_extent,
+    TensorCoord upper_extent,
+    int thread_idx,
+    TensorCoord threadblock_offset = TensorCoord()
+  ): 
+    params_(params)
+  {
+
+    TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset;
+
+    lower_extent_row_ = lower_extent.row();
+    upper_extent_row_ = upper_extent.row();
+    thread_start_row_ = thread_offset.row();
+
+    // Initialize predicates
+    CUTLASS_PRAGMA_UNROLL
+    for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) {
+
+      mask_.predicates[c] = ((thread_offset.column() 
+        + ThreadMap::Delta::kColumn * c) < upper_extent.column()) &&
+        ((thread_offset.column() + ThreadMap::Delta::kColumn * c) >= lower_extent.column());
+    }
+
+    // Initialize internal state counter
+    state_[0] = state_[1] = state_[2] = 0;
+  }
+
+  /// Advances to the next position to load or store
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorPredicates &operator++() {
+
+    ++state_[0];
+    thread_start_row_ += ThreadMap::Shape::kRow;
+
+    if (state_[0] == ThreadMap::Count::kRow) {
+
+      state_[0] = 0;
+      ++state_[1];
+
+      thread_start_row_ += (ThreadMap::Shape::kGroup - 1) *
+        ThreadMap::Shape::kRow * ThreadMap::Count::kRow;
+
+      if (state_[1] == ThreadMap::Count::kGroup) {
+
+        state_[1] = 0;
+        ++state_[2];
+
+        thread_start_row_ += ThreadMap::Count::kGroup *
+          ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow;
+
+        if (state_[2] == ThreadMap::Count::kCluster) {
+          state_[2] = 0;
+        }
+      }
+    }
+
+    return *this;
+  }
+
+  ///< Efficiently disables all accesses guarded by mask
+  CUTLASS_DEVICE void clear_mask() {
+    mask_.clear();
+  }
+
+  ///< Efficiently enables all accesses guarded by mask
+  CUTLASS_DEVICE void enable_mask() {
+    mask_.enable();
+  }
+
+  ///< Gets the mask
+  CUTLASS_DEVICE void get_mask(Mask &mask) {
+    mask = mask_;
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void set_mask(Mask const &mask) {
+    mask_ = mask;
+  }
+
+  ///< Gets lower_extent_row_
+  CUTLASS_DEVICE Index get_lower_extent_row() {
+    return lower_extent_row_;
+  }
+
+  ///< Gets upper_extent_row_
+  CUTLASS_DEVICE Index get_upper_extent_row() {
+    return upper_extent_row_;
+  }
+
+  ///< Gets thread_start_row_
+  CUTLASS_DEVICE Index get_thread_start_row() {
+    return thread_start_row_;
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/predicated_tile_iterator_strided_dgrad.h

@@ -0,0 +1,479 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+  \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
+
+  The epilogue rearranges the result of a matrix product through shared memory to match canonical
+  tensor layouts in global memory. Epilogues support conversion and reduction operations.
+
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/layout/tensor.h"
+#include "cutlass/matrix_shape.h"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/transform/pitch_linear_thread_map.h"
+#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
+#include "cutlass/arch/arch.h"
+#include "cutlass/arch/memory.h"
+#include "cutlass/conv/conv2d_problem_size.h"
+#include "cutlass/epilogue/threadblock/predicated_tile_iterator_params.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace epilogue {
+namespace threadblock {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Tile iterator used to load and store output tile from global memory in epilogue.
+///
+/// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator
+///
+template <
+  typename ThreadMap_,       ///< Thread map (conept: OutputTileThreadMap)
+  typename Element_          ///< Element data type
+>
+class PredicatedTileIteratorStridedDgrad {
+public:
+  using ThreadMap = ThreadMap_;
+  using Shape = typename ThreadMap::Shape;
+
+  using Element = Element_;
+
+  using Layout = layout::RowMajor;
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = MatrixCoord;
+
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+  static int const kThreads = ThreadMap::kThreads;
+  static int const kIterations = ThreadMap::Count::kTile;
+
+  static_assert( ThreadMap::Iterations::kRow > 0,"ThreadMap::Iterations::kRow must be > 0");
+  static_assert( ThreadMap::Iterations::kGroup > 0,"ThreadMap::Iterations::kGroup must be > 0");
+  static_assert( ThreadMap::Iterations::kCluster > 0,"ThreadMap::Iterations::kCluster must be > 0");
+  static_assert( ThreadMap::Iterations::kColumn > 0,"ThreadMap::Iterations::kColumn must be > 0");
+
+  /// Fragment object
+  using Fragment = Array<
+    Element, 
+    ThreadMap::Iterations::kColumn * 
+    ThreadMap::Iterations::kRow * 
+    ThreadMap::Iterations::kGroup * 
+    ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>;
+
+  /// Memory access size
+  using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess>;
+
+  //
+  // Parameters struct
+  //
+
+  /// Uses a non-template class
+  struct Params : PredicatedTileIteratorParams {
+
+    /// Convolution problem size
+    cutlass::conv::Conv2dProblemSize problem_size;
+    int tiled_rows_per_filter;
+
+    CUTLASS_HOST_DEVICE
+    Params() { }
+
+    CUTLASS_HOST_DEVICE
+    Params(Layout const &layout, cutlass::conv::Conv2dProblemSize problem_size_, int threadblock_row): 
+      problem_size(problem_size_), 
+      PredicatedTileIteratorParams(
+        layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess,
+        make_OutputTileThreadMapDesc<ThreadMap>()
+      ) 
+    {
+  
+      int tile_m_per_filter = strided_dgrad_tile_m_per_filter(problem_size, threadblock_row);
+
+      tiled_rows_per_filter = tile_m_per_filter * threadblock_row;
+    }
+  };
+
+  /// Mask object
+  struct Mask {
+
+    static int const kCount = ThreadMap::Iterations::kColumn;
+
+    /// Predicate state
+    bool predicates[kCount];
+
+    //
+    // Mask
+    //
+    CUTLASS_HOST_DEVICE
+    Mask() {
+      enable();
+    }
+
+    ///< Efficiently disables all accesses guarded by mask
+    CUTLASS_HOST_DEVICE void clear() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = false;
+      }
+    }
+
+    ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask
+    CUTLASS_DEVICE void enable() {
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kCount; ++i) {
+        predicates[i] = true;
+      }
+    }
+  };
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Parameters structure containing reference and precomputed state.
+  Params params_;
+
+  /// Byte-level pointer
+  uint8_t *byte_pointer_;
+
+  /// Array of boolean values to contain steady-state predicates
+  Mask mask_;
+
+  /// Extent of the matrix tile in rows
+  Index extent_row_;
+
+  /// Starting Dx h and w dimension for strided dgrad mapping
+  int start_h_, start_w_;
+
+  /// Effective Dy P and Q dimensions for strided dgrad mapping
+  int p_, q_;
+
+  /// A thread's starting row position (assuming steady-state predicates have been computed)
+  Index thread_start_row_;
+
+  /// A thread's starting column position (assuming steady-state predicates have been computed)
+  Index thread_start_column_;
+
+  /// Internal state counter
+  int state_[3];
+ 
+  //
+  // Static asserts about internal strides
+  //
+
+  static_assert(sizeof(extent_row_) == 4, "Expected 32b extents");
+  static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents");
+  static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides");
+
+private:
+
+  //
+  // Methods
+  //
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  CUTLASS_DEVICE
+  PredicatedTileIteratorStridedDgrad(
+    Params const & params,
+    Element *pointer,
+    TensorCoord extent,
+    int thread_idx,
+    FastDivmod const &stride_h_divmod, FastDivmod const &stride_w_divmod,
+    int start_r, int start_s,
+    TensorCoord threadblock_offset = TensorCoord()
+  ): 
+    params_(params)
+  {
+
+    TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset;
+
+    int r = start_r;
+    int s = start_s;
+
+    if (params_.problem_size.mode == cutlass::conv::Mode::kConvolution) {
+      r = (params_.problem_size.R - 1 - r);
+      s = (params_.problem_size.S - 1 - s);
+    }
+
+    // compute starting coordinates in Dx start_h_ and start_w_
+    strided_dgrad_starting_coords(
+      params_.problem_size, 
+      stride_h_divmod, stride_w_divmod, 
+      r, s, 
+      start_h_, start_w_);
+
+    p_ = (params_.problem_size.H - start_h_ + params_.problem_size.stride_h - 1) / params_.problem_size.stride_h;
+    q_ = (params_.problem_size.W - start_w_ + params_.problem_size.stride_w - 1) / params_.problem_size.stride_w;
+
+    extent_row_ = extent.row();
+    thread_start_row_ = thread_offset.row();
+    thread_start_column_ = thread_offset.column();
+
+    // Initialize predicates
+    CUTLASS_PRAGMA_UNROLL
+    for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) {
+
+      mask_.predicates[c] = ((thread_offset.column() 
+        + ThreadMap::Delta::kColumn * c) < extent.column());
+    }
+
+    // Null pointer performs no accesses
+    if (!pointer) {
+      mask_.clear();
+    }
+
+    // Initialize pointer
+    byte_pointer_ = reinterpret_cast<uint8_t *>(pointer);
+
+    // Initialize internal state counter
+    state_[0] = state_[1] = state_[2] = 0;
+  }
+
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    byte_pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load_with_byte_offset(Fragment &frag, int64_t byte_offset) {
+
+    uint8_t *byte_pointer = byte_pointer_;
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int frag_row_idx = 
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          int row_offset = row * ThreadMap::Delta::kRow 
+            + group * ThreadMap::Delta::kGroup 
+            + cluster * ThreadMap::Delta::kCluster;
+
+          // remapping rows to find the mapped_row_offset
+          int npq_offset = (row_offset + thread_start_row_) % params_.tiled_rows_per_filter;
+
+          // (STEP 4.a) [order NHW rows to be loaded and stored in output Dx NHWxC layout]
+          int n = npq_offset / (p_ * q_); 
+          int residual = npq_offset % (p_ * q_);
+          int p = residual / q_;
+          int q = residual % q_;
+        
+          int mapped_row_offset = n * (params_.problem_size.H * params_.problem_size.W) +
+                                  (start_h_ + p * params_.problem_size.stride_h) * params_.problem_size.W +
+                                  (start_w_ + q * params_.problem_size.stride_w);
+          bool row_guard = mapped_row_offset < extent_row_;
+
+          int64_t row_byte_offset = mapped_row_offset * params_.stride;
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+
+            int64_t column_byte_offset = (thread_start_column_ + column * ThreadMap::Delta::kColumn) * (sizeof_bits<Element>::value / 8);
+
+            bool guard = row_guard && mask_.predicates[column];
+
+            cutlass::arch::global_load<
+              AccessType, 
+              sizeof(AccessType)
+            >(
+                frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn +
+                         column],
+                (void *)(byte_pointer + row_byte_offset + column_byte_offset + byte_offset),
+                guard);
+          }
+        }
+      }
+    }
+  }
+
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load(Fragment &frag) {
+
+    load_with_byte_offset(frag, 0);
+  }
+
+  /// Stores a fragment to memory
+  CUTLASS_DEVICE
+  void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) {
+    uint8_t *byte_pointer = byte_pointer_;
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int frag_row_idx = 
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          int row_offset = row * ThreadMap::Delta::kRow 
+            + group * ThreadMap::Delta::kGroup 
+            + cluster * ThreadMap::Delta::kCluster;
+
+          // remapping rows to find the mapped_row_offset
+          int npq_offset = (row_offset + thread_start_row_) % params_.tiled_rows_per_filter;
+
+          // (STEP 4.a) [order NHW rows to be loaded and stored in output Dx NHWxC layout]
+          int n = npq_offset / (p_ * q_); 
+          int residual = npq_offset % (p_ * q_);
+          int p = residual / q_;
+          int q = residual % q_;
+        
+          int mapped_row_offset = n * (params_.problem_size.H * params_.problem_size.W) +
+                                  (start_h_ + p * params_.problem_size.stride_h) * params_.problem_size.W +
+                                  (start_w_ + q * params_.problem_size.stride_w);
+          bool row_guard = mapped_row_offset < extent_row_;
+
+          int64_t row_byte_offset = mapped_row_offset * params_.stride;
+          
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+
+            int64_t column_byte_offset = (thread_start_column_ + column * ThreadMap::Delta::kColumn) * (sizeof_bits<Element>::value / 8);
+
+            bool guard = row_guard && mask_.predicates[column];
+
+            cutlass::arch::global_store<AccessType, sizeof(AccessType) >(
+                frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column],
+                (void *)(byte_pointer + row_byte_offset + column_byte_offset + byte_offset),
+                guard);            
+          }
+        }
+      }
+    }
+  }
+
+
+  /// Stores a fragment to memory
+  CUTLASS_DEVICE
+  void store(Fragment const &frag) {
+
+    store_with_byte_offset(frag, 0);
+  }
+
+  /// Advances to the next position to load or store
+  CUTLASS_HOST_DEVICE
+  PredicatedTileIteratorStridedDgrad &operator++() {
+
+    ++state_[0];
+
+    thread_start_row_ += ThreadMap::Shape::kRow;
+    
+    if (state_[0] == ThreadMap::Count::kRow) {
+
+      state_[0] = 0;
+      ++state_[1];
+
+      thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * 
+        ThreadMap::Shape::kRow * ThreadMap::Count::kRow;
+
+      if (state_[1] == ThreadMap::Count::kGroup) {
+
+        state_[1] = 0;
+        ++state_[2];
+
+        thread_start_row_ += ThreadMap::Count::kGroup * 
+          ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow;
+
+        if (state_[2] == ThreadMap::Count::kCluster) {
+          state_[2] = 0;
+        }
+      }
+    }
+
+    return *this;
+  }
+
+  ///< Efficiently disables all accesses guarded by mask
+  CUTLASS_DEVICE void clear_mask() {
+    mask_.clear();
+  }
+
+  ///< Efficiently enables all accesses guarded by mask
+  CUTLASS_DEVICE void enable_mask() {
+    mask_.enable();
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void get_mask(Mask &mask) {
+    mask = mask_;
+  }
+
+  ///< Sets the mask
+  CUTLASS_DEVICE void set_mask(Mask const &mask) {
+    mask_ = mask;
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/shared_load_iterator.h

@@ -0,0 +1,223 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+  \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
+
+  The epilogue rearranges the result of a matrix product through shared memory to match canonical
+  tensor layouts in global memory. Epilogues support conversion and reduction operations.
+
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/matrix_shape.h"
+#include "cutlass/tensor_ref.h"
+
+#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace threadblock {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tile iterator used to load output tile from shared memory in epilogue.
+///
+/// Satisfies: ReadableTileIterator
+///
+template <
+  typename ThreadMap_,       ///< Thread map (conept: OutputTileThreadMap)
+  typename Element_,         ///< Element data type
+  int MaxAlignment = ThreadMap_::kElementsPerAccess * sizeof_bits<Element_>::value / 8
+>
+class SharedLoadIterator {
+public:
+  using ThreadMap = ThreadMap_;
+  using Shape = typename ThreadMap::TileShape;
+
+  using Element = Element_;
+
+  using Layout = layout::RowMajor;
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = MatrixCoord;
+
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+
+  static int const kMinAlignment = ThreadMap_::kElementsPerAccess * sizeof_bits<Element_>::value / 8;
+
+  static int const kAlignment = (MaxAlignment < kMinAlignment ? MaxAlignment : kMinAlignment);
+
+  static int const kThreads = ThreadMap::kThreads;
+
+  /// Fragment object
+  using Fragment = Array<
+    Element, 
+    ThreadMap::Iterations::kColumn * 
+    ThreadMap::Iterations::kRow * 
+    ThreadMap::Iterations::kGroup * 
+    ThreadMap::Iterations::kCluster * 
+    ThreadMap::kElementsPerAccess>;
+
+  /// Memory access size
+  using AccessType = AlignedArray<
+    Element, 
+    ThreadMap::kElementsPerAccess, 
+    kAlignment>;
+
+  /// Vector type used for SMEM loads
+  using LoadType = AlignedArray<
+    Element,
+    const_min(128 / sizeof_bits<Element>::value, ThreadMap::kElementsPerAccess),
+    const_min(16, kAlignment)
+  >;
+
+  static int const kLoadsPerAccess = AccessType::kElements / LoadType::kElements;
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Byte-level pointer
+  uint8_t *byte_pointer_;
+
+  /// Stride along adjacent rows
+  int stride_;
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  CUTLASS_DEVICE
+  SharedLoadIterator(
+    TensorRef ref,
+    int thread_idx
+  ):
+    byte_pointer_(reinterpret_cast<uint8_t *>(ref.data())),
+    stride_((ref.stride(0) * sizeof_bits<Element>::value) / 8) {
+
+    TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx);
+
+    // Initialize pointer
+    byte_pointer_ +=
+      thread_offset.row() * stride_ + 
+      thread_offset.column() * sizeof(AccessType) / kElementsPerAccess;
+  }
+
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    byte_pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
+  }
+
+  CUTLASS_DEVICE
+  void add_tile_offset(TensorCoord const &offset) {
+    byte_pointer_ += 
+      offset.row() * Shape::kRow * stride_ + 
+      offset.column() * Shape::kColumn * sizeof_bits<Element>::value / 8;
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const {
+
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          uint8_t const *byte_pointer = byte_pointer_ + 
+            row * ThreadMap::Delta::kRow * stride_ + 
+            group * ThreadMap::Delta::kGroup* stride_ + 
+            cluster * ThreadMap::Delta::kCluster * stride_ +
+            pointer_offset * sizeof_bits<Element>::value / 8;
+
+          int frag_row_idx = 
+            (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          LoadType *frag_ptr = reinterpret_cast<LoadType *>(&frag);
+          LoadType const *memory_pointer = reinterpret_cast<LoadType const *>(byte_pointer);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+            
+            int frag_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column;
+
+            CUTLASS_PRAGMA_UNROLL
+            for (int v = 0; v < kLoadsPerAccess; ++v) {
+              frag_ptr[frag_idx * kLoadsPerAccess + v] = 
+                memory_pointer[(column * ThreadMap::Delta::kColumn / kElementsPerAccess) * kLoadsPerAccess + v];
+            }
+          }
+        }
+      }
+    }
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void set_smem_base_address(Index address) {
+  }
+
+  /// Loads a fragment
+  CUTLASS_DEVICE
+  void load(Fragment &frag) const {
+
+    load_with_pointer_offset(frag, 0);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace epilogue
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/shared_load_iterator_mixed.h

@@ -0,0 +1,594 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+  \brief Epilogue for threadblock scoped GEMMs using Tensor Ops optimized for mixed-precision.
+
+  This assumes the shared memory tile is in a permuted layout which avoids bank conflicts on loading.
+
+  When the fragment is loaded into registers, it matches the row-major thread map assumed by
+  the predicated tile iterator writing to global memory.
+
+  The epilogue rearranges the result of a matrix product through shared memory to match canonical
+  tensor layouts in global memory. Epilogues support conversion and reduction operations.
+
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/matrix_shape.h"
+#include "cutlass/tensor_ref.h"
+
+#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace threadblock {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tile iterator used to load output tile from shared memory in epilogue.
+///
+/// Satisfies: ReadableTileIterator
+///
+template <
+  typename ThreadMap_,       ///< Thread map (conept: OutputTileThreadMap)
+  typename Element_,         ///< Accumulator data type
+  int ElementSizeBits_,      ///< Size of accumulator in bits
+  int OutputSizeBits_,       ///< Size of output element in bits
+  int ElementsPerAccess,     ///< Vector length of output vector
+  int ContiguousLanes,       ///< Number of lanes in the warp writing to contiguous elements
+                             ///  in the global memory tensor
+  bool EightBitsOutputOrLess = (OutputSizeBits_ <= 8)
+>
+class SharedLoadIteratorMixed;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tile iterator used to load output tile from shared memory in epilogue.
+///
+/// Satisfies: ReadableTileIterator
+///
+template <
+  typename ThreadMap_,       ///< Thread map (conept: OutputTileThreadMap)
+  typename Element_          ///< Accumulator data type
+>
+class SharedLoadIteratorMixed<ThreadMap_, Element_, 32, 16, 8, 8, false> {
+public:
+  using ThreadMap = ThreadMap_;
+  using Shape = typename ThreadMap::Shape;
+
+  using Element = Element_;
+
+  using Layout = layout::RowMajor;
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = MatrixCoord;
+
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+
+  static int const kAlignment = ThreadMap::kElementsPerAccess * sizeof_bits<Element_>::value / 8;
+
+  static int const kThreads = ThreadMap::kThreads;
+
+  /// Fragment object
+  using Fragment = Array<
+    Element, 
+    ThreadMap::Iterations::kColumn * 
+    ThreadMap::Iterations::kRow * 
+    ThreadMap::Iterations::kGroup * 
+    ThreadMap::Iterations::kCluster * 
+    ThreadMap::kElementsPerAccess>;
+
+  /// Memory access size
+  using AccessType = AlignedArray<
+    Element, 
+    ThreadMap::kElementsPerAccess, 
+    kAlignment>;
+
+  /// Vector type used for SMEM loads
+  using LoadType = AlignedArray<
+    Element,
+    const_min(128 / sizeof_bits<Element>::value, ThreadMap::kElementsPerAccess),
+    const_min(16, kAlignment)
+  >;
+
+  static int const kLoadsPerAccess = AccessType::kElements / LoadType::kElements;
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Byte-level pointer
+  LoadType const *pointers_[kLoadsPerAccess];
+
+  /// Stride along adjacent rows in units of LoadType
+  int stride_;
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  CUTLASS_DEVICE
+  SharedLoadIteratorMixed(
+    TensorRef ref,
+    int thread_idx
+  ):
+    stride_((ref.stride(0) / LoadType::kElements)) {
+
+    TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx);
+
+    // Initialize pointers
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      pointers_[i] = reinterpret_cast<LoadType const *>(ref.data());
+
+      int col_idx = (thread_offset.column() / kElementsPerAccess) * kLoadsPerAccess;
+      int bank_offset = (col_idx * int(sizeof(LoadType)) / 128) % kLoadsPerAccess;
+
+      col_idx += (bank_offset + i) % kLoadsPerAccess;
+
+      pointers_[i] += thread_offset.row() * stride_ + col_idx;
+    }
+  }
+
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      pointers_[i] += pointer_offset / LoadType::kElements;
+    }
+  }
+
+  CUTLASS_DEVICE
+  void add_tile_offset(TensorCoord const &offset) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      pointers_[i] += 
+        offset.row() * Shape::kRow * stride_ + 
+        offset.column() * Shape::kColumn / LoadType::kElements;
+    }
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const {
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int row_ptr_offset =
+            row * ThreadMap::Delta::kRow * stride_ + 
+            group * ThreadMap::Delta::kGroup* stride_ + 
+            cluster * ThreadMap::Delta::kCluster * stride_ +
+            pointer_offset / LoadType::kElements;
+
+          int frag_row_idx = (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          LoadType *frag_ptr = reinterpret_cast<LoadType *>(&frag);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+            
+            int frag_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column;
+
+            CUTLASS_PRAGMA_UNROLL
+            for (int v = 0; v < kLoadsPerAccess; ++v) {
+           
+              int vector_idx = (column * ThreadMap::Delta::kColumn / kElementsPerAccess * kLoadsPerAccess); 
+
+              LoadType const *memory_pointer = pointers_[v] + row_ptr_offset;
+            
+              frag_ptr[frag_idx * kLoadsPerAccess + v] = memory_pointer[vector_idx];
+            }
+          }
+        }
+      }
+    }
+  }
+
+  /// Set base smem address
+  CUTLASS_DEVICE
+  void set_smem_base_address(Index address) {}
+
+  /// Loads a fragment
+  CUTLASS_DEVICE
+  void load(Fragment &frag) const {
+
+    load_with_pointer_offset(frag, 0);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for
+///   int32_t x 16 => int8_t/int4b_t x 16 and
+///   float x 16 => float_e4m3_t/float_e5m2_t x 16
+template <
+  typename ThreadMap_,      ///< Thread map (concept: OutputTileThreadMap)
+  typename Element_,
+  int OutputSizeBits_       ///< Size of output element in bits
+>
+class SharedLoadIteratorMixed<ThreadMap_, Element_, 32, OutputSizeBits_, 16, 8, true> {
+public:
+  using ThreadMap = ThreadMap_;
+  using Shape = typename ThreadMap::Shape;
+
+  using Element = Element_;
+  static_assert(sizeof_bits<Element>::value == 32, "Element size in bits must be 32.");
+
+  using Layout = layout::RowMajor;
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = MatrixCoord;
+
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+
+  static int const kAlignment = 16;
+
+  static int const kThreads = ThreadMap::kThreads;
+
+  /// Fragment object
+  using Fragment = Array<
+    Element, 
+    ThreadMap::Iterations::kColumn * 
+    ThreadMap::Iterations::kRow * 
+    ThreadMap::Iterations::kGroup * 
+    ThreadMap::Iterations::kCluster * 
+    ThreadMap::kElementsPerAccess>;
+
+  /// Memory access size
+  using AccessType = AlignedArray<
+    Element, 
+    16, 
+    kAlignment>;
+
+  /// Vector type used for SMEM loads
+  using LoadType = AlignedArray<
+    Element,
+    4,
+    16
+  >;
+
+  static int const kLoadsPerAccess = 4;
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Byte-level pointer
+  LoadType const *pointers_[kLoadsPerAccess];
+
+  /// Stride along adjacent rows in units of LoadType
+  int stride_;
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  CUTLASS_DEVICE
+  SharedLoadIteratorMixed(
+    TensorRef ref,
+    int thread_idx
+  ):
+    stride_((ref.stride(0) / LoadType::kElements)) {
+
+    TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx);
+    
+    // Initialize pointers
+    LoadType const *base_ptr = reinterpret_cast<LoadType const *>(ref.data()) + thread_offset.row() * stride_;
+      
+    int lane_col_idx = thread_offset.column() / 16;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      int lane_offset = (lane_col_idx % 2) * 4 | ((lane_col_idx / 2) * 8) | ((lane_col_idx / 2) ^ i);
+ 
+      pointers_[i] = base_ptr + lane_offset;
+    }
+  }
+
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      pointers_[i] += pointer_offset / LoadType::kElements;
+    }
+  }
+
+  CUTLASS_DEVICE
+  void add_tile_offset(TensorCoord const &offset) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      pointers_[i] += 
+        offset.row() * Shape::kRow * stride_ + 
+        offset.column() * Shape::kColumn / LoadType::kElements;
+    }
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load_with_pointer_offset(Fragment &frag, Index pointer_offset) {
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int row_ptr_offset =
+            row * ThreadMap::Delta::kRow * stride_ + 
+            group * ThreadMap::Delta::kGroup* stride_ + 
+            cluster * ThreadMap::Delta::kCluster * stride_ +
+            pointer_offset / LoadType::kElements;
+
+          int frag_row_idx = (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          LoadType *frag_ptr = reinterpret_cast<LoadType *>(&frag);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+            
+            int frag_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column;
+
+            CUTLASS_PRAGMA_UNROLL
+            for (int v = 0; v < kLoadsPerAccess; ++v) {
+           
+              LoadType const *memory_pointer = pointers_[v];
+            
+              frag_ptr[frag_idx * kLoadsPerAccess + v] = memory_pointer[row_ptr_offset];
+            }
+          }
+        }
+      }
+    }
+  }
+
+  /// Set base smem address
+  CUTLASS_DEVICE
+  void set_smem_base_address(Index address) {}
+
+  /// Loads a fragment
+  CUTLASS_DEVICE
+  void load(Fragment &frag) {
+
+    load_with_pointer_offset(frag, 0);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for:
+///   int32_t x 8 => int8_t/int4b_t x 8 and
+///   float x 8 => float_e4m3_t/float_e5m2_t x 8
+template <
+  typename ThreadMap_,      ///< Thread map (concept: OutputTileThreadMap)
+  typename Element_,
+  int OutputSizeBits_
+>
+class SharedLoadIteratorMixed<ThreadMap_, Element_, 32, OutputSizeBits_, 8, 8, true> {
+public:
+  using ThreadMap = ThreadMap_;
+  using Shape = typename ThreadMap::Shape;
+
+  using Element = Element_;
+  static_assert(sizeof_bits<Element>::value == 32, "Element size in bits must be 32.");
+
+  using Layout = layout::RowMajor;
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = MatrixCoord;
+
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+
+  static int const kAlignment = 8;
+
+  static int const kThreads = ThreadMap::kThreads;
+
+  /// Fragment object
+  using Fragment = Array<
+    Element, 
+    ThreadMap::Iterations::kColumn * 
+    ThreadMap::Iterations::kRow * 
+    ThreadMap::Iterations::kGroup * 
+    ThreadMap::Iterations::kCluster * 
+    ThreadMap::kElementsPerAccess>;
+
+  /// Memory access size
+  using AccessType = AlignedArray<
+    Element, 
+    8, 
+    kAlignment>;
+
+  /// Vector type used for SMEM loads
+  using LoadType = AlignedArray<
+    Element,
+    4,
+    16
+  >;
+
+  static int const kLoadsPerAccess = 2;
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Byte-level pointer
+  LoadType const *pointers_[kLoadsPerAccess];
+
+  /// Stride along adjacent rows in units of LoadType
+  int stride_;
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  CUTLASS_DEVICE
+  SharedLoadIteratorMixed(
+    TensorRef ref,
+    int thread_idx
+  ):
+    stride_((ref.stride(0) / LoadType::kElements)) {
+
+    TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx);
+    
+    // Initialize pointers
+    LoadType const *base_ptr = reinterpret_cast<LoadType const *>(ref.data()) + thread_offset.row() * stride_;
+      
+    int lane_col_idx = thread_offset.column() / 8;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      int lane_offset = (lane_col_idx % 8) * 2 | ((lane_col_idx / 4) ^ i);
+
+      pointers_[i] = base_ptr + lane_offset;
+    }
+  }
+
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      pointers_[i] += pointer_offset / LoadType::kElements;
+    }
+  }
+
+  CUTLASS_DEVICE
+  void add_tile_offset(TensorCoord const &offset) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      pointers_[i] += 
+        offset.row() * Shape::kRow * stride_ + 
+        offset.column() * Shape::kColumn / LoadType::kElements;
+    }
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load_with_pointer_offset(Fragment &frag, Index pointer_offset) {
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+
+          int row_ptr_offset =
+            row * ThreadMap::Delta::kRow * stride_ + 
+            group * ThreadMap::Delta::kGroup* stride_ + 
+            cluster * ThreadMap::Delta::kCluster * stride_ +
+            pointer_offset / LoadType::kElements;
+
+          int frag_row_idx = (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
+
+          LoadType *frag_ptr = reinterpret_cast<LoadType *>(&frag);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
+            
+            int frag_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column;
+
+            CUTLASS_PRAGMA_UNROLL
+            for (int v = 0; v < kLoadsPerAccess; ++v) {
+           
+              LoadType const *memory_pointer = pointers_[v];
+            
+              frag_ptr[frag_idx * kLoadsPerAccess + v] = memory_pointer[row_ptr_offset];
+            }
+          }
+        }
+      }
+    }
+  }
+
+  /// Set base smem address
+  CUTLASS_DEVICE
+  void set_smem_base_address(Index address) {}
+
+  /// Loads a fragment
+  CUTLASS_DEVICE
+  void load(Fragment &frag) {
+
+    load_with_pointer_offset(frag, 0);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace epilogue
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/threadblock/shared_load_iterator_pitch_linear.h

@@ -0,0 +1,194 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+  \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
+
+  This assumes the shared memory tile is in a permuted layout which avoids bank conflicts on loading.
+  
+  When the fragment is loaded into registers, it matches the row-major thread map assumed by
+  the predicated tile iterator writing to global memory.
+
+  The epilogue rearranges the result of a matrix product through shared memory to match canonical
+  tensor layouts in global memory. Epilogues support conversion and reduction operations.
+*/
+
+#pragma once
+
+#include "cutlass/array.h"
+#include "cutlass/cutlass.h"
+#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
+#include "cutlass/transform/pitch_linear_thread_map.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/matrix_shape.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/tensor_ref.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace threadblock {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tile iterator used to load output tile from shared memory in epilogue.
+///
+/// Satisfies: ReadableTileIterator
+///
+template <typename ThreadMap_,  ///< Thread map (conept: PitchLinearThreadMap)
+          typename Element_,    ///< Element data type
+          int MaxAlignment = ThreadMap_::kElementsPerAccess *sizeof_bits<Element_>::value / 8>
+class SharedLoadIteratorPitchLinear {
+ public:
+  using ThreadMap = ThreadMap_;
+  using Element = Element_;
+
+  using Layout = layout::RowMajor;
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = MatrixCoord;
+
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+
+  static int const kMinAlignment =
+      ThreadMap_::kElementsPerAccess * sizeof_bits<Element_>::value / 8;
+
+  static int const kAlignment = (MaxAlignment < kMinAlignment ? MaxAlignment : kMinAlignment);
+
+  static int const kThreads = ThreadMap::kThreads;
+
+  /// Fragment object
+  using Fragment = Array<Element, ThreadMap::Iterations::kCount * kElementsPerAccess>;
+
+  /// Memory access size
+  using AccessType = AlignedArray<Element, kElementsPerAccess, kAlignment>;
+
+  /// Vector type used for SMEM loads
+  using LoadType =
+      AlignedArray<Element,
+                   const_min(128 / sizeof_bits<Element>::value, ThreadMap::kElementsPerAccess),
+                   const_min(16, kAlignment)>;
+
+  static int const kLoadsPerAccess = AccessType::kElements / LoadType::kElements;
+
+ private:
+  //
+  // Data members
+  //
+
+  /// Byte-level pointer
+  uint8_t *byte_pointer_;
+
+  /// Stride along adjacent rows
+  int stride_;
+
+  /// Base address offset
+  Index base_smem_address_;
+
+ public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  CUTLASS_DEVICE
+  SharedLoadIteratorPitchLinear(TensorRef ref, int thread_idx)
+      : byte_pointer_(reinterpret_cast<uint8_t *>(ref.data())),
+        stride_((ref.stride(0) * sizeof_bits<Element>::value) / 8),
+        base_smem_address_(0) {
+    TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx);
+
+    // Initialize pointer
+    // thread_offset.row() is contiguous dim
+    // thread_offset.column() is stride dim
+    byte_pointer_ += thread_offset.row() * sizeof(AccessType) / kElementsPerAccess+
+                     thread_offset.column() * stride_ ;
+  }
+
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    byte_pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
+  }
+
+  CUTLASS_DEVICE
+  void add_tile_offset(TensorCoord const &offset) {
+    byte_pointer_ +=
+        offset.row() * ThreadMap::StorageShape::kContiguous * sizeof(AccessType) / kElementsPerAccess +
+        offset.column() * ThreadMap::StorageShape::kStrided * stride_;
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const {
+    CUTLASS_PRAGMA_UNROLL
+    for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) {
+        uint8_t const *byte_pointer =
+            byte_pointer_ + s * ThreadMap::Delta::kStrided * stride_ +
+            c * ThreadMap::Delta::kContiguous * ThreadMap::kElementsPerAccess *
+                sizeof_bits<Element>::value / 8 +
+            pointer_offset * sizeof_bits<Element>::value / 8 + base_smem_address_;
+
+        int frag_base_idx = s * ThreadMap::Iterations::kContiguous + c;
+
+        LoadType *frag_ptr = reinterpret_cast<LoadType *>(&frag);
+
+        LoadType const *memory_pointer = reinterpret_cast<LoadType const *>(byte_pointer);
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int v = 0; v < kLoadsPerAccess; ++v) {
+          frag_ptr[frag_base_idx * kLoadsPerAccess + v] = memory_pointer[v];
+        }
+      }
+    }
+  }
+
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void set_smem_base_address(Index address) { base_smem_address_ = address; }
+
+  /// Loads a fragment
+  CUTLASS_DEVICE
+  void load(Fragment &frag) const { load_with_pointer_offset(frag, 0); }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace threadblock
+}  // namespace epilogue
+}  // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/fragment_iterator_complex_tensor_op.h

@@ -0,0 +1,187 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief This defines a "fragment" iterator for visiting the fragments of an accumulator tile
+      that participate in one warp-level store operation.
+
+      Typically, the accumulator tile is the largest single block of register-backed storage 
+      within the kernel. Storing it to memory is best accomplished by partitioning it into
+      smaller tiles and storing these sequentially.
+
+      Round trips through shared memory during the Epilogue phase require partitioning, as
+      shared memory capacity is typically insufficient for a threadblock's total accumulator
+      size.
+*/
+
+#pragma once
+
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+
+#include "cutlass/epilogue/warp/tensor_op_policy.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// 
+template <
+  typename WarpShape,         ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename OperatorShape,     ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  typename OperatorElementC,  ///< matrix multiply operation data type (concept: data type)
+  typename OperatorFragmentC, ///< matrix multiply operation fragment (concept: Array)
+  typename Layout             ///< target shared memory layout
+>
+class FragmentIteratorComplexTensorOp;
+
+////////////////////////////////////////////////////////////////////////////////
+
+
+/// Partial specialization for row-major shared memory
+template <
+  typename WarpShape_,         ///< shape of the warp-level GEMM tile
+  typename OperatorShape_,     ///< underlying real-valued matrix multiply operation shape (concept: gemm::GemmShape)
+  typename OperatorElementC_,  ///< underlying real-valued matrix multiply operation data type
+  typename OperatorFragmentC_  ///< underlying real-valued matrix multiply operation fragment (concept: Array)
+>
+class FragmentIteratorComplexTensorOp<WarpShape_, OperatorShape_, OperatorElementC_, OperatorFragmentC_, layout::RowMajor> {
+public:
+
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using OperatorElementC = OperatorElementC_;
+  using OperatorFragmentC = OperatorFragmentC_;
+  using Layout = layout::RowMajor;
+
+  using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    complex<OperatorElementC>, 
+    Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
+
+  static int const kRealIndex = 0;
+
+  /// Offset into the accumulator fragment
+  static int const kImaginaryIndex = 
+    OperatorFragmentC::kElements * Policy::OperatorCount::kRow * Policy::OperatorCount::kColumn;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = Array<OperatorElementC, 2 * kImaginaryIndex>;
+
+  /// This is the complete warp-level accumulator tile.
+  using OutputAccumulatorTile = Array<complex<OperatorElementC>, kImaginaryIndex>;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+private:
+
+  /// Internal access type
+  using AccessType = Array<OperatorElementC, Policy::kElementsPerAccess>;
+
+  using FragmentAccessType = Array<complex<OperatorElementC>, Policy::kElementsPerAccess>;
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Accumulator tile
+  AccessType const *accumulators_;
+
+  /// Internal index
+  int index_;
+
+public:
+
+  /// Constructs an iterator
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorComplexTensorOp(AccumulatorTile const &accum): 
+    accumulators_(reinterpret_cast<AccessType const *>(&accum)), 
+    index_(0) {
+
+  }
+
+  /// Increments
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorComplexTensorOp &operator++() {
+    ++index_;
+    return *this;
+  }
+
+  /// Decrements
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorComplexTensorOp &operator--() {
+    --index_;
+    return *this;
+  }
+
+  /// Loads a fragment from the referenced part of the accumulator tile
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag, int index_offset = 0) const {
+
+    int index = index_ + index_offset;
+
+    FragmentAccessType *frag_ptr = reinterpret_cast<FragmentAccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
+
+      int accumulator_access_offset = 
+        index + n * Policy::kAccumulatorColumnStride / Policy::kElementsPerAccess;
+
+      auto const & real_accum_array = accumulators_[accumulator_access_offset + kRealIndex];
+      auto const & imag_accum_array = accumulators_[accumulator_access_offset + kImaginaryIndex / Policy::kElementsPerAccess];
+
+      // Pack real and imaginary parts into a structure. This is likely to result in MOVs
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < Policy::kElementsPerAccess; ++i) {
+
+        frag_ptr[n][i].real() = real_accum_array[i];
+        frag_ptr[n][i].imag() = imag_accum_array[i]; 
+      }
+    }
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/fragment_iterator_gaussian_complex_tensor_op.h

@@ -0,0 +1,194 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief This defines a "fragment" iterator for visiting the fragments of an accumulator tile
+      that participate in one warp-level store operation.
+
+      Typically, the accumulator tile is the largest single block of register-backed storage 
+      within the kernel. Storing it to memory is best accomplished by partitioning it into
+      smaller tiles and storing these sequentially.
+
+      Round trips through shared memory during the Epilogue phase require partitioning, as
+      shared memory capacity is typically insufficient for a threadblock's total accumulator
+      size.
+*/
+
+#pragma once
+
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+
+#include "cutlass/epilogue/warp/tensor_op_policy.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// 
+template <
+  typename WarpShape,         ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename OperatorShape,     ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  typename OperatorElementC,  ///< matrix multiply operation data type (concept: data type)
+  typename OperatorFragmentC, ///< matrix multiply operation fragment (concept: Array)
+  typename Layout             ///< target shared memory layout
+>
+class FragmentIteratorGaussianComplexTensorOp;
+
+////////////////////////////////////////////////////////////////////////////////
+
+
+/// Partial specialization for row-major shared memory
+template <
+  typename WarpShape_,         ///< shape of the warp-level GEMM tile
+  typename OperatorShape_,     ///< underlying real-valued matrix multiply operation shape (concept: gemm::GemmShape)
+  typename OperatorElementC_,  ///< underlying real-valued matrix multiply operation data type
+  typename OperatorFragmentC_  ///< underlying real-valued matrix multiply operation fragment (concept: Array)
+>
+class FragmentIteratorGaussianComplexTensorOp<WarpShape_, OperatorShape_, OperatorElementC_, OperatorFragmentC_, layout::RowMajor> {
+public:
+
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using OperatorElementC = OperatorElementC_;
+  using OperatorFragmentC = OperatorFragmentC_;
+  using Layout = layout::RowMajor;
+
+  using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    complex<OperatorElementC>, 
+    Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
+
+  /// Size of one part of accumulator of 3-part accumulator in units of number of OperatorElementC
+  static int const kElementsAccumulatorPerPart = 
+    OperatorFragmentC::kElements * Policy::OperatorCount::kRow * Policy::OperatorCount::kColumn;
+
+  /// Offset into the accumulator fragment part 1
+  static int const kPart1Index = kElementsAccumulatorPerPart * 0;
+
+  /// Offset into the accumulator fragment part 2
+  static int const kPart2Index = kElementsAccumulatorPerPart * 1;
+
+  /// Offset into the accumulator fragment part 3
+  static int const kPart3Index = kElementsAccumulatorPerPart * 2;
+
+  /// This is the complete warp-level accumulator tile holding part1, part2, and part3
+  using AccumulatorTile = Array<OperatorElementC, kElementsAccumulatorPerPart * 3>;
+
+  /// This is the complete warp-level accumulator tile holding final output of complex<T> type 
+  using OutputAccumulatorTile = Array<complex<OperatorElementC>, kElementsAccumulatorPerPart>;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+private:
+
+  /// Internal access type
+  using AccessType = Array<OperatorElementC, Policy::kElementsPerAccess>;
+
+  using FragmentAccessType = Array<complex<OperatorElementC>, Policy::kElementsPerAccess>;
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Accumulator tile
+  AccessType const *accumulators_;
+
+  /// Internal index
+  int index_;
+
+public:
+
+  /// Constructs an iterator
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorGaussianComplexTensorOp(AccumulatorTile const &accum): 
+    accumulators_(reinterpret_cast<AccessType const *>(&accum)), 
+    index_(0) {
+  }
+
+  /// Increments
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorGaussianComplexTensorOp &operator++() {
+    ++index_;
+    return *this;
+  }
+
+  /// Decrements
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorGaussianComplexTensorOp &operator--() {
+    --index_;
+    return *this;
+  }
+
+  /// Loads a fragment from the referenced part of the accumulator tile
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag, int index_offset = 0) const {
+
+    int index = index_ + index_offset;
+
+    FragmentAccessType *frag_ptr = reinterpret_cast<FragmentAccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
+
+      int accumulator_access_offset = 
+        index + n * Policy::kAccumulatorColumnStride / Policy::kElementsPerAccess;
+
+      auto const & part1_accum_array = accumulators_[accumulator_access_offset + kPart1Index];
+      auto const & part2_accum_array = accumulators_[accumulator_access_offset + kPart2Index / Policy::kElementsPerAccess];
+      auto const & part3_accum_array = accumulators_[accumulator_access_offset + kPart3Index / Policy::kElementsPerAccess];
+
+      // Pack parts 1, 2, and 3 into a structure. This is likely to result in MOVs
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < Policy::kElementsPerAccess; ++i) {
+
+        frag_ptr[n][i].real() = part1_accum_array[i] - part3_accum_array[i];
+        frag_ptr[n][i].imag() = part1_accum_array[i] + part2_accum_array[i]; 
+      }
+    }
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/fragment_iterator_simt.h

@@ -0,0 +1,164 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief This defines a "fragment" iterator for visiting the fragments of an accumulator tile
+      that participate in one warp-level store operation.
+
+      Typically, the accumulator tile is the largest single block of register-backed storage 
+      within the kernel. Storing it to memory is best accomplished by partitioning it into
+      smaller tiles and storing these sequentially.
+
+      Round trips through shared memory during the Epilogue phase require partitioning, as
+      shared memory capacity is typically insufficient for a threadblock's total accumulator
+      size.
+*/
+
+#pragma once
+
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+
+#include "cutlass/epilogue/warp/simt_policy.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fragment iterator for SIMT accumulator arrangements
+template <
+  typename WarpShape,             ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename Operator,              ///< matrix multiply operation (concept: arch::Mma)
+  typename Layout,                ///< target shared memory layout
+  typename MmaSimtPolicy          ///< policy defining lane arrangement (concept: MmaSimtPolicy)
+>
+class FragmentIteratorSimt;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for row-major shared memory
+template <
+  typename WarpShape_,     ///< shape of the warp-level GEMM tile
+  typename Operator_ ,     ///< matrix multiply operator (concept: arch::Mma)
+  typename MmaSimtPolicy_  ///< policy defining lane arrangement (concept: MmaSimtPolicy)
+>
+class FragmentIteratorSimt<WarpShape_, Operator_, layout::RowMajor, MmaSimtPolicy_> {
+public:
+
+  using WarpShape = WarpShape_;
+  using Operator = Operator_;
+  using Layout = layout::RowMajor;
+
+  /// Policy for warp-level epilogue components
+  using Policy = SimtPolicy<WarpShape, Operator, Layout, MmaSimtPolicy_>;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    typename Operator::ElementC, 
+    Policy::kElementsPerIteration>;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = Array<
+    typename Operator::ElementC, 
+    Policy::kAccumulatorElementCount>;
+
+  using OutputAccumulatorTile = AccumulatorTile;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+private:
+
+  /// Internal access type
+  using AccessType = Array<typename Operator::ElementC, Policy::kElementsPerAccess>;
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Accumulator tile
+  AccessType const *accumulators_;
+
+  /// Internal index
+  int index_;
+
+public:
+
+  /// Constructs an iterator
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorSimt(AccumulatorTile const &accum): 
+    accumulators_(reinterpret_cast<AccessType const *>(&accum)), 
+    index_(0) {
+
+  }
+
+  /// Increments
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorSimt &operator++() {
+    ++index_;
+    return *this;
+  }
+
+  /// Decrements
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorSimt &operator--() {
+    --index_;
+    return *this;
+  }
+
+  /// Loads a fragment from the referenced part of the accumulator tile
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag, int index_offset = 0) const {
+
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::kAccessesPerIteration; ++n) {
+
+      int accumulator_access_offset = index_ * Policy::kAccessesPerIteration + n;
+
+      frag_ptr[n] = accumulators_[accumulator_access_offset];
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/fragment_iterator_tensor_op.h

@@ -0,0 +1,378 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief This defines a "fragment" iterator for visiting the fragments of an accumulator tile
+      that participate in one warp-level store operation.
+
+      Typically, the accumulator tile is the largest single block of register-backed storage 
+      within the kernel. Storing it to memory is best accomplished by partitioning it into
+      smaller tiles and storing these sequentially.
+
+      Round trips through shared memory during the Epilogue phase require partitioning, as
+      shared memory capacity is typically insufficient for a threadblock's total accumulator
+      size.
+*/
+
+#pragma once
+
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+
+#include "cutlass/epilogue/warp/tensor_op_policy.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// 
+template <
+  typename WarpShape,         ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename OperatorShape,     ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  typename OperatorElementC,  ///< matrix multiply operation data type (concept: data type)
+  typename OperatorFragmentC, ///< matrix multiply operation fragment (concept: Array)
+  typename Layout             ///< target shared memory layout
+>
+class FragmentIteratorTensorOp;
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for row-major shared memory
+template <
+  typename WarpShape_,         ///< shape of the warp-level GEMM tile
+  typename OperatorShape_,     ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  typename OperatorElementC_,  ///< matrix multiply operation data type (concept: data type)
+  typename OperatorFragmentC_  ///< matrix multiply operation fragment (concept: Array)
+>
+class FragmentIteratorTensorOp<WarpShape_, OperatorShape_, OperatorElementC_, OperatorFragmentC_, layout::RowMajor> {
+public:
+
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using OperatorElementC = OperatorElementC_;
+  using OperatorFragmentC = OperatorFragmentC_;
+  using Layout = layout::RowMajor;
+
+  using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    OperatorElementC, 
+    Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = Array<
+    OperatorElementC, 
+    OperatorFragmentC::kElements * Policy::OperatorCount::kRow * Policy::OperatorCount::kColumn>;
+
+  using OutputAccumulatorTile = AccumulatorTile;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+  using TileIterations = typename Policy::TileIterations;
+  static int const kIterationsPerTile = kIterations / TileIterations::kCount;
+
+private:
+
+  /// Internal access type
+  using AccessType = Array<OperatorElementC, Policy::kElementsPerAccess>;
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Accumulator tile
+  AccessType const *accumulators_;
+
+  /// Internal index
+  int index_;
+
+public:
+
+  /// Constructs an iterator
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorTensorOp(AccumulatorTile const &accum): 
+    accumulators_(reinterpret_cast<AccessType const *>(&accum)), 
+    index_(0) {
+  }
+
+  /// Increments
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorTensorOp &operator++() {
+    ++index_;
+    return *this;
+  }
+
+  /// Decrements
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorTensorOp &operator--() {
+    --index_;
+    return *this;
+  }
+
+  /// Loads a fragment from the referenced part of the accumulator tile
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag, int index_offset = 0) const {
+
+    int index = index_ + index_offset;
+
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
+
+      int accumulator_access_offset = 
+        index + n * Policy::kAccumulatorColumnStride / Policy::kElementsPerAccess;
+
+      frag_ptr[n] = accumulators_[accumulator_access_offset];
+    }
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for col-major shared memory
+/// Only works for 168x tensor core kernels
+template <
+  typename WarpShape_,         ///< shape of the warp-level GEMM tile
+  typename OperatorShape_,     ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  typename OperatorElementC_,  ///< matrix multiply operation data type (concept: data type)
+  typename OperatorFragmentC_  ///< matrix multiply operation fragment (concept: Array)
+>
+class FragmentIteratorTensorOp<WarpShape_, OperatorShape_, OperatorElementC_, OperatorFragmentC_, layout::ColumnMajor> {
+public:
+
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using OperatorElementC = OperatorElementC_;
+  using OperatorFragmentC = OperatorFragmentC_;
+  using Layout = layout::ColumnMajor;
+
+  using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    OperatorElementC, 
+    4 * Policy::OperatorCount::kRow * Policy::kElementsPerAccess>;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = Array<
+    OperatorElementC, 
+    OperatorFragmentC::kElements * Policy::OperatorCount::kRow * Policy::OperatorCount::kColumn>;
+
+  using OutputAccumulatorTile = AccumulatorTile;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+  using TileIterations = typename Policy::TileIterations;
+  static int const kIterationsPerTile = kIterations / TileIterations::kCount;
+
+private:
+
+  /// Internal access type
+  using AccessType = Array<OperatorElementC, Policy::kElementsPerAccess>;
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Accumulator tile
+  AccessType const *accumulators_;
+
+  /// Internal index
+  int index_;
+
+public:
+
+  /// Constructs an iterator
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorTensorOp(AccumulatorTile const &accum): 
+    accumulators_(reinterpret_cast<AccessType const *>(&accum)), 
+    index_(0) {
+  }
+
+  /// Increments
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorTensorOp &operator++() {
+    ++index_;
+    return *this;
+  }
+
+  /// Decrements
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorTensorOp &operator--() {
+    --index_;
+    return *this;
+  }
+
+  /// Loads a fragment from the referenced part of the accumulator tile
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag, int index_offset = 0) const {
+
+    int index = index_ + index_offset;
+
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Policy::kAccumulatorRowStride; ++i) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int m = 0; m < (Policy::OperatorCount::kRow * 2); ++m) {
+
+        int accumulator_access_offset = 
+          index * Policy::kAccumulatorColumnStride + m * Policy::kAccumulatorRowStride / Policy::kElementsPerAccess + i;
+
+        frag_ptr[m + i * Policy::OperatorCount::kRow * 2] = accumulators_[accumulator_access_offset];
+      }
+    }
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Dedicated to interleaved layout
+template <
+    /// shape of the warp-level GEMM tile
+    typename WarpShape_,
+    /// matrix multiply operator shape (concept: gemm::GemmShape)
+    typename OperatorShape_,
+    /// matrix multiply operator data type (concept: data type)
+    typename OperatorElementC_,
+    /// matrix multiply operator fragment (concept: Array)
+    typename OperatorFragmentC_,
+    /// number of interleaved k
+    int InterleavedK>
+class FragmentIteratorTensorOp<WarpShape_, OperatorShape_, OperatorElementC_, OperatorFragmentC_,
+                               layout::ColumnMajorInterleaved<InterleavedK>> {
+ public:
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using OperatorElementC = OperatorElementC_;
+  using OperatorFragmentC = OperatorFragmentC_;
+  static int const kInterleavedK = InterleavedK;
+  using Layout = layout::ColumnMajorInterleaved<kInterleavedK>;
+
+  using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment =
+      Array<OperatorElementC,
+            Policy::kElementsPerAccess * InterleavedK / OperatorShape::kN>;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile =
+      Array<OperatorElementC, OperatorFragmentC::kElements *
+                                  Policy::OperatorCount::kRow *
+                                  Policy::OperatorCount::kColumn>;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+  using TileIterations = typename Policy::TileIterations;
+  static int const kIterationsPerTile = kIterations / TileIterations::kCount;
+
+ private:
+  /// Internal access type
+  using AccessType =
+      Array<OperatorElementC, Policy::kElementsPerAccess>;
+
+ private:
+  //
+  // Data members
+  //
+
+  /// Accumulator tile
+  AccessType const *accumulators_;
+
+  /// Internal index
+  int index_;
+
+ public:
+  /// Constructs an iterator
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorTensorOp(AccumulatorTile const &accum)
+      : accumulators_(reinterpret_cast<AccessType const *>(&accum)),
+        index_(0) {}
+
+  /// Increments
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorTensorOp &operator++() {
+    ++index_;
+    return *this;
+  }
+
+  /// Decrements
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorTensorOp &operator--() {
+    --index_;
+    return *this;
+  }
+
+  /// Loads a fragment from the referenced part of the accumulator tile
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag, int index_offset = 0) const {
+    int index = index_ + index_offset;
+
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < (InterleavedK / OperatorShape::kN); ++n) {
+      int index_m = index % (Policy::OperatorCount::kRow *
+                             Policy::kIterationsPerInstruction);
+      int index_n = index / (Policy::OperatorCount::kRow *
+                             Policy::kIterationsPerInstruction);
+      int accumulator_access_offset =
+          (index_m / Policy::kIterationsPerInstruction) *
+              (Policy::OperatorCount::kColumn *
+               Policy::kIterationsPerInstruction) +
+          (index_m % Policy::kIterationsPerInstruction) +
+          index_n * (InterleavedK / OperatorShape::kN) *
+              Policy::kIterationsPerInstruction +
+          n * Policy::kIterationsPerInstruction;
+
+      frag_ptr[n] = accumulators_[accumulator_access_offset];
+    }
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/fragment_iterator_volta_tensor_op.h

@@ -0,0 +1,269 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief This defines a "fragment" iterator for visiting the fragments of an accumulator tile
+      that participate in one warp-level store operation.
+
+      Typically, the accumulator tile is the largest single block of register-backed storage 
+      within the kernel. Storing it to memory is best accomplished by partitioning it into
+      smaller tiles and storing these sequentially.
+
+      Round trips through shared memory during the Epilogue phase require partitioning, as
+      shared memory capacity is typically insufficient for a threadblock's total accumulator
+      size.
+*/
+
+#pragma once
+
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/gemm/gemm.h"
+
+#include "cutlass/epilogue/warp/volta_tensor_op_policy.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 
+template <
+  typename WarpShape,             ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename InterleavedTileShape,  ///< shape of indivisible instruction-level arrangement (concept: GemmShape)
+  typename ElementC,              ///< Accumulator layout
+  typename Layout                 ///< target shared memory layout
+>
+class FragmentIteratorVoltaTensorOp;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for row-major shared memory
+template <
+  typename WarpShape_         ///< shape of warp-level GEMM (concept: MatrixShape)
+>
+class FragmentIteratorVoltaTensorOp<WarpShape_, gemm::GemmShape<32, 32, 4>, half_t, layout::RowMajor> {
+public:
+
+  using WarpShape = WarpShape_;
+  using InterleavedTileShape = gemm::GemmShape<32, 32, 4>;
+  using ElementC = half_t;
+  using Layout = layout::RowMajor;
+
+  /// Policy operator
+  using Policy = VoltaTensorOpPolicy<WarpShape, InterleavedTileShape, ElementC, Layout>;
+
+  /// Array type for aligned memory accesses
+  using AccessType = typename Policy::AccessType;
+  
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = typename Policy::Fragment;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = typename Policy::AccumulatorTile;
+
+  using OutputAccumulatorTile = AccumulatorTile;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+private:
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Accumulator tile
+  AccessType const *accumulators_;
+
+  /// Internal index
+  int index_;
+
+public:
+
+  /// Constructs an iterator
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorVoltaTensorOp(AccumulatorTile const &accum): 
+    accumulators_(reinterpret_cast<AccessType const *>(&accum)), 
+    index_(0) {
+
+  }
+
+  /// Increments
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorVoltaTensorOp &operator++() {
+    ++index_;
+    return *this;
+  }
+
+  /// Decrements
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorVoltaTensorOp &operator--() {
+    --index_;
+    return *this;
+  }
+
+  /// Loads a fragment from the referenced part of the accumulator tile
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag, int index_offset = 0) const {
+
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    static int const kAccessesPerMma = Policy::kElementsPerMma / Policy::kElementsPerAccess;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int tile_n = 0; tile_n < Policy::TileIterations::kColumn; ++tile_n) {
+      
+      int tile_access_idx = 
+        (tile_n * Policy::TileIterations::kRow + (index_ & 2) / 2) * Policy::MmaIterations::kCount * kAccessesPerMma;
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn * kAccessesPerMma; ++mma_n) {
+
+        int mma_access_idx = ((mma_n & 1) * 2 + (index_ & 1)) * kAccessesPerMma + (mma_n & 2) / 2;
+
+        frag_ptr[tile_n * Policy::MmaIterations::kColumn * kAccessesPerMma +
+          mma_n] = accumulators_[tile_access_idx + mma_access_idx];
+      }
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for row-major shared memory
+template <
+  typename WarpShape_         ///< shape of warp-level GEMM (concept: MatrixShape)
+>
+class FragmentIteratorVoltaTensorOp<WarpShape_, gemm::GemmShape<32, 32, 4>, float, layout::RowMajor> {
+public:
+
+  using WarpShape = WarpShape_;
+  using InterleavedTileShape = gemm::GemmShape<32, 32, 4>;
+  using ElementC = float;
+  using Layout = layout::RowMajor;
+
+  /// Policy operator
+  using Policy = VoltaTensorOpPolicy<WarpShape, InterleavedTileShape, ElementC, Layout>;
+
+  /// Array type for aligned memory accesses
+  using AccessType = typename Policy::AccessType;
+  
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = typename Policy::Fragment;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = typename Policy::AccumulatorTile;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+private:
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Accumulator tile
+  AccessType const *accumulators_;
+
+  /// Internal index
+  int index_;
+
+public:
+
+  /// Constructs an iterator
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorVoltaTensorOp(AccumulatorTile const &accum): 
+    accumulators_(reinterpret_cast<AccessType const *>(&accum)), 
+    index_(0) {
+  }
+
+  /// Increments
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorVoltaTensorOp &operator++() {
+    ++index_;
+    return *this;
+  }
+
+  /// Decrements
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorVoltaTensorOp &operator--() {
+    --index_;
+    return *this;
+  }
+
+  /// Loads a fragment from the referenced part of the accumulator tile
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag, int index_offset = 0) const {
+
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    int const kRegsPerMmaRow = 2;
+      
+    CUTLASS_PRAGMA_UNROLL
+    for (int reg_row = 0; reg_row < Policy::kRowsPerMmaTile; ++reg_row) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int tile_n = 0; tile_n < Policy::TileIterations::kColumn; ++tile_n) {
+    
+        CUTLASS_PRAGMA_UNROLL
+        for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn * 2; ++mma_n) {
+
+          int mma_idx = (index_ & 1) + (index_ & 2) * Policy::MmaIterations::kCount / 2 +
+            (tile_n * Policy::TileIterations::kRow) * Policy::MmaIterations::kCount + (mma_n & 1) * 2;
+
+          int reg_offset = reg_row * kRegsPerMmaRow + (mma_n & 2) * 2;
+          int reg_idx = mma_idx * Policy::kElementsPerMma + reg_offset;
+
+          *frag_ptr = accumulators_[reg_idx / Policy::kElementsPerAccess];
+          ++frag_ptr;
+        }
+      }
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/fragment_iterator_wmma_tensor_op.h

@@ -0,0 +1,158 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief This defines a "fragment" iterator for visiting the fragments of an accumulator tile
+      that participate in one warp-level store operation.
+
+      Typically, the accumulator tile is the largest single block of register-backed storage 
+      within the kernel. Storing it to memory is best accomplished by partitioning it into
+      smaller tiles and storing these sequentially.
+
+      Round trips through shared memory during the Epilogue phase require partitioning, as
+      shared memory capacity is typically insufficient for a threadblock's total accumulator
+      size.
+*/
+
+#pragma once
+
+#include "cutlass/wmma_array.h"
+#include "cutlass/layout/matrix.h"
+
+#include "cutlass/epilogue/warp/wmma_tensor_op_policy.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// 
+template <
+  typename WarpShape,         ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename OperatorShape,     ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  typename OperatorElementC,  ///< matrix multiply operation data type (concept: data type)
+  typename OperatorFragmentC, ///< matrix multiply operation fragment (concept: nvcuda::cuda::fragment)
+  typename Layout             ///< target shared memory layout
+>
+class FragmentIteratorWmmaTensorOp;
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for row-major shared memory
+template <
+  typename WarpShape_,         ///< shape of the warp-level GEMM tile
+  typename OperatorShape_,     ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  typename OperatorElementC_,  ///< matrix multiply operation data type (concept: data type)
+  typename OperatorFragmentC_  ///< matrix multiply operation fragment (concept: nvcuda::cuda::fragment)
+>
+class FragmentIteratorWmmaTensorOp<WarpShape_, OperatorShape_, OperatorElementC_, OperatorFragmentC_, layout::RowMajor> {
+public:
+
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using OperatorElementC = OperatorElementC_;
+  using OperatorFragmentC = OperatorFragmentC_;
+  using Layout = layout::RowMajor;
+
+  using Policy = WmmaTensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = WmmaFragmentArray<OperatorFragmentC, Policy::OperatorCount::kColumn>;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = WmmaFragmentArray<OperatorFragmentC, Policy::OperatorCount::kCount>;
+
+  using OutputAccumulatorTile = AccumulatorTile;
+
+private:
+
+  /// Internal access type
+  using AccessType = WmmaFragmentArray<OperatorFragmentC, Policy::kWmmaFragmentsPerAccess>;
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Accumulator tile
+  AccessType const *accumulators_;
+
+  /// Internal index
+  int index_;
+
+public:
+
+  /// Constructs an iterator
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorWmmaTensorOp(AccumulatorTile const &accum): 
+    accumulators_(reinterpret_cast<AccessType const *>(&accum)), 
+    index_(0) { 
+  }
+
+  /// Increments
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorWmmaTensorOp &operator++() {
+    ++index_;
+    return *this;
+  }
+
+  /// Decrements
+  CUTLASS_HOST_DEVICE
+  FragmentIteratorWmmaTensorOp &operator--() {
+    --index_;
+    return *this;
+  }
+
+  /// Loads a fragment from the referenced part of the accumulator tile
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag, int index_offset = 0) const {
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for(int n=0; n < Policy::OperatorCount::kColumn; n++) {
+      
+      int accumulator_access_offset = index_ * Policy::OperatorCount::kColumn + n;
+
+      frag_ptr[n] = accumulators_[accumulator_access_offset];
+    }
+  }
+};
+
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////
+

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/simt_policy.h

@@ -0,0 +1,107 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Defines basic structures needed for implementing the warp-scoped phase of the epilogue.
+          These quantities assume a 'column-major' arrangement of SimtOp instructions, of which
+          a row-oriented slice is visible per iteration.
+*/
+
+#pragma once
+
+#include "cutlass/matrix_shape.h"
+#include "cutlass/layout/matrix.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename WarpShape,            ///< shape of warp-level GEMM (concept: GemmShape)
+  typename Operator,             ///< matrix multiply operation (concept: arch::Mma)
+  typename Layout,               ///< destination layout in shared memory
+  typename MmaSimtPolicy         ///< policy defining lane arrangement (concept: MmaSimtPolicy)
+>
+struct SimtPolicy;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for row-major
+template <
+  typename WarpShape_,           ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename Operator_,            ///< matrix multiply operation (concept: arch::Mma)
+  typename MmaSimtPolicy_        ///< policy defining lane arrangement (concept: MmaSimtPolicy)
+>
+struct SimtPolicy<WarpShape_, Operator_, layout::RowMajor, MmaSimtPolicy_> {
+
+  using WarpShape = WarpShape_;
+  using Operator = Operator_;
+  using MmaSimtPolicy = MmaSimtPolicy_;
+
+  static_assert(!(WarpShape::kM % MmaSimtPolicy::WarpShape::kRow), "Divisibility");
+  static_assert(!(WarpShape::kN % MmaSimtPolicy::WarpShape::kColumn), "Divisibility");
+
+  /// Number of iterations
+  static int const kIterations = WarpShape::kM / MmaSimtPolicy::WarpShape::kRow;
+
+  /// Number of accumulators written per iteration
+  static int const kElementsPerIteration = 
+    (WarpShape::kN / MmaSimtPolicy::WarpShape::kColumn);
+
+  /// Total number of accumulators
+  static int const kAccumulatorElementCount = kElementsPerIteration * kIterations;
+
+  /// Number of consecutive elements
+  static int const kElementsPerAccess = MmaSimtPolicy::LaneMmaShape::kN;
+
+  /// Number of rows per epilogue iteration
+  static int const kRowsPerIteration = MmaSimtPolicy::WarpShape::kRow;
+
+  /// Number of accesses made in one iteration
+  static int const kAccessesPerIteration = kElementsPerIteration / kElementsPerAccess;
+
+  /// Number of elements in between accumulator chunks of (LaneMmaShape::kM x LaneMmaShape::kN)
+  using Delta = MatrixShape<
+    MmaSimtPolicy::WarpShape::kRow * MmaSimtPolicy::LaneMmaShape::kM,
+    MmaSimtPolicy::WarpShape::kColumn * MmaSimtPolicy::LaneMmaShape::kN
+  >;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/tensor_op_policy.h

@@ -0,0 +1,189 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Defines basic structures needed for implementing the warp-scoped phase of the epilogue.
+          These quantities assume a 'column-major' arrangement of TensorOp instructions, of which
+          a row-oriented slice is visible per iteration.
+*/
+
+#pragma once
+
+#include "cutlass/matrix_shape.h"
+#include "cutlass/layout/matrix.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Policy details related to the epilogue
+template <
+  typename WarpShape,     ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename OperatorShape, ///< matrix multiply operation shape (concept: gemm:GemmShape)
+  typename Layout         ///< target shared memory layout
+>
+struct TensorOpPolicy; 
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for row-major
+template <
+  typename WarpShape,           ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename OperatorShape        ///< matrix multiply operation shape (concept: gemm::GemmShape)
+>
+struct TensorOpPolicy<WarpShape, OperatorShape, layout::RowMajor> {
+
+  /// Number of operations
+  using OperatorCount = MatrixShape<
+    (WarpShape::kM + OperatorShape::kM - 1) / OperatorShape::kM,
+    (WarpShape::kN + OperatorShape::kN - 1) / OperatorShape::kN
+  >;
+
+  //
+  // Hard-coded constants regarding Tensor Operations
+  //
+
+  static int const kElementsPerAccess = 2;
+  static int const kRowsPerIteration = 8;
+  static bool const kDivisible = 
+    !(WarpShape::kM % OperatorShape::kM) && !(WarpShape::kN % OperatorShape::kN);
+
+  //
+  // Derived quantities
+  //
+
+  // Number of 'externally visible' iterations per actual instruction
+  static int const kIterationsPerInstruction = OperatorShape::kM / kRowsPerIteration;
+
+  // Number of externally visible iterations
+  static int const kIterations = OperatorCount::kRow * kIterationsPerInstruction;
+
+  using TileIterations = MatrixShape<kIterations, 1>;
+
+  static int const kAccumulatorRowStride = kElementsPerAccess;
+  static int const kAccumulatorColumnStride = kElementsPerAccess * OperatorCount::kRow * kIterationsPerInstruction;
+
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for row-major
+template <
+  typename WarpShape,           ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename OperatorShape        ///< matrix multiply operation shape (concept: gemm::GemmShape)
+>
+struct TensorOpPolicy<WarpShape, OperatorShape, layout::ColumnMajor> {
+
+  /// Number of operations
+  using OperatorCount = MatrixShape<
+    (WarpShape::kM + OperatorShape::kM - 1) / OperatorShape::kM,
+    (WarpShape::kN + OperatorShape::kN - 1) / OperatorShape::kN
+  >;
+
+  //
+  // Hard-coded constants regarding Tensor Operations
+  //
+
+  static int const kElementsPerAccess = 1;
+  static int const kColumnsPerIteration = 8;
+  static bool const kDivisible = 
+    !(WarpShape::kM % OperatorShape::kM) && !(WarpShape::kN % OperatorShape::kN);
+
+  //
+  // Derived quantities
+  //
+
+  // Number of 'externally visible' iterations per actual instruction
+  static int const kIterationsPerInstruction = OperatorShape::kN / kColumnsPerIteration;
+
+  // Number of externally visible iterations
+  static int const kIterations = OperatorCount::kColumn * kIterationsPerInstruction;
+
+  using TileIterations = MatrixShape<kIterations, 1>;
+
+  // Hard code for 16x8
+  static int const kAccumulatorRowStride = 2;
+  static int const kAccumulatorColumnStride = 4 * OperatorCount::kRow;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for column-major-interleaved
+template <
+    typename WarpShape,  ///< shape of warp-level GEMM (concept: MatrixShape)
+    typename OperatorShape,   ///< matrix multiply operation (concept: arch::Mma)
+    int InterleavedK     ///< number of interleaved k
+    >
+struct TensorOpPolicy<WarpShape, OperatorShape,
+                      layout::ColumnMajorInterleaved<InterleavedK> > {
+  /// Number of operations
+  using OperatorCount = MatrixShape<WarpShape::kM / OperatorShape::kM,
+                                    WarpShape::kN / OperatorShape::kN>;
+
+  //
+  // Hard-coded constants regarding Tensor Operations
+  //
+
+  static int const kElementsPerAccess = 2;
+  static int const kRowsPerIteration = 8;
+
+  //
+  // Derived quantities
+  //
+
+  // Number of 'externally visible' iterations per actual instruction
+  static int const kIterationsPerInstruction =
+      OperatorShape::kM / kRowsPerIteration;
+
+  // Number of externally visible iterations
+  static int const kIterations = WarpShape::kN / InterleavedK *
+                                 OperatorCount::kRow *
+                                 kIterationsPerInstruction;
+
+  static int const kElementsPerIteration = InterleavedK / OperatorShape::kN * kElementsPerAccess;
+
+  static int const kAccessPerIteration = kElementsPerIteration / kElementsPerAccess;
+
+  // Number of externally visible iterations
+  //static int const kTileIterations = OperatorCount::kRow * kIterationsPerInstruction;
+  using TileIterations = MatrixShape<1, WarpShape::kN / InterleavedK>;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/tile_iterator_simt.h

@@ -0,0 +1,785 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief 
+*/
+
+#pragma once
+
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/layout/pitch_linear.h"
+
+#include "cutlass/epilogue/warp/simt_policy.h"
+
+#define CUTLASS_SIMT_EPILOGUE_USE_SCALAR_STORES 1
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template for reading and writing tiles of accumulators to shared memory
+template <
+  typename WarpShape,     ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename Operator,      ///< matrix multiply operation (concept: arch::Mma)
+  typename Element,       ///< data type of element to be written
+  typename Layout,        ///< target shared memory layout
+  typename MmaSimtPolicy          ///< policy defining lane arrangement (concept: MmaSimtPolicy)
+>
+class TileIteratorSimt;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template for reading and writing tiles of accumulators to shared memory
+template <
+  typename WarpShape_,     ///< shape of warp-level GEMM (concept: GemmShape)
+  typename Operator_,      ///< matrix multiply operation (concept: arch::Mma)
+  typename Element_,       ///< data type of element to be written
+  typename MmaSimtPolicy_         ///< policy defining lane arrangement (concept: MmaSimtPolicy)
+>
+class TileIteratorSimt<WarpShape_, Operator_, Element_, layout::RowMajor, MmaSimtPolicy_> {
+public:
+
+  using WarpShape = WarpShape_;
+  using Operator = Operator_;
+  using Element = Element_;
+  using Layout = layout::RowMajor;
+
+  using TensorRef = TensorRef<Element, Layout>;         ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;                      ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  using Policy = SimtPolicy<WarpShape, Operator, Layout, MmaSimtPolicy_>;
+
+  /// Shape of the tile in memory
+  using Shape = MatrixShape<
+    Policy::kRowsPerIteration,
+    WarpShape::kN
+  >;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    typename Operator::ElementC, 
+    Policy::kElementsPerIteration>;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = Array<
+    typename Operator::ElementC, 
+    Policy::kAccumulatorElementCount>;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+  /// Padding quantity
+  using Padding = MatrixShape<
+    0,
+    4 * Policy::kElementsPerAccess
+#if CUTLASS_SIMT_EPILOGUE_USE_SCALAR_STORES
+    + 1
+#endif
+  >;
+
+private:
+
+#if CUTLASS_SIMT_EPILOGUE_USE_SCALAR_STORES
+  /// Storage type for accessing memory
+  using AccessType = AlignedArray<
+    Element, 
+    1
+  >;
+
+#else
+  /// Storage type for accessing memory
+  using AccessType = AlignedArray<
+    Element, 
+    Policy::kElementsPerAccess
+  >;
+#endif
+
+  //
+  // Data members
+  //
+
+  /// Internal pointer to memory
+  AccessType *pointer_;
+
+  /// Internal layout object
+  Layout layout_;
+
+public:
+
+  /// Default constructor
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimt(): pointer_(nullptr) { }
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimt(
+    TensorRef const &ref,
+    unsigned lane_id
+  ):
+    pointer_(reinterpret_cast<AccessType *>(ref.data())),
+    layout_(ref.stride()[0] / AccessType::kElements) { 
+
+    auto lane_layout = Policy::MmaSimtPolicy::get_lane_layout();
+    MatrixCoord lane_offset = lane_layout.inverse(lane_id);
+
+    pointer_ += layout_({
+      lane_offset.row(),
+      lane_offset.column() * Policy::kElementsPerAccess / int(AccessType::kElements)
+    });
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimt & add_pointer_offset(Index pointer_offset) {
+    pointer_ += pointer_offset / AccessType::kElements;
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimt & add_tile_offset(TensorCoord const &tile_offset) {
+
+    pointer_ += layout_({
+      tile_offset.row() * Shape::kRow, 
+      (tile_offset.column() * Shape::kColumn / int(AccessType::kElements))
+    });
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimt & operator+=(TensorCoord const &tile_offset) {
+
+    add_tile_offset(tile_offset);
+    
+    return *this;
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+#if CUTLASS_SIMT_EPILOGUE_USE_SCALAR_STORES
+      // de-vectorized stores
+      using ScalarAccessType = AlignedArray<Element, 1>;
+      ScalarAccessType const *scalarFragPtr = reinterpret_cast<ScalarAccessType const *>(&frag);
+      ScalarAccessType *scalarPointer = reinterpret_cast<ScalarAccessType *>(pointer_) + pointer_offset;
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int n = 0; n < Policy::kAccessesPerIteration; ++n) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int s = 0; s < Policy::kElementsPerAccess; s++) {
+          scalarPointer[n * Policy::MmaSimtPolicy::WarpShape::kColumn * Policy::kElementsPerAccess + s] = scalarFragPtr[n * Policy::kElementsPerAccess + s];
+        }
+      }
+#else
+    // original vector stores
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::kAccessesPerIteration; ++n) {
+      pointer_[n * Policy::MmaSimtPolicy::WarpShape::kColumn + pointer_offset / int(AccessType::kElements)] = frag_ptr[n];
+    }
+#endif
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) {
+    store_with_pointer_offset(frag, 0);
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const {
+
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::kAccessesPerIteration; ++n) {
+      frag_ptr[n] = pointer_[n * Policy::MmaSimtPolicy::WarpShape::kColumn + pointer_offset / int(AccessType::kElements)];
+    }
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag) const {
+    load_with_pointer_offset(frag, 0);
+  }
+
+  /// Set smem base address
+  CUTLASS_HOST_DEVICE
+  void set_smem_base_address(Index address) {
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template for reading and writing tiles of accumulators to shared memory
+template <typename WarpShape_,     ///< shape of warp-level GEMM (concept: GemmShape)
+          typename Operator_,      ///< matrix multiply operation (concept: arch::Mma)
+          typename Element_,       ///< data type of element to be written
+          typename Layout_,         ///< target shared memory layout
+          typename MmaSimtPolicy_  ///< policy defining lane arrangement (concept: MmaSimtPolicy)
+          >
+class TileIteratorSimtDirectConv {
+ public:
+
+  using WarpShape = WarpShape_;
+  using Operator = Operator_;
+  using Element = Element_;
+  using Layout = layout::RowMajor;
+
+  using TensorRef = TensorRef<Element, Layout>;  ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;               ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  using Policy = SimtPolicy<WarpShape, Operator, Layout, MmaSimtPolicy_>;
+
+  /// Shape of the tile in memory
+  using Shape = MatrixShape<Policy::kRowsPerIteration, WarpShape::kN>;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<typename Operator::ElementC, Policy::kElementsPerIteration>;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = Array<typename Operator::ElementC, Policy::kAccumulatorElementCount>;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+  /// Padding quantity
+  using Padding = MatrixShape<0,
+                              0
+                              >;
+
+private:
+  /// Storage type for accessing memory
+  using AccessType = AlignedArray<
+    Element, 
+    Policy::kElementsPerAccess
+  >;
+
+  //
+  // Data members
+  //
+
+  /// Internal pointer to memory
+  AccessType *pointer_;
+
+  /// Internal layout object
+  Layout layout_;
+
+  /// Base smem offset;
+  Index base_smem_address_;
+
+ public:
+  /// Default constructor
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtDirectConv() : pointer_(nullptr) {}
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtDirectConv(
+    TensorRef const &ref,
+    unsigned lane_id
+  ):
+    pointer_(reinterpret_cast<AccessType *>(ref.data())),
+    layout_(ref.stride()[0] / AccessType::kElements) {
+
+    auto lane_layout = Policy::MmaSimtPolicy::get_lane_layout();
+    MatrixCoord lane_offset = lane_layout.inverse(lane_id);
+
+    pointer_ += layout_({
+      lane_offset.row(),
+      lane_offset.column() * Policy::kElementsPerAccess / int(AccessType::kElements)
+    });
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtDirectConv & add_pointer_offset(Index pointer_offset) {
+    pointer_ += pointer_offset / AccessType::kElements;
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtDirectConv & add_tile_offset(TensorCoord const &tile_offset) {
+
+    pointer_ += layout_({
+      tile_offset.row() * Shape::kRow, 
+      (tile_offset.column() * Shape::kColumn / int(AccessType::kElements))
+    });
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtDirectConv & operator+=(TensorCoord const &tile_offset) {
+
+    add_tile_offset(tile_offset);
+    
+    return *this;
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+
+    // original vector stores
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+    AccessType * load_pointer_ = reinterpret_cast<AccessType *>(reinterpret_cast<uint8_t *>(pointer_) + base_smem_address_);
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::kAccessesPerIteration; ++n) {
+      load_pointer_[n * Policy::MmaSimtPolicy::WarpShape::kColumn + pointer_offset / int(AccessType::kElements)] = frag_ptr[n];
+    }
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) {
+    store_with_pointer_offset(frag, 0);
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const {
+
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::kAccessesPerIteration; ++n) {
+      frag_ptr[n] = pointer_[n * Policy::MmaSimtPolicy::WarpShape::kColumn + pointer_offset / int(AccessType::kElements)];
+    }
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag) const {
+    load_with_pointer_offset(frag, 0);
+  }
+
+  /// Set smem base address
+  CUTLASS_HOST_DEVICE
+  void set_smem_base_address(Index address){
+    base_smem_address_ = address;
+  }
+
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+/// Template for reading and writing tiles of accumulators to shared memory
+template <typename WarpShape_,               ///< shape of warp-level GEMM (concept: GemmShape)
+          typename ThreadOutputShape_,       /// Size of the matrix to load (concept: TensorNHWC)
+          typename ThreadBlockOutputShape_,  /// Size of the matrix to load (concept: TensorNHWC)
+          typename Operator_,                ///< matrix multi ply operation (concept: arch::Mma)
+          typename Element_,                 ///< data type of element to be written
+          typename Layout_,                  ///< target shared memory layout
+          typename MmaSimtPolicy_            ///< policy defining lane arrangement (concept: MmaSimtPolicy)
+          >
+class TileIteratorSimtDirect2dConv {
+ public:
+  using WarpShape = WarpShape_;
+  using ThreadOutputShape = ThreadOutputShape_;
+  using ThreadBlockOutputShape = ThreadBlockOutputShape_;
+  using Operator = Operator_;
+  using Element = Element_;
+  using Layout = layout::RowMajor;
+  using MmaSimtPolicy = MmaSimtPolicy_;
+
+  using TensorRef = TensorRef<Element, Layout>;  ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;               ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  // Thread-level shape of a fragment
+  using ThreadShape = MatrixShape<ThreadOutputShape::kNHW, ThreadOutputShape::kC>;
+
+  static_assert(!(ThreadShape::kColumn % MmaSimtPolicy::LaneMmaShape::kN),
+                "Thread-level GEMM must be divisible by Policy::LaneMmaShape.");
+
+  using ThreadTileCount = MatrixShape<ThreadBlockOutputShape::kH / ThreadOutputShape::kH,
+                                      ThreadBlockOutputShape::kW / ThreadOutputShape::kW>;
+
+  using Iterations =
+      MatrixShape<ThreadShape::kRow, ThreadShape::kColumn / MmaSimtPolicy::LaneMmaShape::kN>;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = typename Operator::FragmentC;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = AccumulatorTile;
+
+  /// Padding quantity
+  using Padding = MatrixShape<0, 0>;
+
+ private:
+  // Storage type for accessing memory
+  using AccessType = AlignedArray<Element, MmaSimtPolicy::LaneMmaShape::kN>;
+  //
+  // Data members
+  //
+
+  /// Internal pointer to memory
+  AccessType *pointer_;
+
+  /// Internal layout object
+  Layout layout_;
+
+  /// Base smem offset;
+  Index base_smem_address_;
+
+ public:
+  /// Default constructor
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtDirect2dConv() : pointer_(nullptr) {}
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtDirect2dConv(TensorRef const &ref, unsigned thread_id, unsigned lane_id)
+      : pointer_(reinterpret_cast<AccessType *>(ref.data())),
+        layout_(ref.stride()[0] / AccessType::kElements) {
+  
+    auto lane_layout = MmaSimtPolicy::get_lane_layout();
+
+    MatrixCoord lane_offset = lane_layout.inverse(lane_id);
+
+    // Get base HW offset of current threads
+    const int threadgroup = thread_id / (ThreadBlockOutputShape::kC / ThreadOutputShape::kC);
+    const int base_p = (threadgroup / (ThreadTileCount::kColumn)) * ThreadOutputShape::kH;
+    const int base_q = (threadgroup % (ThreadTileCount::kColumn)) * ThreadOutputShape::kW;
+
+    const int row_offset = base_p * ThreadBlockOutputShape::kW + base_q;
+
+    pointer_ += layout_(
+        {row_offset,
+         lane_offset.column() * MmaSimtPolicy::LaneMmaShape::kN / int(AccessType::kElements)});
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtDirect2dConv &add_pointer_offset(Index pointer_offset) {
+    pointer_ += pointer_offset / AccessType::kElements;
+    return *this;
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+    AccessType *storer_pointer_ =
+        reinterpret_cast<AccessType *>(reinterpret_cast<uint8_t *>(pointer_) + base_smem_address_);
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int h = 0; h < ThreadOutputShape::kH; ++h) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int w = 0; w < ThreadOutputShape::kW; ++w) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int col = 0; col < Iterations::kColumn; ++col) {
+          int offset = (w + h * ThreadBlockOutputShape::kW) *
+                           (ThreadBlockOutputShape::kC / AccessType::kElements) +
+                       col;
+          storer_pointer_[offset + pointer_offset / int(AccessType::kElements)] =
+              frag_ptr[w + h * ThreadOutputShape::kW + col];
+        }
+      }
+    }
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); }
+
+  /// Set smem base address
+  CUTLASS_HOST_DEVICE
+  void set_smem_base_address(Index address) { base_smem_address_ = address; }
+};
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template for reading and writing tiles of accumulators to shared memory
+template <
+  typename WarpShape_,        ///< shape of warp-level GEMM (concept: GemmShape)
+  typename Operator_,         ///< matrix multiply operation (concept: arch::Mma)
+  typename Element_,          ///< data type of element to be written
+  typename Layout_,            ///< target shared memory layout
+  typename MmaSimtPolicy_     ///< policy defining lane arrangement (concept: MmaSimtPolicy)
+>
+class TileIteratorSimtCanonical {
+public:
+
+  using WarpShape = WarpShape_;
+  using Operator = Operator_;
+  using Element = Element_;
+  using Layout = Layout_;
+
+  using TensorRef = TensorRef<Element, Layout>;         ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;                      ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  using Policy = SimtPolicy<WarpShape, Operator, Layout, MmaSimtPolicy_>;
+
+  /// Shape of the tile in memory
+  using Shape = MatrixShape<
+    Policy::kRowsPerIteration,
+    WarpShape::kN
+  >;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    typename Operator::ElementC, 
+    Policy::kElementsPerIteration>;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = Array<
+    typename Operator::ElementC, 
+    Policy::kAccumulatorElementCount>;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+  /// Padding quantity
+  using Padding = MatrixShape<
+    0,
+    4 * Policy::kElementsPerAccess + 1
+  >;
+
+private:
+
+  /// Storage type for accessing memory
+  using AccessType = AlignedArray<
+    Element, 
+    1
+  >;
+
+  //
+  // Data members
+  //
+
+  /// Internal pointer to memory
+  AccessType *pointer_;
+
+  /// Internal layout object
+  Layout layout_;
+
+  /// Guard to indicate whether the shape is divisible
+  bool divisible_;
+
+  /// Extent of the output tensor
+  MatrixCoord extent_;
+
+  /// Thread offset
+  MatrixCoord thread_offset_;
+
+public:
+
+  /// Default constructor
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtCanonical(): pointer_(nullptr) { }
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtCanonical(
+    TensorRef const &ref,
+    unsigned lane_id
+  ):
+    pointer_(reinterpret_cast<AccessType *>(ref.data())),
+    layout_(ref.stride()[0] / AccessType::kElements),
+    divisible_(true),
+    extent_(WarpShape::kM, WarpShape::kN) { 
+
+    auto lane_layout = Policy::MmaSimtPolicy::get_lane_layout();
+    MatrixCoord lane_offset = lane_layout.inverse(lane_id);
+
+    thread_offset_ = {
+      lane_offset.row() * Shape::kRow, 
+      lane_offset.column() * Policy::kElementsPerAccess
+    };
+
+    pointer_ += layout_({
+      lane_offset.row() * Shape::kRow,
+      lane_offset.column() * Policy::kElementsPerAccess / int(AccessType::kElements)
+    });
+  }
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtCanonical(
+    TensorRef const &ref,
+    TensorCoord const &extent,
+    unsigned lane_id
+  ):
+    pointer_(reinterpret_cast<AccessType *>(ref.data())),
+    layout_(ref.stride()[0] / AccessType::kElements),
+    divisible_(false),
+    extent_(extent) { 
+
+    auto lane_layout = Policy::MmaSimtPolicy::get_lane_layout();
+    MatrixCoord lane_offset = lane_layout.inverse(lane_id);
+
+    thread_offset_ = {
+      lane_offset.row() * Shape::kRow, 
+      lane_offset.column() * Policy::kElementsPerAccess
+    };
+
+    pointer_ += layout_({
+      lane_offset.row() * Shape::kRow,
+      lane_offset.column() * Policy::kElementsPerAccess / int(AccessType::kElements)
+    });
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtCanonical & add_pointer_offset(Index pointer_offset) {
+    pointer_ += pointer_offset / AccessType::kElements;
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtCanonical & add_tile_offset(TensorCoord const &tile_offset) {
+
+    MatrixCoord coord_offset(
+      tile_offset.row(), 
+      tile_offset.column() * Shape::kColumn
+    );
+
+    thread_offset_ += coord_offset;
+
+    pointer_ += layout_({
+      coord_offset.row(), 
+      coord_offset.column()
+    });
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtCanonical & operator+=(TensorCoord const &tile_offset) {
+
+    add_tile_offset(tile_offset);
+    
+    return *this;
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+
+    // de-vectorized stores
+    using ScalarAccessType = AlignedArray<Element, 1>;
+    ScalarAccessType const *scalarFragPtr = reinterpret_cast<ScalarAccessType const *>(&frag);
+    ScalarAccessType *scalarPointer = reinterpret_cast<ScalarAccessType *>(pointer_) + pointer_offset;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::kAccessesPerIteration; ++n) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int s = 0; s < Policy::kElementsPerAccess; s++) {
+        
+        int ptr_idx = n * Policy::MmaSimtPolicy::WarpShape::kColumn * Policy::kElementsPerAccess + s;
+        int frag_idx = n * Policy::kElementsPerAccess + s;
+        
+        int col = thread_offset_.column() + ptr_idx;
+
+        if (divisible_ || (thread_offset_.row() < extent_.row() && col < extent_.column())) {
+          scalarPointer[ptr_idx] = scalarFragPtr[frag_idx];
+        }
+      }
+    }
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) {
+    store_with_pointer_offset(frag, 0);
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const {
+
+      // de-vectorized loads
+      using ScalarAccessType = AlignedArray<Element, 1>;
+      ScalarAccessType *scalarFragPtr = reinterpret_cast<ScalarAccessType *>(&frag);
+      ScalarAccessType const *scalarPointer = reinterpret_cast<ScalarAccessType const*>(pointer_) + pointer_offset;
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int n = 0; n < Policy::kAccessesPerIteration; ++n) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int s = 0; s < Policy::kElementsPerAccess; s++) {
+          
+          int ptr_idx = n * Policy::MmaSimtPolicy::WarpShape::kColumn * Policy::kElementsPerAccess + s;
+          int frag_idx = n * Policy::kElementsPerAccess + s;
+          
+          int col = thread_offset_.column() + ptr_idx;
+
+          if (divisible_ || (thread_offset_.row() < extent_.row() && col < extent_.column())) {
+            scalarFragPtr[frag_idx] = scalarPointer[ptr_idx];
+          }
+        }
+      }
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag) const {
+    load_with_pointer_offset(frag, 0);
+  }
+
+  CUTLASS_HOST_DEVICE
+  TileIteratorSimtCanonical & operator++() {
+    return add_tile_offset({1, 0});
+  }
+
+  /// Set smem base address
+  CUTLASS_HOST_DEVICE
+  void set_smem_base_address(Index address) {
+  }
+};
+
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/tile_iterator_tensor_op.h

@@ -0,0 +1,671 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief 
+*/
+
+#pragma once
+
+#include "cutlass/array.h"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/layout/pitch_linear.h"
+
+#include "cutlass/epilogue/warp/tensor_op_policy.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template for reading and writing tiles of accumulators to shared memory
+template <
+  typename WarpShape,     ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename OperatorShape, ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  typename Element,       ///< data type of element to be written
+  typename Layout         ///< target shared memory layout
+>
+class TileIteratorTensorOp;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template for reading and writing tiles of accumulators to shared memory
+template <
+  typename WarpShape_,     ///< shape of warp-level GEMM (concept: GemmShape)
+  typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  typename Element_        ///< data type of element to be written
+>
+class TileIteratorTensorOp<WarpShape_, OperatorShape_, Element_, layout::RowMajor> {
+public:
+
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using Element = Element_;
+  using Layout = layout::RowMajor;
+
+  using TensorLayout = Layout;
+  using TensorRef = TensorRef<Element, Layout>;         ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;                      ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  /// Shape of the tile in memory
+  using Shape = MatrixShape<
+    Policy::kRowsPerIteration,
+    WarpShape::kN
+  >;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    Element, 
+    Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
+
+  /// This is the complete warp-level accumulator tile.
+  //using AccumulatorTile = typename Operator::FragmentC;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+  /// Number of times this iterator can be incremented
+  using TileIterations = typename Policy::TileIterations;
+
+  // Internal constants
+  struct Detail {
+    static int const kLanesInQuad = 4;
+  };
+
+  /// Padding quantity
+  using Padding = MatrixShape<
+    0,
+    Detail::kLanesInQuad * Policy::kElementsPerAccess>;
+
+private:
+
+  /// Storage type for accessing memory
+  using AccessType = AlignedArray<Element, Policy::kElementsPerAccess>;
+
+  //
+  // Data members
+  //
+
+  /// Internal pointer to memory
+  AccessType *pointer_;
+
+  /// Internal layout object
+  Layout layout_;
+
+  /// Thread offset
+  MatrixCoord thread_offset_;
+
+public:
+
+  /// Default constructor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOp(): pointer_(nullptr) { }
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOp(
+    TensorRef const &ref,
+    unsigned lane_id
+  ):
+    pointer_(reinterpret_cast<AccessType *>(ref.data())),
+    layout_(ref.stride()[0] / Policy::kElementsPerAccess) {
+
+    int quad_id = (lane_id / Detail::kLanesInQuad); 
+    int lane_in_quad = (lane_id % Detail::kLanesInQuad);
+
+    thread_offset_ = {
+      quad_id, lane_in_quad * Policy::kElementsPerAccess
+    };
+
+    pointer_ += layout_({thread_offset_.row(), thread_offset_.column() / Policy::kElementsPerAccess});
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOp & add_pointer_offset(Index pointer_offset) {
+    pointer_ += pointer_offset / Policy::kElementsPerAccess;
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOp & add_tile_offset(TensorCoord const &tile_offset) {
+
+    MatrixCoord coord_offset(
+      tile_offset.row() * Shape::kRow, 
+      tile_offset.column() * Shape::kColumn
+    );
+
+    thread_offset_ += coord_offset;
+
+    pointer_ += layout_({
+      coord_offset.row(),
+      coord_offset.column() / Policy::kElementsPerAccess
+    });
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOp & operator+=(TensorCoord const &tile_offset) {
+    add_tile_offset(tile_offset);
+    return *this;
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
+      pointer_[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess] = frag_ptr[n];
+    }
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) {
+    store_with_pointer_offset(frag, 0);
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const {
+
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
+      frag_ptr[n] = pointer_[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess];
+    }
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag) const {
+    load_with_pointer_offset(frag, 0);
+  }
+
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOp & operator++() {
+    return add_tile_offset({1, 0});
+  }
+  
+  /// Set smem base address
+  CUTLASS_HOST_DEVICE
+  void set_smem_base_address(Index address) {
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template for reading and writing tiles of accumulators to shared memory
+template <
+  typename WarpShape_,     ///< shape of warp-level GEMM (concept: GemmShape)
+  typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  typename Element_,       ///< data type of element to be written
+  int InterleavedK         ///< number of interleaved k
+>
+class TileIteratorTensorOp<WarpShape_, OperatorShape_, Element_, 
+                            layout::ColumnMajorInterleaved<InterleavedK> > {
+public:
+
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using Element = Element_;
+  using Layout = layout::ColumnMajorInterleaved<InterleavedK>;
+  using TensorLayout = Layout;                ///< shared memory tensor ref layout
+
+  using TensorRef = TensorRef<Element, TensorLayout>;         ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;                      ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  /// Shape of the tile in memory
+  using Shape = MatrixShape<
+//    Policy::kRowsPerIteration,
+    WarpShape::kM,
+    InterleavedK
+  >;
+
+  /// This is the fragment size produced by one tile
+  using Fragment = Array<
+    Element, 
+    Policy::OperatorCount::kRow * Policy::kIterationsPerInstruction 
+        * Policy::kElementsPerIteration>;
+
+  /// This is the fragment size produced by one iteration
+//  using Fragment = Array<
+//    Element, Policy::kElementsPerIteration >;
+
+  /// This is the complete warp-level accumulator tile.
+  //using AccumulatorTile = typename Operator::FragmentC;
+
+  /// Number of times this iterator can be incremented
+  using TileIterations = typename Policy::TileIterations;
+
+  // Internal constants
+  struct Detail {
+    static int const kLanesInQuad = 4;
+  };
+
+  /// Padding quantity
+  using Padding = MatrixShape<
+    0,
+    Detail::kLanesInQuad * Policy::kElementsPerIteration>;
+
+private:
+
+  /// Storage type for accessing memory
+  using AccessType = AlignedArray<Element, Policy::kElementsPerAccess>;
+
+  //
+  // Data members
+  //
+
+  /// Internal pointer to memory
+  AccessType *pointer_;
+
+  /// Internal layout object
+  TensorLayout layout_;
+
+  /// Thread offset
+  MatrixCoord thread_offset_;
+
+public:
+
+  /// Default constructor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOp(): pointer_(nullptr) { }
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOp(
+    TensorRef const &ref,
+    unsigned lane_id
+  ):
+    pointer_(reinterpret_cast<AccessType *>(ref.data())),
+    layout_(ref.stride()[0]) {
+
+    int quad_id = (lane_id / Detail::kLanesInQuad); 
+    int lane_in_quad = (lane_id % Detail::kLanesInQuad);
+
+    thread_offset_ = {
+      quad_id, lane_in_quad * Policy::kElementsPerIteration
+    };
+
+    pointer_ += (layout_({thread_offset_.row(), thread_offset_.column()}) / Policy::kElementsPerAccess);
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOp & add_pointer_offset(Index pointer_offset) {
+    pointer_ += pointer_offset / Policy::kElementsPerAccess;
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOp & add_tile_offset(TensorCoord const &tile_offset) {
+
+    MatrixCoord coord_offset(
+      tile_offset.row() * Shape::kRow, 
+      tile_offset.column() * Shape::kColumn
+    );
+
+    thread_offset_ += coord_offset;
+
+    pointer_ += (layout_({
+      coord_offset.row(),
+      coord_offset.column()
+    }) / Policy::kElementsPerAccess);
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOp & operator+=(TensorCoord const &tile_offset) {
+    add_tile_offset(tile_offset);
+    return *this;
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+      
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::OperatorCount::kRow * Policy::kIterationsPerInstruction; n++ ) {
+
+      AccessType *ptr = pointer_ + layout_({n * Policy::kRowsPerIteration, 0}) / Policy::kElementsPerAccess;
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int a = 0; a < Policy::kAccessPerIteration; ++a) {
+        ptr[a + pointer_offset / Policy::kElementsPerAccess] = frag_ptr[n * Policy::kAccessPerIteration + a];
+
+//        printf("store thread %d, address %p, bank %ld\n", threadIdx.x, pointer_+a+n*Detail::kLanesInQuad, 
+//            ((long long)(pointer_+a+n*Detail::kLanesInQuad)>>2)&0x1f);
+      }
+    }
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) {
+    store_with_pointer_offset(frag, 0);
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const {
+
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::OperatorCount::kRow * Policy::kIterationsPerInstruction; n++ ) {
+
+      AccessType *ptr = pointer_ + layout_({n * Policy::kRowsPerIteration, 0}) / Policy::kElementsPerAccess;
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int a = 0; a < Policy::kAccessPerIteration; ++a) {
+        frag_ptr[n * Policy::kAccessPerIteration + a] = ptr[a + pointer_offset / Policy::kElementsPerAccess];
+      }
+    }
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag) const {
+    load_with_pointer_offset(frag, 0);
+  }
+
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOp & operator++() {
+    return add_tile_offset({0, 1});
+  }
+
+  /// Set smem base address
+  CUTLASS_HOST_DEVICE
+  void set_smem_base_address(Index address) {
+  }
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template for reading and writing tiles of accumulators to shared memory
+template <
+  typename WarpShape_,     ///< shape of warp-level GEMM (concept: GemmShape)
+  typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  typename Element_,       ///< data type of element to be written
+  typename Layout_
+>
+class TileIteratorTensorOpCanonical {
+public:
+
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using Element = Element_;
+  using Layout = Layout_;
+
+  using TensorRef = TensorRef<Element, Layout>;         ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;                      ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  static int const kAccessSize = 1;
+  static int const kAccessCount = Policy::kElementsPerAccess / kAccessSize;
+
+  /// Shape of the tile in memory
+  using Shape = MatrixShape<
+    Policy::kRowsPerIteration,
+    WarpShape::kN
+  >;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    Element, 
+    Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
+
+  /// This is the complete warp-level accumulator tile.
+  //using AccumulatorTile = typename Operator::FragmentC;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+  // Internal constants
+  struct Detail {
+    static int const kLanesInQuad = 4;
+  };
+
+  /// Padding quantity
+  using Padding = MatrixShape<
+    0,
+    Detail::kLanesInQuad * Policy::kElementsPerAccess>;
+
+private:
+
+  /// Storage type for accessing memory
+  using AccessType = AlignedArray<Element, kAccessSize>;
+
+  //
+  // Data members
+  //
+
+  /// Internal pointer to memory
+  AccessType *pointer_;
+
+  /// Internal layout object
+  Layout layout_;
+
+  /// Guard to indicate whether the shape is divisible
+  bool divisible_;
+
+  /// Extent of the output tensor
+  MatrixCoord extent_;
+
+  /// Thread offset
+  MatrixCoord thread_offset_;
+
+public:
+
+  /// Default constructor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpCanonical(): pointer_(nullptr) { }
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpCanonical(
+    TensorRef const &ref,
+    unsigned lane_id
+  ):
+    pointer_(reinterpret_cast<AccessType *>(ref.data())),
+    layout_(ref.stride()[0]),
+    divisible_(true),
+    extent_(WarpShape::kM, WarpShape::kN) {
+
+    int quad_id = (lane_id / Detail::kLanesInQuad); 
+    int lane_in_quad = (lane_id % Detail::kLanesInQuad);
+
+    thread_offset_ = {
+      quad_id, lane_in_quad * Policy::kElementsPerAccess
+    };
+
+    pointer_ += layout_({thread_offset_.row(), thread_offset_.column()});
+  }
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpCanonical(
+    TensorRef const &ref,
+    TensorCoord const &extent,
+    unsigned lane_id
+  ):
+    pointer_(reinterpret_cast<AccessType *>(ref.data())),
+    layout_(ref.stride()[0]),
+    divisible_(false),
+    extent_(extent) {
+
+    int quad_id = (lane_id / Detail::kLanesInQuad); 
+    int lane_in_quad = (lane_id % Detail::kLanesInQuad);
+
+    thread_offset_ = {
+      quad_id, lane_in_quad * Policy::kElementsPerAccess
+    };
+
+    pointer_ += layout_({thread_offset_.row(), thread_offset_.column()});
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpCanonical & add_pointer_offset(Index pointer_offset) {
+    pointer_ += pointer_offset;
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpCanonical & add_tile_offset(TensorCoord const &tile_offset) {
+
+    MatrixCoord coord_offset(
+      tile_offset.row() * Shape::kRow, 
+      tile_offset.column() * Shape::kColumn
+    );
+
+    thread_offset_ += coord_offset;
+
+    pointer_ += layout_({
+      coord_offset.row(),
+      coord_offset.column()
+    });
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpCanonical & operator+=(TensorCoord const &tile_offset) {
+    add_tile_offset(tile_offset);
+    return *this;
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int a = 0; a < kAccessCount; ++a) {
+
+        int ptr_idx = n * Detail::kLanesInQuad * kAccessCount + pointer_offset + a;
+        int frag_idx = n * kAccessCount + a;
+
+        int col = thread_offset_.column() + n * Detail::kLanesInQuad * Policy::kElementsPerAccess + a;
+
+        if (divisible_ || (thread_offset_.row() < extent_.row() && col < extent_.column())) {
+          pointer_[ptr_idx] = frag_ptr[frag_idx];
+        }
+      }
+    }
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) {
+    store_with_pointer_offset(frag, 0);
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const {
+
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+    
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int a = 0; a < kAccessCount; ++a) {
+
+        int ptr_idx = n * Detail::kLanesInQuad * kAccessCount + pointer_offset + a;
+        int frag_idx = n * kAccessCount + a;
+        
+        int col = thread_offset_.column() + n * Detail::kLanesInQuad * Policy::kElementsPerAccess + a;
+
+        if (divisible_ || (thread_offset_.row() < extent_.row() && col < extent_.column())) {
+          frag_ptr[frag_idx] = pointer_[ptr_idx];
+        }
+      }
+    }
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag) const {
+    load_with_pointer_offset(frag, 0);
+  }
+
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpCanonical & operator++() {
+    return add_tile_offset({1, 0});
+  }
+  
+  /// Set smem base address
+  CUTLASS_HOST_DEVICE
+  void set_smem_base_address(Index address) {
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/tile_iterator_tensor_op_mixed.h

@@ -0,0 +1,1089 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief 
+*/
+
+#pragma once
+
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/layout/pitch_linear.h"
+
+#include "cutlass/arch/memory_sm75.h"
+#include "cutlass/epilogue/warp/tensor_op_policy.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// This is an optimization available on CUDA 11.2 and beyond that eliminates branches in the epilogue.
+#define CUTLASS_EPILOGUE_WARP_TILE_ITERATOR_TENSOR_OP_MIXED_OPTIMIZATION_ENABLED ((__CUDACC_VER_MAJOR__ * 10 + __CUDACC_VER_MINOR__) >= 112)
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template for reading and writing tiles of accumulators to shared memory. This is optimized
+/// for mixed-precision epilogues in which the accumulators are 32b in width, but the output
+/// data type is smaller. 
+template <
+  typename WarpShape_,            ///< shape of warp-level GEMM (concept: GemmShape)
+  typename OperatorShape_,        ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  typename Element_,              ///< data type of accumulator element
+  int ElementSizeBits,            ///< Size of accumulator element in bits
+  int OutputSizeBits,             ///< Size of output element in bits
+  int OutputElementCount,         ///< number of elements in output vector
+  int ContiguousLanes,            ///< Number of consecutive lanes writing to contiguous memory
+  bool EightBitsOutputOrLess = (OutputSizeBits <= 8)
+>
+class TileIteratorTensorOpMixed {
+public:
+
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using Element = Element_;
+  using Layout = layout::RowMajor;
+  static int const kOutputElementCount = OutputElementCount;
+
+  using TensorRef = TensorRef<Element, Layout>;         ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;                      ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  /// Shape of the tile in memory
+  using Shape = MatrixShape<
+    Policy::kRowsPerIteration,
+    WarpShape::kN
+  >;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    Element, 
+    Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
+
+  /// This is the complete warp-level accumulator tile.
+  //using AccumulatorTile = typename Operator::FragmentC;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+  // Internal constants
+  struct Detail {
+    static int const kLanesInQuad = 4;
+
+    /// Number of pointers needed to write accumulators
+    static int const kPointerCount = 
+      (OutputElementCount * sizeof_bits<Element>::value) / (const_min(128, OutputElementCount * sizeof_bits<Element>::value));
+
+    // Currently support max 4 ptr
+    static constexpr int kMaxPointerCount{4};
+
+    static_assert(kPointerCount <= kMaxPointerCount, "Can only accommodate four pointers at present.");
+    static_assert(sizeof(Element) == 4, "This can only be used with 32b accumulator data types (f32, s32).");
+  };
+
+  /// Padding quantity
+  using Padding = MatrixShape<
+    0,
+    Detail::kLanesInQuad * Policy::kElementsPerAccess>;
+
+private:
+
+  /// Storage type for accessing memory
+  using AccessType = AlignedArray<Element, Policy::kElementsPerAccess>;
+
+  //
+  // Data members
+  //
+
+  /// Internal pointer to memory
+  AccessType *pointers_[Detail::kPointerCount] = {nullptr};
+
+  /// Stride in units of AccessType
+  int stride_{0};
+
+  /// Logical column in which warp tile is aligned
+  int warp_column_{0};
+
+public:
+
+  /// Default constructor
+  TileIteratorTensorOpMixed() = default;
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed(
+    TensorRef const &ref,
+    unsigned lane_id
+  ):
+    stride_(ref.stride()[0] / Policy::kElementsPerAccess),
+    warp_column_(0) { 
+
+    int quad_id = (lane_id / Detail::kLanesInQuad); 
+    int lane_in_quad = (lane_id % Detail::kLanesInQuad);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
+      AccessType *ptr = reinterpret_cast<AccessType *>(ref.data()) + quad_id * stride_;
+      int column_idx = (lane_in_quad % 2) + (((lane_in_quad / 2) + i) % Detail::kPointerCount) * 2;
+
+      ptr += column_idx;
+
+      pointers_[i % Detail::kPointerCount] = ptr;
+    }
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & add_pointer_offset(Index pointer_offset) {
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
+      pointers_[i] += pointer_offset / Policy::kElementsPerAccess;
+    }
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & add_tile_offset(TensorCoord const &tile_offset) {
+    
+    CUTLASS_PRAGMA_UNROLL
+    for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
+      pointers_[i] += tile_offset.row() * Shape::kRow * stride_ + 
+        tile_offset.column() * Shape::kColumn / Policy::kElementsPerAccess;
+    }
+
+    warp_column_ += tile_offset.column() * Shape::kColumn;
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & operator+=(TensorCoord const &tile_offset) {
+    return add_tile_offset(tile_offset);
+  }
+
+  /// Store
+  CUTLASS_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    AccessType *ptr = pointers_[0];
+
+#if CUTLASS_EPILOGUE_WARP_TILE_ITERATOR_TENSOR_OP_MIXED_OPTIMIZATION_ENABLED
+
+    // When the optimization is enabled, small tiles require separate logic.
+    bool kN32_optimization = (WarpShape::kN * Detail::kLanesInQuad * Policy::kElementsPerAccess * sizeof_bits<Element>::value) % 1024 == 0;
+    if (kN32_optimization) {
+      
+      int ptr_idx = ((warp_column_ * sizeof_bits<Element>::value) / 1024) % Detail::kPointerCount;
+      
+      if (ptr_idx == 0) {
+        ptr = pointers_[0];
+      } else if (ptr_idx == 1) {
+	if constexpr (AccessType::kElements >= 2) {
+          ptr = pointers_[1];
+	}
+      } else if (ptr_idx == 2) {
+	if constexpr (AccessType::kElements >= 3) {
+          ptr = pointers_[2];
+	}
+      } else if (ptr_idx == 3) {
+	if constexpr (AccessType::kElements >= 4) {
+          ptr = pointers_[3];
+	}
+      }
+    }
+
+#endif
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int64_t n = 0; n < Policy::OperatorCount::kColumn; ++n) {
+      
+#if CUTLASS_EPILOGUE_WARP_TILE_ITERATOR_TENSOR_OP_MIXED_OPTIMIZATION_ENABLED
+
+      //
+      // When the optimization is enabled, this expression suffices to obtain the SMEM pointer.
+      //
+      if (WarpShape::kN == 64) {
+        ptr = pointers_[n / 4];
+      }
+      else if (!kN32_optimization)
+#endif
+      {
+        // This is the reference implementation
+        int column_idx = warp_column_ + n * Detail::kLanesInQuad * Policy::kElementsPerAccess;
+        int ptr_idx = ((column_idx * sizeof_bits<Element>::value) / 1024) % Detail::kPointerCount;
+  
+        if (ptr_idx == 0) {
+          ptr = pointers_[0 % Detail::kPointerCount];
+        }
+        else if (ptr_idx == 1) {
+          ptr = pointers_[1 % Detail::kPointerCount];
+        }
+        else if (ptr_idx == 2) {
+          ptr = pointers_[2 % Detail::kPointerCount];
+        }
+        else if (ptr_idx == 3) {
+          ptr = pointers_[3 % Detail::kPointerCount];
+        }
+      }
+
+      int offset = n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess;
+      ptr[offset] = frag_ptr[n];
+    }
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) {
+    store_with_pointer_offset(frag, 0);
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const {
+
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int64_t n = 0; n < Policy::OperatorCount::kColumn; ++n) {
+
+      int column_idx = warp_column_ + n * Detail::kLanesInQuad * Policy::kElementsPerAccess;
+      int ptr_idx = ((column_idx * sizeof_bits<Element>::value) / 1024) % Detail::kPointerCount;
+
+      AccessType const *smem_ptr = pointers_[ptr_idx];
+      frag_ptr[n] = smem_ptr[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess];
+    }
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag) const {
+    load_with_pointer_offset(frag, 0);
+  }
+  
+  /// Set smem base address
+  CUTLASS_HOST_DEVICE
+  void set_smem_base_address(Index address) {
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for int32_t x 16 => int8_t/int4b_t x 16
+template <
+  typename WarpShape_,            ///< shape of warp-level GEMM (concept: GemmShape)
+  typename OperatorShape_,        ///< matrix multiply operation shape (concept: gemm::GemmShape),
+  int OutputSizeBits              ///< Size of output element in bits
+>
+class TileIteratorTensorOpMixed<WarpShape_, OperatorShape_, int32_t, 32, OutputSizeBits, 16, 8, true> {
+public:
+
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using Element = int32_t;
+  using Layout = layout::RowMajor;
+  static int const kOutputElementCount = 16;
+
+  using TensorRef = TensorRef<Element, Layout>;         ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;                      ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  /// Shape of the tile in memory
+  using Shape = MatrixShape<
+    Policy::kRowsPerIteration,
+    WarpShape::kN
+  >;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    Element, 
+    Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
+
+  /// This is the complete warp-level accumulator tile.
+  //using AccumulatorTile = typename Operator::FragmentC;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+  // Internal constants
+  struct Detail {
+    static int const kLanesInQuad = 4;
+
+    /// Number of pointers needed to write accumulators
+    static int const kPointerCount = 2;
+
+    /// Offsets added 
+    static int const kOffsetCount = 4;
+
+    static_assert(sizeof(Element) == 4, "This can only be used with 32b accumulator data types (f32, s32).");
+  };
+
+  /// Padding quantity
+  using Padding = MatrixShape<0, Detail::kLanesInQuad * 2>;
+
+private:
+
+  /// Storage type for accessing memory
+  using AccessType = AlignedArray<Element, 2>;
+
+  //
+  // Data members
+  //
+
+  /// Internal pointer to memory
+  AccessType *pointers_[Detail::kPointerCount] = {nullptr};
+
+  /// Stride in units of AccessType
+  int stride_{0};
+
+  /// Uniform offset in bytes added to warp tile iterator
+  int uniform_offset_[Detail::kOffsetCount] = {0};
+
+public:
+
+  /// Default constructor
+  TileIteratorTensorOpMixed() = default;
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed(
+    TensorRef const &ref,
+    unsigned lane_id
+  ):
+    stride_(ref.stride()[0] / AccessType::kElements) { 
+
+    int quad_id = (lane_id / Detail::kLanesInQuad); 
+    int lane_in_quad = (lane_id % Detail::kLanesInQuad);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Detail::kPointerCount; ++i) {
+      AccessType *ptr = reinterpret_cast<AccessType *>(ref.data()) + quad_id * stride_;
+      int column_idx = lane_in_quad ^ (i * 2);
+
+      ptr += column_idx;
+    
+      if (i == 0) {
+        pointers_[0] = ptr;
+      }
+      else if (i == 1) {
+        pointers_[1] = ptr;
+      }
+    }
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Detail::kOffsetCount; ++i) {
+      uniform_offset_[i] = (i ^ 0) * 4 * sizeof(AccessType);
+    }
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & add_pointer_offset(Index pointer_offset) {
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
+      pointers_[i] += pointer_offset / AccessType::kElements;
+    }
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & add_tile_offset(TensorCoord const &tile_offset) {
+    
+    int ptr_offset = tile_offset.row() * Shape::kRow * stride_ + 
+      tile_offset.column() * Shape::kColumn / AccessType::kElements;
+
+    pointers_[0] += ptr_offset;
+    pointers_[1] += ptr_offset;
+    
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Detail::kOffsetCount; ++i) {
+      uniform_offset_[i] = (i ^ tile_offset.column()) * 4 * sizeof(AccessType);
+    }
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & operator+=(TensorCoord const &tile_offset) {
+    return add_tile_offset(tile_offset);
+  }
+
+  /// Store
+  CUTLASS_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
+
+      int ptr_idx = (n / 4);
+      int offset_idx = (n % 4);
+
+      AccessType *ptr;
+      if (ptr_idx == 0) {
+        ptr = pointers_[0];
+      }
+      else if (ptr_idx == 1) {
+        ptr = pointers_[1];
+      }
+
+      int offset = (n / 4) * 16 + pointer_offset / AccessType::kElements;
+
+#if 0
+      //
+      // Using inline PTX to avoid generic memory
+      //
+      AccessType *smem_ptr = pointers_[ptr_idx];
+      smem_ptr[offset] = frag_ptr[n];
+#else
+      uint32_t smem_addr = arch::cutlass_get_smem_pointer(ptr);
+      uint32_t const *data = reinterpret_cast<uint32_t const *>(frag_ptr + n);
+      uint32_t offset_in_bytes = offset * sizeof(AccessType) + uniform_offset_[offset_idx];
+
+      asm volatile(
+        "{ .reg .u32 smem_ptr; add.u32 smem_ptr, %0, %1; st.shared.v2.u32 [smem_ptr], {%2, %3}; }\n"
+        : : "r"(smem_addr), "r"(offset_in_bytes), "r"(data[0]), "r"(data[1])
+      );
+#endif
+    }
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) {
+    store_with_pointer_offset(frag, 0);
+  }
+
+  /// Set smem base address
+  CUTLASS_HOST_DEVICE
+  void set_smem_base_address(Index address) {
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for int32_t x 8 => int8_t/int4b_t x 8
+template <
+  typename WarpShape_,            ///< shape of warp-level GEMM (concept: GemmShape)
+  typename OperatorShape_,        ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  int OutputSizeBits              ///< Size of output element in bits
+>
+class TileIteratorTensorOpMixed<WarpShape_, OperatorShape_, int32_t, 32, OutputSizeBits, 8, 8, true> {
+public:
+
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using Element = int32_t;
+  using Layout = layout::RowMajor;
+  static int const kOutputElementCount = 8;
+
+  using TensorRef = TensorRef<Element, Layout>;         ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;                      ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  /// Shape of the tile in memory
+  using Shape = MatrixShape<
+    Policy::kRowsPerIteration,
+    WarpShape::kN
+  >;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    Element, 
+    Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
+
+  /// This is the complete warp-level accumulator tile.
+  //using AccumulatorTile = typename Operator::FragmentC;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+  // Internal constants
+  struct Detail {
+    static int const kLanesInQuad = 4;
+
+    /// Number of pointers needed to write accumulators
+    static int const kPointerCount = 2;
+
+    static_assert(sizeof(Element) == 4, "This can only be used with 32b accumulator data types (f32, s32).");
+  };
+
+  /// Padding quantity
+  using Padding = MatrixShape<0, Detail::kLanesInQuad * 2>;
+
+private:
+
+  /// Storage type for accessing memory
+  using AccessType = AlignedArray<Element, 2>;
+
+  //
+  // Data members
+  //
+
+  /// Internal pointer to memory
+  AccessType *pointers_[Detail::kPointerCount] = {nullptr};
+
+  /// Stride in units of AccessType
+  int stride_{0};
+
+public:
+
+  /// Default constructor
+  TileIteratorTensorOpMixed() = default;
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed(
+    TensorRef const &ref,
+    unsigned lane_id
+  ):
+    stride_(ref.stride()[0] / AccessType::kElements) { 
+
+    int quad_id = (lane_id / Detail::kLanesInQuad); 
+    int lane_in_quad = (lane_id % Detail::kLanesInQuad);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Detail::kPointerCount; ++i) {
+      AccessType *ptr = reinterpret_cast<AccessType *>(ref.data()) + quad_id * stride_;
+      int column_idx = lane_in_quad ^ (i * 2);
+
+      ptr += column_idx;
+    
+      if (i == 0) {
+        pointers_[0] = ptr;
+      }
+      else if (i == 1) {
+        pointers_[1] = ptr;
+      }
+    }
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & add_pointer_offset(Index pointer_offset) {
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
+      pointers_[i] += pointer_offset / AccessType::kElements;
+    }
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & add_tile_offset(TensorCoord const &tile_offset) {
+    
+    int ptr_offset = tile_offset.row() * Shape::kRow * stride_ + 
+      tile_offset.column() * Shape::kColumn / AccessType::kElements;
+
+    pointers_[0] += ptr_offset;
+    pointers_[1] += ptr_offset;
+   
+    if (tile_offset.column() % 2) {
+      auto tmp = pointers_[0];
+      pointers_[0] = pointers_[1];
+      pointers_[1] = tmp;
+    }
+ 
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & operator+=(TensorCoord const &tile_offset) {
+    return add_tile_offset(tile_offset);
+  }
+
+  /// Store
+  CUTLASS_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
+
+      int ptr_idx = (n / 4);
+
+      AccessType *ptr;
+      if (ptr_idx == 0) {
+        ptr = pointers_[0];
+      }
+      else if (ptr_idx == 1) {
+        ptr = pointers_[1];
+      }
+
+      int offset = (n / 4) * 16 + pointer_offset / AccessType::kElements + (n % 4) * 4;
+
+#if 0
+      //
+      // Using inline PTX to avoid generic memory
+      //
+      AccessType *smem_ptr = pointers_[ptr_idx];
+      smem_ptr[offset] = frag_ptr[n];
+#else
+      uint32_t smem_addr = arch::cutlass_get_smem_pointer(ptr);
+      uint32_t const *data = reinterpret_cast<uint32_t const *>(frag_ptr + n);
+      uint32_t offset_in_bytes = offset * sizeof(AccessType);
+
+      asm volatile(
+        "{ .reg .u32 smem_ptr; add.u32 smem_ptr, %0, %1; st.shared.v2.u32 [smem_ptr], {%2, %3}; }\n"
+        : : "r"(smem_addr), "r"(offset_in_bytes), "r"(data[0]), "r"(data[1])
+      );
+#endif
+    }
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) {
+    store_with_pointer_offset(frag, 0);
+  }
+
+  /// Set smem base address
+  CUTLASS_HOST_DEVICE
+  void set_smem_base_address(Index address) {
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for float x 16 => float_e4m3_t/float_e5m2_t x 16
+template <
+  typename WarpShape_,            ///< shape of warp-level GEMM (concept: GemmShape)
+  typename OperatorShape_         ///< matrix multiply operation shape (concept: gemm::GemmShape),
+>
+class TileIteratorTensorOpMixed<WarpShape_, OperatorShape_, float, 32, 8, 16, 8> {
+public:
+
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using Element = float;
+  using Layout = layout::RowMajor;
+  static int const kOutputElementCount = 16;
+
+  using TensorRef = TensorRef<Element, Layout>;         ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;                      ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  /// Shape of the tile in memory
+  using Shape = MatrixShape<
+    Policy::kRowsPerIteration,
+    WarpShape::kN
+  >;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    Element,
+    Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
+
+  /// This is the complete warp-level accumulator tile.
+  //using AccumulatorTile = typename Operator::FragmentC;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+  // Internal constants
+  struct Detail {
+    static int const kLanesInQuad = 4;
+
+    /// Number of pointers needed to write accumulators
+    static int const kPointerCount = 2;
+
+    /// Offsets added
+    static int const kOffsetCount = 4;
+
+    static_assert(sizeof(Element) == 4, "This can only be used with 32b accumulator data types (f32, s32).");
+  };
+
+  /// Padding quantity
+  using Padding = MatrixShape<0, Detail::kLanesInQuad * 2>;
+
+private:
+
+  /// Storage type for accessing memory
+  using AccessType = AlignedArray<Element, 2>;
+
+  //
+  // Data members
+  //
+
+  /// Internal pointer to memory
+  AccessType *pointers_[Detail::kPointerCount] = {nullptr};
+
+  /// Stride in units of AccessType
+  int stride_{0};
+
+  /// Uniform offset in bytes added to warp tile iterator
+  int uniform_offset_[Detail::kOffsetCount] = {0};
+
+public:
+
+  /// Default constructor
+  TileIteratorTensorOpMixed() = default;
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed(
+    TensorRef const &ref,
+    unsigned lane_id
+  ):
+    stride_(ref.stride()[0] / AccessType::kElements) {
+
+    int quad_id = (lane_id / Detail::kLanesInQuad);
+    int lane_in_quad = (lane_id % Detail::kLanesInQuad);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Detail::kPointerCount; ++i) {
+      AccessType *ptr = reinterpret_cast<AccessType *>(ref.data()) + quad_id * stride_;
+      int column_idx = lane_in_quad ^ (i * 2);
+
+      ptr += column_idx;
+
+      if (i == 0) {
+        pointers_[0] = ptr;
+      }
+      else if (i == 1) {
+        pointers_[1] = ptr;
+      }
+    }
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Detail::kOffsetCount; ++i) {
+      uniform_offset_[i] = (i ^ 0) * 4 * sizeof(AccessType);
+    }
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & add_pointer_offset(Index pointer_offset) {
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
+      pointers_[i] += pointer_offset / AccessType::kElements;
+    }
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & add_tile_offset(TensorCoord const &tile_offset) {
+
+    int ptr_offset = tile_offset.row() * Shape::kRow * stride_ +
+      tile_offset.column() * Shape::kColumn / AccessType::kElements;
+
+    pointers_[0] += ptr_offset;
+    pointers_[1] += ptr_offset;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Detail::kOffsetCount; ++i) {
+      uniform_offset_[i] = (i ^ tile_offset.column()) * 4 * sizeof(AccessType);
+    }
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & operator+=(TensorCoord const &tile_offset) {
+    return add_tile_offset(tile_offset);
+  }
+
+  /// Store
+  CUTLASS_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
+
+      int ptr_idx = (n / 4);
+      int offset_idx = (n % 4);
+
+      AccessType *ptr;
+      if (ptr_idx == 0) {
+        ptr = pointers_[0];
+      }
+      else if (ptr_idx == 1) {
+        ptr = pointers_[1];
+      }
+
+      int offset = (n / 4) * 16 + pointer_offset / AccessType::kElements;
+
+#if 0
+      //
+      // Using inline PTX to avoid generic memory
+      //
+      AccessType *smem_ptr = pointers_[ptr_idx];
+      smem_ptr[offset] = frag_ptr[n];
+#else
+      uint32_t smem_addr = arch::cutlass_get_smem_pointer(ptr);
+      uint32_t const *data = reinterpret_cast<uint32_t const *>(frag_ptr + n);
+      uint32_t offset_in_bytes = offset * sizeof(AccessType) + uniform_offset_[offset_idx];
+
+      asm volatile(
+        "{ .reg .u32 smem_ptr; add.u32 smem_ptr, %0, %1; st.shared.v2.u32 [smem_ptr], {%2, %3}; }\n"
+        : : "r"(smem_addr), "r"(offset_in_bytes), "r"(data[0]), "r"(data[1])
+      );
+#endif
+    }
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) {
+    store_with_pointer_offset(frag, 0);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for float x 8 => float_e4m3_t/float_e5m2_t x 8
+template <
+  typename WarpShape_,            ///< shape of warp-level GEMM (concept: GemmShape)
+  typename OperatorShape_         ///< matrix multiply operation shape (concept: gemm::GemmShape)
+>
+class TileIteratorTensorOpMixed<WarpShape_, OperatorShape_, float, 32, 8, 8, 8> {
+public:
+
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using Element = float;
+  using Layout = layout::RowMajor;
+  static int const kOutputElementCount = 8;
+
+  using TensorRef = TensorRef<Element, Layout>;         ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;                      ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  /// Shape of the tile in memory
+  using Shape = MatrixShape<
+    Policy::kRowsPerIteration,
+    WarpShape::kN
+  >;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    Element,
+    Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
+
+  /// This is the complete warp-level accumulator tile.
+  //using AccumulatorTile = typename Operator::FragmentC;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+  // Internal constants
+  struct Detail {
+    static int const kLanesInQuad = 4;
+
+    /// Number of pointers needed to write accumulators
+    static int const kPointerCount = 2;
+
+    static_assert(sizeof(Element) == 4, "This can only be used with 32b accumulator data types (f32, s32).");
+  };
+
+  /// Padding quantity
+  using Padding = MatrixShape<0, Detail::kLanesInQuad * 2>;
+
+private:
+
+  /// Storage type for accessing memory
+  using AccessType = AlignedArray<Element, 2>;
+
+  //
+  // Data members
+  //
+
+  /// Internal pointer to memory
+  AccessType *pointers_[Detail::kPointerCount] = {nullptr};
+
+  /// Stride in units of AccessType
+  int stride_{0};
+
+public:
+
+  /// Default constructor
+  TileIteratorTensorOpMixed() = default;
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed(
+    TensorRef const &ref,
+    unsigned lane_id
+  ):
+    stride_(ref.stride()[0] / AccessType::kElements) {
+
+    int quad_id = (lane_id / Detail::kLanesInQuad);
+    int lane_in_quad = (lane_id % Detail::kLanesInQuad);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Detail::kPointerCount; ++i) {
+      AccessType *ptr = reinterpret_cast<AccessType *>(ref.data()) + quad_id * stride_;
+      int column_idx = lane_in_quad ^ (i * 2);
+
+      ptr += column_idx;
+
+      if (i == 0) {
+        pointers_[0] = ptr;
+      }
+      else if (i == 1) {
+        pointers_[1] = ptr;
+      }
+    }
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & add_pointer_offset(Index pointer_offset) {
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
+      pointers_[i] += pointer_offset / AccessType::kElements;
+    }
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & add_tile_offset(TensorCoord const &tile_offset) {
+
+    int ptr_offset = tile_offset.row() * Shape::kRow * stride_ +
+      tile_offset.column() * Shape::kColumn / AccessType::kElements;
+
+    pointers_[0] += ptr_offset;
+    pointers_[1] += ptr_offset;
+
+    if (tile_offset.column() % 2) {
+      auto tmp = pointers_[0];
+      pointers_[0] = pointers_[1];
+      pointers_[1] = tmp;
+    }
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorTensorOpMixed & operator+=(TensorCoord const &tile_offset) {
+    return add_tile_offset(tile_offset);
+  }
+
+  /// Store
+  CUTLASS_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
+
+      int ptr_idx = (n / 4);
+
+      AccessType *ptr;
+      if (ptr_idx == 0) {
+        ptr = pointers_[0];
+      }
+      else if (ptr_idx == 1) {
+        ptr = pointers_[1];
+      }
+
+      int offset = (n / 4) * 16 + pointer_offset / AccessType::kElements + (n % 4) * 4;
+
+#if 0
+      //
+      // Using inline PTX to avoid generic memory
+      //
+      AccessType *smem_ptr = pointers_[ptr_idx];
+      smem_ptr[offset] = frag_ptr[n];
+#else
+      uint32_t smem_addr = arch::cutlass_get_smem_pointer(ptr);
+      uint32_t const *data = reinterpret_cast<uint32_t const *>(frag_ptr + n);
+      uint32_t offset_in_bytes = offset * sizeof(AccessType);
+
+      asm volatile(
+        "{ .reg .u32 smem_ptr; add.u32 smem_ptr, %0, %1; st.shared.v2.u32 [smem_ptr], {%2, %3}; }\n"
+        : : "r"(smem_addr), "r"(offset_in_bytes), "r"(data[0]), "r"(data[1])
+      );
+#endif
+    }
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) {
+    store_with_pointer_offset(frag, 0);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+#undef CUTLASS_EPILOGUE_WARP_TILE_ITERATOR_TENSOR_OP_MIXED_OPTIMIZATION_ENABLED
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/tile_iterator_volta_tensor_op.h

@@ -0,0 +1,440 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief 
+*/
+
+#pragma once
+
+#include "cutlass/array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/layout/pitch_linear.h"
+
+#include "cutlass/epilogue/warp/tensor_op_policy.h"
+#include "cutlass/epilogue/warp/volta_tensor_op_policy.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template for reading and writing tiles of accumulators to shared memory
+template <
+  typename WarpShape,             ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename InterleavedTileShape,  ///< shape of indivisible instruction-level arrangement (concept: GemmShape)
+  typename ElementC,              ///< Accumulator layout
+  typename Layout                 ///< target shared memory layout
+>
+struct TileIteratorVoltaTensorOp; 
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template for reading and writing tiles of accumulators to shared memory
+template <
+  typename WarpShape_         ///< shape of warp-level GEMM (concept: MatrixShape)
+>
+struct TileIteratorVoltaTensorOp<WarpShape_, gemm::GemmShape<32, 32, 4>, half_t, layout::RowMajor> {
+public:
+
+  using WarpShape = WarpShape_;
+  using InterleavedTileShape = gemm::GemmShape<32, 32, 4>;
+  using Element = half_t;
+  using Layout = layout::RowMajor;
+
+  using TensorRef = TensorRef<Element, Layout>;         ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;                      ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  using Policy = VoltaTensorOpPolicy<WarpShape, InterleavedTileShape, Element, Layout>;
+
+  /// Shape of the tile in memory
+  using Shape = MatrixShape<
+    Policy::kRowsPerIteration,
+    WarpShape::kN
+  >;
+
+  /// Array type for aligned memory accesses
+  using AccessType = typename Policy::AccessType;
+  
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = typename Policy::Fragment;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = typename Policy::AccumulatorTile;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+  /// Number of elements per access
+  static int const kElementsPerAccess = Policy::kElementsPerAccess;
+
+  // Internal constants
+  struct Detail {
+    static int const kLanesInQuad = 4;
+    static int const kRowsPerQuad = 4;
+    static int const kColumnsPerQuad = 8;
+    static int const kAccessesPerQuad = kColumnsPerQuad / Policy::kElementsPerAccess;
+    static int const kAccessQuadDelta = 16;
+  };
+
+  /// Padding quantity
+  using Padding = MatrixShape<
+    0,
+    Policy::kElementsPerAccess>;
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Internal pointer to memory
+  AccessType *pointer_;
+
+  /// Internal layout object
+  Layout layout_;
+
+public:
+
+  /// Default constructor
+  CUTLASS_HOST_DEVICE
+  TileIteratorVoltaTensorOp(): pointer_(nullptr) { }
+
+  /// Constructor from TensorRef
+  CUTLASS_DEVICE
+  TileIteratorVoltaTensorOp(
+    TensorRef const &ref,
+    unsigned lane_id
+  ):
+    pointer_(reinterpret_cast<AccessType *>(ref.data())),
+    layout_(ref.stride()[0] / Policy::kElementsPerAccess) { 
+
+    int quad_id = lane_id / Detail::kLanesInQuad;
+    int lane_in_quad = (lane_id % Detail::kLanesInQuad);
+
+    int quad_row_idx = ((quad_id & 4) >> 1) + (quad_id & 1);
+    int quad_col_idx = ((quad_id & 2) >> 1);
+
+    int row = quad_row_idx * Detail::kRowsPerQuad + lane_in_quad;
+    int column = quad_col_idx * Detail::kColumnsPerQuad;
+
+    pointer_ += layout_({row, column / kElementsPerAccess});
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorVoltaTensorOp & add_pointer_offset(Index pointer_offset) {
+    pointer_ += pointer_offset / Policy::kElementsPerAccess;
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorVoltaTensorOp & add_tile_offset(TensorCoord const &tile_offset) {
+
+    pointer_ += layout_({
+      tile_offset.row() * Shape::kRow, 
+      tile_offset.column() * Shape::kColumn / Policy::kElementsPerAccess});
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorVoltaTensorOp & operator+=(TensorCoord const &tile_offset) {
+    add_tile_offset(tile_offset);
+    return *this;
+  }
+
+  /// Store
+  CUTLASS_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int tile_idx = 0; tile_idx < Policy::TileIterations::kColumn; ++tile_idx) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int access_idx = 0; access_idx < Policy::kAccessesPerInterleavedTile; ++access_idx) {
+
+        int access_quad = access_idx / 2;
+        int access = access_idx % 2;
+
+        int ptr_offset = tile_idx * InterleavedTileShape::kN / Policy::kElementsPerAccess +
+          access_quad * Detail::kAccessQuadDelta / Policy::kElementsPerAccess + 
+          access + pointer_offset / Policy::kElementsPerAccess;
+
+        int frag_idx = tile_idx * Policy::kAccessesPerInterleavedTile + access_idx;
+
+        AccessType access_vector = frag_ptr[frag_idx];
+
+        pointer_[ptr_offset] = access_vector;
+      }
+    }
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) {
+    store_with_pointer_offset(frag, 0);
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int tile_idx = 0; tile_idx < Policy::TileIterations::kColumn; ++tile_idx) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int access_idx = 0; access_idx < Policy::kAccessesPerInterleavedTile; ++access_idx) {
+
+        int access_quad = access_idx / 2;
+        int access = access_idx % 2;
+
+        int ptr_offset = tile_idx * Detail::kTileDelta + access_quad * Detail::kAccessQuadDelta + 
+          access + pointer_offset / Policy::kElementsPerAccess;
+
+        int frag_idx = tile_idx * Policy::kAccessesPerInterleavedTile + access_idx;
+
+        frag_ptr[frag_idx] = pointer_[ptr_offset];
+      }
+    }
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load(Fragment const &frag) {
+    load_with_pointer_offset(frag, 0);
+  }
+  
+  /// Set smem base address
+  CUTLASS_HOST_DEVICE
+  void set_smem_base_address(Index address) {
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template for reading and writing tiles of accumulators to shared memory
+template <
+  typename WarpShape_         ///< shape of warp-level GEMM (concept: MatrixShape)
+>
+struct TileIteratorVoltaTensorOp<WarpShape_, gemm::GemmShape<32, 32, 4>, float, layout::RowMajor> {
+public:
+
+  using WarpShape = WarpShape_;
+  using InterleavedTileShape = gemm::GemmShape<32, 32, 4>;
+  using Element = float;
+  using Layout = layout::RowMajor;
+
+  using TensorRef = TensorRef<Element, Layout>;         ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;                      ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  using Policy = VoltaTensorOpPolicy<WarpShape, InterleavedTileShape, Element, Layout>;
+
+  /// Shape of the tile in memory
+  using Shape = MatrixShape<
+    Policy::kRowsPerIteration,
+    WarpShape::kN
+  >;
+
+  /// Array type for aligned memory accesses
+  using AccessType = typename Policy::AccessType;
+  
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = typename Policy::Fragment;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = typename Policy::AccumulatorTile;
+
+  /// Number of times this iterator can be incremented
+  static int const kIterations = Policy::kIterations;
+
+  /// Number of elements per access
+  static int const kElementsPerAccess = Policy::kElementsPerAccess;
+
+  // Internal constants
+  struct Detail {
+    static int const kLanesInQuad = 4;
+    static int const kRowsPerQuad = 4;
+    static int const kColumnsPerQuad = 8;
+    static int const kAccessesPerQuad = kColumnsPerQuad / Policy::kElementsPerAccess;
+    static int const kAccessQuadDelta = 16;
+  };
+
+  /// Padding quantity
+  using Padding = MatrixShape<
+    0,
+    Policy::kElementsPerAccess>;
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Internal pointer to memory
+  AccessType *pointer_;
+
+  /// Internal layout object
+  Layout layout_;
+
+public:
+
+  /// Default constructor
+  CUTLASS_HOST_DEVICE
+  TileIteratorVoltaTensorOp(): pointer_(nullptr) { }
+
+  /// Constructor from TensorRef
+  CUTLASS_DEVICE
+  TileIteratorVoltaTensorOp(
+    TensorRef const &ref,
+    unsigned lane_id
+  ):
+    pointer_(reinterpret_cast<AccessType *>(ref.data())),
+    layout_(ref.stride()[0] / Policy::kElementsPerAccess) { 
+
+    int quad_id = lane_id / Detail::kLanesInQuad;
+    int lane_in_quad = (lane_id % Detail::kLanesInQuad);
+
+    int const kQuadRowDelta = 4;
+    int const kQuadColumnDelta = 2 * Policy::MmaIterations::kColumn;
+
+    int quad_row_offset = ((quad_id & 4) / 2 + (quad_id & 1)) * kQuadRowDelta;
+    int quad_column_offset = (quad_id & 2) / 2 * kQuadColumnDelta;
+
+    int thread_row_offset = (lane_in_quad & 1);
+    int thread_column_offset = (lane_in_quad & 2) / 2;
+
+    int row = quad_row_offset + thread_row_offset;
+    int column = quad_column_offset + thread_column_offset;
+
+    pointer_ += layout_({row, column});
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorVoltaTensorOp & add_pointer_offset(Index pointer_offset) {
+    pointer_ += pointer_offset / Policy::kElementsPerAccess;
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorVoltaTensorOp & add_tile_offset(TensorCoord const &tile_offset) {
+
+    pointer_ += layout_({
+      tile_offset.row() * Shape::kRow, 
+      tile_offset.column() * Shape::kColumn / Policy::kElementsPerAccess});
+
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorVoltaTensorOp & operator+=(TensorCoord const &tile_offset) {
+    add_tile_offset(tile_offset);
+    return *this;
+  }
+
+  /// Store
+  CUTLASS_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+
+    AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
+
+    int const kAccessesPerRow = Policy::TileIterations::kColumn * Policy::MmaIterations::kColumn * 2;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int row_idx = 0; row_idx < Policy::kRowsPerMmaTile; ++row_idx) {
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int access_idx = 0; access_idx < kAccessesPerRow; ++access_idx) {
+
+        int frag_idx = row_idx * kAccessesPerRow + access_idx;
+
+        int ptr_column_offset = (access_idx & 1) * 2 + 
+          (access_idx & 2) * Policy::MmaIterations::kColumn * 2 + 
+          (access_idx & 4) * Policy::MmaIterations::kColumn * 2;
+
+        int ptr_row_offset = row_idx * 2;
+
+        int ptr_offset = layout_({ptr_row_offset, ptr_column_offset}) + pointer_offset / Policy::kElementsPerAccess;
+
+        pointer_[ptr_offset] = frag_ptr[frag_idx];
+      }
+    }
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) {
+    store_with_pointer_offset(frag, 0);
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+
+    AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
+
+    assert(0);
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load(Fragment const &frag) {
+    load_with_pointer_offset(frag, 0);
+  }
+  
+  /// Set smem base address
+  CUTLASS_HOST_DEVICE
+  void set_smem_base_address(Index address) {
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/tile_iterator_wmma_tensor_op.h

@@ -0,0 +1,224 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief 
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/wmma_array.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/layout/pitch_linear.h"
+#include "cutlass/tensor_ref.h"
+
+#include "cutlass/epilogue/warp/wmma_tensor_op_policy.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template for reading and writing tiles of accumulators to shared memory
+template <
+  typename WarpShape,           ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename OperatorShape,       ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  typename OperatorFragment,    ///< wmma fragment to be written (concept: nvcuda::wmma::fragment)
+  typename Layout               ///< target shared memory layout
+>
+class TileIteratorWmmaTensorOp;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template for reading and writing tiles of accumulators to shared memory
+template <
+  typename WarpShape_,          ///< shape of warp-level GEMM (concept: GemmShape)
+  typename OperatorShape_,      ///< matrix multiply operation shape (concept: gemm::GemmShape)
+  typename OperatorFragment_    ///< wmma fragment to be written (concept: nvcuda::wmma::fragment)
+>
+class TileIteratorWmmaTensorOp<WarpShape_, OperatorShape_, OperatorFragment_, layout::RowMajor> {
+public:
+
+  using WarpShape = WarpShape_;
+  using OperatorShape = OperatorShape_;
+  using OperatorFragment = OperatorFragment_;
+  using Layout = layout::RowMajor;
+
+  //
+  // Derived types
+  //
+  using WmmaDataType = typename OperatorFragment::element_type;
+  using Element = typename cutlass::arch::WmmaToCutlassDataType<WmmaDataType>::Type; ///< Data Type of element stored in nvcuda::wmma::frament         
+  using TensorRef = TensorRef<Element, Layout>;                                      ///< Tensor Reference object
+  using TensorCoord = MatrixCoord;                                                   ///< Logical coordinate in referenced tensor
+  using Index = typename TensorRef::Index;
+  using LongIndex = typename TensorRef::LongIndex;
+
+  using Policy = WmmaTensorOpPolicy<WarpShape, OperatorShape, Layout>;
+
+  /// Shape of the tile in memory
+  using Shape = MatrixShape<
+    Policy::kRowsPerIteration,
+    WarpShape::kN
+  >;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = WmmaFragmentArray<OperatorFragment, Policy::OperatorCount::kColumn * Policy::kWmmaFragmentsPerAccess>;
+
+
+  /// This is the complete warp-level accumulator tile.
+  //using AccumulatorTile = typename Operator::FragmentC;
+
+
+  /// Padding quantity 
+  // (Epilogue shared memory padding for WMMA Gemm kernel is set to run optimaly on Turing)
+  using Padding = MatrixShape<
+    0,
+    4 * Policy::kElementsPerAccess
+  >;
+
+private:
+
+  /// Storage type for accessing memory
+  //using AccessType = AlignedArray<Element, Policy::kElementsPerAccess>;
+
+  //
+  // Data members
+  //
+
+  /// Internal pointer to shared memory
+  TensorRef ref_;
+
+
+public:
+
+  /// Default constructor
+  CUTLASS_HOST_DEVICE
+  TileIteratorWmmaTensorOp(): ref_(nullptr) { 
+
+  }
+
+  /// Constructor from TensorRef
+  CUTLASS_HOST_DEVICE
+  TileIteratorWmmaTensorOp(
+    TensorRef const &ref,
+    unsigned lane_id
+  ): ref_(ref) {
+  }
+
+  /// Adds a pointer offset
+  CUTLASS_HOST_DEVICE
+  TileIteratorWmmaTensorOp & add_pointer_offset(Index pointer_offset) {
+    ref_.add_pointer_offset(pointer_offset);
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorWmmaTensorOp & add_tile_offset(TensorCoord const &tile_offset) {
+    ref_.add_coord_offset({tile_offset.row() * OperatorShape::kM, tile_offset.column() * WarpShape::kN});
+    return *this;
+  }
+
+  ///< advances in units of whole tiles along the logical coordinate space of the tensor
+  CUTLASS_HOST_DEVICE
+  TileIteratorWmmaTensorOp & operator+=(TensorCoord const &tile_offset) {
+    add_tile_offset(tile_offset);
+    return *this;
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
+
+    for(int n=0; n < Policy::OperatorCount::kColumn; n++) {
+      
+      WmmaDataType* ptr = reinterpret_cast<WmmaDataType*> (ref_.data() + ref_.offset({0, n * OperatorShape::kN}) + pointer_offset);
+
+      nvcuda::wmma::store_matrix_sync(
+        ptr, 
+        frag[n], 
+        ref_.stride()[0], 
+        nvcuda::wmma::layout_t::mem_row_major
+      ); 
+    
+    }
+  }
+
+  /// Store
+  CUTLASS_HOST_DEVICE
+  void store(Fragment const &frag) {
+    store_with_pointer_offset(frag, 0);
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const {
+ 
+    for(int n=0; n < Policy::OperatorCount::kColumn; n++) {
+
+      WmmaDataType* ptr = reinterpret_cast<WmmaDataType*> (ref_.data() + ref_.offset({0, n * OperatorShape::kN}) + pointer_offset);
+
+      nvcuda::wmma::load_matrix_sync(         
+        frag[n], 
+        ptr,
+        ref_.stride()[0], 
+        nvcuda::wmma::layout_t::mem_row_major
+      ); 
+    
+    }
+  }
+
+  /// Load
+  CUTLASS_HOST_DEVICE
+  void load(Fragment &frag) const {
+    load_with_pointer_offset(frag, 0);
+  }
+
+  
+  /// Set smem base address
+  CUTLASS_HOST_DEVICE
+  void set_smem_base_address(Index address) {
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/volta_tensor_op_policy.h

@@ -0,0 +1,195 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Defines basic structures needed for implementing the warp-scoped phase of the epilogue.
+          These quantities assume a 'column-major' arrangement of TensorOp instructions, of which
+          a row-oriented slice is visible per iteration.
+*/
+
+#pragma once
+
+#include "cutlass/matrix_shape.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/gemm/gemm.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Policy details related to the epilogue
+template <
+  typename WarpShape,             ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename InterleavedTileShape,  ///< shape of indivisible instruction-level arrangement (concept: GemmShape)
+  typename ElementC,              ///< Accumulator layout
+  typename Layout                 ///< target shared memory layout
+>
+struct VoltaTensorOpPolicy; 
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for row-major
+template <
+  typename WarpShape_          ///< shape of warp-level GEMM (concept: GemmShape)
+>
+struct VoltaTensorOpPolicy<WarpShape_, gemm::GemmShape<32, 32, 4>, half_t, layout::RowMajor> {
+
+  using WarpShape = WarpShape_;
+  using InterleavedTileShape = gemm::GemmShape<32, 32, 4>;
+  using ElementC = half_t;
+  using Layout = layout::RowMajor;
+
+  /// Shape of one warp-levelinstruction
+  using InstructionShape = gemm::GemmShape<16, 16, 4>;
+
+  /// Number of mma operations performed for one 32x32x4 interleaved tile
+  using MmaIterations = MatrixShape<
+    InterleavedTileShape::kM / InstructionShape::kM,
+    InterleavedTileShape::kN / InstructionShape::kN
+  >;
+
+  /// Number of 32x32x4 interleaved tiles performed to cover the warp-level GEMM shape
+  using TileIterations = MatrixShape<
+    WarpShape::kM / InterleavedTileShape::kM,
+    WarpShape::kN / InterleavedTileShape::kN
+  >;
+
+  /// Number of accumulator elements owned by each thread per Mma
+  static int const kElementsPerMma = 8;
+  static int const kRowsPerIteration = 16;
+
+  //
+  // Hard-coded constants regarding Tensor Operations
+  //
+
+  /// Number of accumulator elements stored per memory instruction to shared memory
+  static int const kElementsPerAccess = 4;
+  
+  /// Number of accesses performed per interleaved tile
+  static int const kAccessesPerInterleavedTile = 4;
+
+  /// Total number of iterations needed to cover the entire tile
+  static int const kIterations = TileIterations::kRow * 2;
+
+  //
+  // Derived types
+  //
+
+  /// Array type for aligned memory accesses
+  using AccessType = AlignedArray<ElementC, kElementsPerAccess>;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    ElementC, 
+    kElementsPerAccess * kAccessesPerInterleavedTile * TileIterations::kColumn>;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = Array<
+    ElementC, 
+    TileIterations::kCount * MmaIterations::kCount * kElementsPerMma>;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for row-major
+template <
+  typename WarpShape_          ///< shape of warp-level GEMM (concept: MatrixShape)
+>
+struct VoltaTensorOpPolicy<WarpShape_, gemm::GemmShape<32, 32, 4>, float, layout::RowMajor> {
+
+  using WarpShape = WarpShape_;
+  using InterleavedTileShape = gemm::GemmShape<32, 32, 4>;
+  using ElementC = float;
+  using Layout = layout::RowMajor;
+
+  /// Shape of one warp-levelinstruction
+  using InstructionShape = gemm::GemmShape<16, 16, 4>;
+
+  /// Number of mma operations performed for one 32x32x4 interleaved tile
+  using MmaIterations = MatrixShape<
+    InterleavedTileShape::kM / InstructionShape::kM,
+    InterleavedTileShape::kN / InstructionShape::kN
+  >;
+
+  /// Number of 32x32x4 interleaved tiles performed to cover the warp-level GEMM shape
+  using TileIterations = MatrixShape<
+    WarpShape::kM / InterleavedTileShape::kM,
+    WarpShape::kN / InterleavedTileShape::kN
+  >;
+
+  /// Number of accumulator elements owned by each thread per Mma
+  static int const kElementsPerMma = 8;
+  static int const kRowsPerIteration = 16;
+
+  //
+  // Hard-coded constants regarding Tensor Operations
+  //
+
+  /// Number of accumulator elements stored per memory instruction to shared memory
+  static int const kElementsPerAccess = 2;
+  
+  /// Number of accesses performed per interleaved tile
+  static int const kAccessesPerInterleavedTile = 8;
+
+  /// Number of rows per interleaved tile
+  static int const kRowsPerMmaTile = 2;
+
+  /// Total number of iterations needed to cover the entire tile
+  static int const kIterations = TileIterations::kRow * MmaIterations::kRow;
+
+  //
+  // Derived types
+  //
+  
+  /// Array type for aligned memory accesses
+  using AccessType = AlignedArray<ElementC, kElementsPerAccess>;
+
+  /// This is the fragment size produced by one access of the iterator.
+  using Fragment = Array<
+    ElementC, 
+    kElementsPerAccess * kAccessesPerInterleavedTile * TileIterations::kColumn>;
+
+  /// This is the complete warp-level accumulator tile.
+  using AccumulatorTile = Array<
+    ElementC, 
+    TileIterations::kCount * MmaIterations::kCount * kElementsPerMma>;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/epilogue/warp/wmma_tensor_op_policy.h

@@ -0,0 +1,101 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Defines basic structures needed for implementing the warp-scoped phase of the epilogue.
+          These quantities assume a 'column-major' arrangement of TensorOp instructions, of which
+          a row-oriented slice is visible per iteration.
+*/
+
+#pragma once
+
+#include "cutlass/arch/wmma.h"
+#include "cutlass/matrix_shape.h"
+#include "cutlass/layout/matrix.h"
+
+#if defined(CUTLASS_ARCH_WMMA_ENABLED)
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace epilogue {
+namespace warp {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Policy details related to the epilogue
+template <
+  typename WarpShape,     ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename OperatorShape, ///< matrix multiply operation shape (concept: gemm:GemmShape)
+  typename Layout         ///< target shared memory layout
+>
+struct WmmaTensorOpPolicy; 
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for row-major
+template <
+  typename WarpShape,           ///< shape of warp-level GEMM (concept: MatrixShape)
+  typename OperatorShape        ///< matrix multiply operation shape (concept: gemm::GemmShape)
+>
+struct WmmaTensorOpPolicy<WarpShape, OperatorShape, layout::RowMajor> {
+
+  /// Number of operations
+  using OperatorCount = MatrixShape<
+    WarpShape::kM / OperatorShape::kM,
+    WarpShape::kN / OperatorShape::kN
+  >;
+
+  //
+  // Hard-coded constants regarding Tensor Operations
+  //
+  static int const kElementsPerAccess = 2;
+  static int const kRowsPerIteration = OperatorShape::kM;
+  static int const kWmmaFragmentsPerAccess = 1;
+
+  //
+  // Derived quantities
+  //
+
+  // Number of externally visible iterations
+  static int const kIterations = OperatorCount::kRow;
+
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace epilogue
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////
+
+#endif
+

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/exmy_base.h

@@ -0,0 +1,1222 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+
+
+/*!
+  \file
+  \brief Generic floating-point type for ExMy format
+*/
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_size.h"
+#include "cutlass/platform/platform.h"
+
+// #define CUTLASS_DEBUG_TRACE_LEVEL 2
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+ // Helper functions
+namespace detail {
+
+template <class Src, class Dst>
+CUTLASS_HOST_DEVICE
+Dst copy_bits(Src src)
+{
+  Dst dst;
+  static_assert(sizeof(Src) <= sizeof(Dst), "Dst type should be at least the same size as Src type");
+  static_assert(cutlass::platform::is_trivially_copyable<Dst>::value, "Dst type should be trivially copyable");
+  static_assert(cutlass::platform::is_trivially_copyable<
+    /*cutlass::platform::remove_cvref_t< */ Dst /* > */ >::value, "Dst type should be trivially copyable");
+  memcpy(&dst, &src, sizeof(src));
+  return dst;
+}
+
+enum class NanInfEncoding
+{
+  // IEEE-754 style NaN. Exponent bits are
+  // all ones, and at least one bit of mantissa is one
+  IEEE_754,
+  // Canonical NaN. There is only one value representing NaN and
+  // no Inf is defined.
+  CANONICAL_ONLY,
+  // No NaN or Inf encoded.
+  NONE
+};
+
+enum class FpEncoding
+{
+  E11M52, // double
+  E8M23,  // float
+  E5M2,   // FP8
+  E4M3,   // FP8
+  UE4M3,  // FP8 
+  UE8M0,  // FP8
+  E3M2,   // FP6
+  E2M3,   // FP6
+  E2M1,   // FP4
+};
+
+//////
+
+#if (CUTLASS_CXX17_OR_LATER)
+template<uint32_t NumExpBits, uint32_t NumMantissaBits>
+constexpr int exponent_bias_cxx17() {
+  if CUTLASS_CONSTEXPR_IF_CXX17 (NumExpBits == 0) {
+    static_assert(NumMantissaBits <= static_cast<uint32_t>(cutlass::platform::numeric_limits<int32_t>::max()));
+    return -1 * static_cast<int>(NumMantissaBits);
+  }
+  else {
+    return static_cast<int>((1 << (NumExpBits - 1))) - 1;
+  }
+
+  CUTLASS_GCC_UNREACHABLE;
+}
+#endif
+
+namespace impl {
+template<uint32_t NumExpBitsMinusOne>
+constexpr int shift_num_bits_expression_cxx11() {
+#if (CUTLASS_CXX17_OR_LATER)
+  static_assert(NumExpBitsMinusOne <= 31u);
+#endif
+  return NumExpBitsMinusOne > 31u ? 31u : NumExpBitsMinusOne;
+}
+
+template<uint32_t NumExpBitsMinusOne>
+constexpr int inner_shift_expression_cxx11() {
+  return static_cast<int>((1u << shift_num_bits_expression_cxx11<NumExpBitsMinusOne>()) - 1u);
+}
+
+} // namespace impl
+
+// C++11 equivalent of exponent_bias_cxx17()
+template<uint32_t NumExpBits, uint32_t NumMantissaBits>
+constexpr int exponent_bias_cxx11() {
+#if (CUTLASS_CXX17_OR_LATER)
+  return exponent_bias_cxx17<NumExpBits, NumMantissaBits>();
+#else
+  return (NumExpBits == 0) ?
+    -1 * static_cast<int>(NumMantissaBits) : impl::inner_shift_expression_cxx11<NumExpBits - 1u>();
+#endif
+}
+
+// C++11 equivalent of maximum_exponent_cxx17()
+template<uint32_t NumExpBits, uint32_t NumMantissaBits, NanInfEncoding NaNEncoding>
+constexpr int maximum_exponent_cxx11() {
+  return
+    ((NumExpBits == 0) ?
+      (0 - exponent_bias_cxx11<NumExpBits, NumMantissaBits>()) :
+      ((NaNEncoding == NanInfEncoding::IEEE_754) ?
+        ((static_cast<int>((1 << NumExpBits)) - 2) - exponent_bias_cxx11<NumExpBits, NumMantissaBits>()) :
+        ((NaNEncoding == NanInfEncoding::CANONICAL_ONLY) ?
+          ((NumMantissaBits > 0) ?
+            static_cast<int>((1 << NumExpBits)) - 1 - exponent_bias_cxx11<NumExpBits, NumMantissaBits>() :
+            static_cast<int>((1 << NumExpBits)) - 2 - exponent_bias_cxx11<NumExpBits, NumMantissaBits>()
+          ) :
+          (static_cast<int>((1 << NumExpBits)) - 1 - exponent_bias_cxx11<NumExpBits, NumMantissaBits>())
+        )
+      )
+    );
+}
+
+#if (CUTLASS_CXX17_OR_LATER)
+template<uint32_t NumExpBits, uint32_t NumMantissaBits, NanInfEncoding NaNEncoding>
+constexpr int maximum_exponent_cxx17() {
+  constexpr int exp_bias = exponent_bias_cxx17<NumExpBits, NumMantissaBits>();
+  if CUTLASS_CONSTEXPR_IF_CXX17 (NumExpBits == 0) {
+    // If no exponent bits, return fixed hidden bias
+    return 0 - exp_bias;
+  }
+  else {
+    if CUTLASS_CONSTEXPR_IF_CXX17 (NaNEncoding == NanInfEncoding::IEEE_754) {
+      // We have IEEE style NaN and infinity
+      // All values when exp_bits = 1...1s are used.
+      int max_exp_bits = static_cast<int>((1 << NumExpBits)) - 2;
+      return max_exp_bits - exp_bias;
+    }
+    else {
+      // There are no cases where we have Inf without IEEE_754_Nan
+
+      // If we have a canonical NaN. Only exp=1..1 and mantissa=1..1
+      // value has a special meaning. If we also have at least one mantissa
+      // bit, then maximum exponent is 1...1 - exponent_bias
+      if CUTLASS_CONSTEXPR_IF_CXX17 (NaNEncoding == NanInfEncoding::CANONICAL_ONLY) {
+        if CUTLASS_CONSTEXPR_IF_CXX17 (NumMantissaBits > 0) {
+          int max_exp_bits = static_cast<int>((1 << NumExpBits)) - 1;
+          return max_exp_bits - exp_bias;
+        }
+        else { // no mantissa bits
+          int max_exp_bits = static_cast<int>((1 << NumExpBits)) - 2;
+          return max_exp_bits - exp_bias;
+        }
+      }
+      // No NaNs or infs
+      int max_exp_bits = static_cast<int>((1 << NumExpBits)) - 1;
+      return max_exp_bits - exp_bias;
+    }
+  }
+
+  CUTLASS_GCC_UNREACHABLE;
+}
+#endif
+
+template<uint32_t NumExpBits, uint32_t NumMantissaBits>
+constexpr int minimum_exponent_cxx11() {
+  return
+    ((NumExpBits == 0) ?
+      0 - exponent_bias_cxx11<NumExpBits, NumMantissaBits>() :
+      ((NumMantissaBits > 0) ?
+        1 - exponent_bias_cxx11<NumExpBits, NumMantissaBits>() :
+        0 - exponent_bias_cxx11<NumExpBits, NumMantissaBits>())
+    );
+}
+
+#if (CUTLASS_CXX17_OR_LATER)
+template<uint32_t NumExpBits, uint32_t NumMantissaBits>
+constexpr int minimum_exponent_cxx17() {
+  constexpr int exp_bias = exponent_bias_cxx17<NumExpBits, NumMantissaBits>();
+  constexpr bool has_denorm = (NumMantissaBits > 0);
+  if CUTLASS_CONSTEXPR_IF_CXX17 (NumExpBits == 0) {
+    // If no exponent bits, return fixed hidden bias
+    // Note that minimum and maximum exponents are the same.
+    return 0 - exp_bias;
+  }
+
+  if CUTLASS_CONSTEXPR_IF_CXX17 (has_denorm) {
+    // Exp = 0...0s is reserved for denorm values.
+    return 1 - exp_bias;
+  }
+  return 0 - exp_bias;
+}
+#endif
+
+template<class Storage, uint32_t NumExpBits, uint32_t NumMantissaBits, NanInfEncoding NaNEncoding>
+constexpr Storage max_pos_denormal_value_cxx11() {
+  static_assert(NumExpBits > 0 || NumMantissaBits > 0, "Both NumExpBits and NumMantissaBits can't be zero");
+  return
+    (!(NumMantissaBits > 0) ? Storage(0) : Storage((1ull << NumMantissaBits) - 1));
+}
+
+#if (CUTLASS_CXX17_OR_LATER)
+template<class Storage, uint32_t NumExpBits, uint32_t NumMantissaBits, NanInfEncoding NaNEncoding>
+constexpr Storage max_pos_denormal_value_cxx17() {
+  static_assert(NumExpBits > 0 || NumMantissaBits > 0, "Both NumExpBits and NumMantissaBits can't be zero");
+  constexpr bool has_denorm = (NumMantissaBits > 0);
+  if CUTLASS_CONSTEXPR_IF_CXX17 (!has_denorm) {
+    // If we don't have denormal values, return all 0s
+    return Storage(0);
+  }
+  else {
+    // Case: (NumExpBits > 0 && NumMantissaBits > 0) or (NumExpBits == 0 && NumMantissaBits > 0)
+    return Storage((1ull << NumMantissaBits) - 1);
+  }
+
+  CUTLASS_GCC_UNREACHABLE;
+}
+#endif
+
+
+template<class Storage, uint32_t NumExpBits, uint32_t NumMantissaBits, NanInfEncoding NaNEncoding>
+constexpr Storage min_pos_denormal_value_cxx11() {
+  return (!(NumMantissaBits > 0) ? Storage(0) : Storage(1));
+}
+
+#if (CUTLASS_CXX17_OR_LATER)
+template<class Storage, uint32_t NumExpBits, uint32_t NumMantissaBits, NanInfEncoding NaNEncoding>
+constexpr Storage min_pos_denormal_value_cxx17() {
+  constexpr bool has_denorm = (NumMantissaBits > 0);
+  if CUTLASS_CONSTEXPR_IF_CXX17 (!has_denorm) {
+    // If we don't have denormal values, return all 0s
+    return Storage(0);
+  }
+  // Case: (NumExpBits > 0 && NumMantissaBits > 0) or (NumExpBits == 0 && NumMantissaBits > 0)
+  return Storage(1);
+}
+#endif
+
+template<class Storage, uint32_t NumExpBits, uint32_t NumMantissaBits, NanInfEncoding NaNEncoding>
+constexpr Storage max_pos_normal_value_cxx11() {
+  return
+    ((NumExpBits == 0) ?
+      Storage(0) :
+      ((NumMantissaBits == 0) ?
+        0 :
+        (((NaNEncoding == NanInfEncoding::IEEE_754 || NaNEncoding == NanInfEncoding::NONE) ?
+          ((1ull << NumMantissaBits) - 1) :
+          ((1ull << NumMantissaBits) - 2)))
+      ) | (static_cast<Storage>(
+            maximum_exponent_cxx11<NumExpBits, NumMantissaBits, NaNEncoding>() +
+            exponent_bias_cxx11<NumExpBits, NumMantissaBits>()
+          ) << NumMantissaBits)
+    );
+}
+
+#if (CUTLASS_CXX17_OR_LATER)
+template<class Storage, uint32_t NumExpBits, uint32_t NumMantissaBits, NanInfEncoding NaNEncoding>
+constexpr Storage max_pos_normal_value_cxx17() {
+  if CUTLASS_CONSTEXPR_IF_CXX17 (NumExpBits == 0) {
+    // if there are no exponent bits, we don't have normal values.
+    return Storage(0);
+  }
+  constexpr int exp_bias = exponent_bias_cxx17<NumExpBits, NumMantissaBits>();
+  constexpr int max_exp = maximum_exponent_cxx17<NumExpBits, NumMantissaBits, NaNEncoding>();
+  constexpr int exp = max_exp + exp_bias;
+
+  // place the exponent
+  Storage val = static_cast<Storage>(exp) << NumMantissaBits;
+  // If there are no mantissa bits return the exponent
+  if CUTLASS_CONSTEXPR_IF_CXX17 (NumMantissaBits == 0) {
+    return val;
+  }
+  else {
+    // If the NaN Inf encoding follows IEEE 754 or there is no (NaN and Inf) then mantissa can be all 1..1s
+    if CUTLASS_CONSTEXPR_IF_CXX17 (NaNEncoding == NanInfEncoding::IEEE_754 ||
+                  NaNEncoding == NanInfEncoding::NONE  ) {
+      Storage mantissa = (1ull << NumMantissaBits) - 1;
+      val |= mantissa;
+    }
+    else {
+      // If we have a canonical NaN, then the exponent can be the maximum bit value
+      // but mantissa=1..1s is reserved for NaN.
+      Storage mantissa = (1ull << NumMantissaBits) - 2;
+      val |= mantissa;
+    }
+    return val;
+  }
+
+  CUTLASS_GCC_UNREACHABLE;
+}
+#endif
+
+template<class Storage, uint32_t NumExpBits, uint32_t NumMantissaBits, NanInfEncoding NaNEncoding>
+constexpr Storage min_pos_normal_value_cxx11() {
+  return
+    ((NumExpBits == 0) ?
+      Storage(0) :
+      (Storage((NumMantissaBits > 0) ? 1 : 0) << NumMantissaBits)
+    );
+}
+
+#if (CUTLASS_CXX17_OR_LATER)
+template<class Storage, uint32_t NumExpBits, uint32_t NumMantissaBits, NanInfEncoding NaNEncoding>
+constexpr Storage min_pos_normal_value_cxx17() {
+  constexpr bool has_denorm = (NumMantissaBits > 0);
+
+  if CUTLASS_CONSTEXPR_IF_CXX17 (NumExpBits == 0) {
+    // if there are no exponent bits, we don't have normal values.
+    return Storage(0);
+  }
+  Storage exp = 0;
+  if CUTLASS_CONSTEXPR_IF_CXX17 (has_denorm) {
+    exp = 1;
+  }
+  return static_cast<Storage>(exp << NumMantissaBits);
+}
+#endif
+
+template<class Storage, uint32_t NumExpBits, uint32_t NumMantissaBits, NanInfEncoding NaNEncoding>
+constexpr Storage max_value_cxx11() {
+  return
+    ((NumExpBits > 0) ?
+      max_pos_normal_value_cxx11<Storage, NumExpBits, NumMantissaBits, NaNEncoding>() :
+      max_pos_denormal_value_cxx11<Storage, NumExpBits, NumMantissaBits, NaNEncoding>()
+    );
+}
+
+#if (CUTLASS_CXX17_OR_LATER)
+template<class Storage, uint32_t NumExpBits, uint32_t NumMantissaBits, NanInfEncoding NaNEncoding>
+constexpr Storage max_value_cxx17() {
+  constexpr bool has_normal = (NumExpBits > 0);
+  if CUTLASS_CONSTEXPR_IF_CXX17 (has_normal) {
+    return max_pos_normal_value_cxx17<Storage, NumExpBits, NumMantissaBits, NaNEncoding>();
+  }
+  else {
+    return max_pos_denormal_value_cxx17<Storage, NumExpBits, NumMantissaBits, NaNEncoding>();
+  }
+
+  CUTLASS_GCC_UNREACHABLE;
+}
+#endif
+
+template<class Storage, uint32_t NumExpBits, uint32_t NumMantissaBits, NanInfEncoding NaNEncoding, bool IsSigned>
+constexpr Storage min_value_cxx11() {
+  return
+    (IsSigned ?
+      Storage(1ull << (NumExpBits + NumMantissaBits)) | max_value_cxx11<Storage, NumExpBits, NumMantissaBits, NaNEncoding>() :
+      Storage(0)
+    );
+}
+
+#if (CUTLASS_CXX17_OR_LATER)
+template<class Storage, uint32_t NumExpBits, uint32_t NumMantissaBits, NanInfEncoding NaNEncoding, bool IsSigned>
+constexpr Storage min_value_cxx17() {
+  if (IsSigned) {
+    return Storage(1ull << (NumExpBits + NumMantissaBits)) | max_value_cxx17<Storage, NumExpBits, NumMantissaBits, NaNEncoding>();
+  }
+  else { // Unsigned number
+    return Storage(0);
+  }
+
+  CUTLASS_GCC_UNREACHABLE;
+}
+#endif
+
+template <
+    class StorageType,
+    uint32_t NumBits, uint32_t NumExpBits, uint32_t NumMantissaBits,
+    NanInfEncoding Nan = NanInfEncoding::IEEE_754, bool IsSigned = true>
+struct FpBitRepresentation {
+public:
+
+  using Storage = StorageType;
+
+#if (CUTLASS_CXX17_OR_LATER)
+  static_assert(cutlass::platform::is_unsigned_v<Storage>, "Use an unsigned integer for StorageType");
+#endif
+  static constexpr bool IS_SIGNED = IsSigned;
+  // Canonical NaN is always represented as exponent=11...11 and mantissa=11...11, if it exists
+  static constexpr NanInfEncoding NAN_TYPE = Nan;
+  // Inf is always represented as exponent=11...11 and mantissa=00...00, if it exists
+  static constexpr bool HAS_INF = (NAN_TYPE == NanInfEncoding::IEEE_754);
+  static constexpr bool HAS_NAN = (NAN_TYPE != NanInfEncoding::NONE);
+
+  static constexpr bool HAS_DENORM = (NumMantissaBits > 0);
+  static constexpr bool HAS_NORMAL = !HAS_DENORM;
+
+  static constexpr uint32_t NUM_BITS = NumBits;
+  static constexpr uint32_t NUM_EXPONENT_BITS = NumExpBits;
+  static constexpr uint32_t NUM_MANTISSA_BITS = NumMantissaBits;
+  static_assert(NUM_BITS >= (NUM_EXPONENT_BITS + NUM_MANTISSA_BITS + uint32_t(IS_SIGNED)), "Number of bits do not match");
+
+  static constexpr Storage ONE = Storage(1);
+  static constexpr Storage ZERO = Storage(0);
+
+  // Note: Don't rely on operator precedence. Use parenthesis.
+  static constexpr Storage EXPONENT_MASK = (Storage(1) << Storage(NUM_EXPONENT_BITS)) - ONE;
+  static constexpr Storage MANTISSA_MASK = (Storage(1) << Storage(NUM_MANTISSA_BITS)) - ONE;
+  static constexpr Storage EXPONENT_SHIFT = Storage(NUM_MANTISSA_BITS);
+  static constexpr Storage SIGN_SHIFT = (IS_SIGNED) ? Storage(NUM_MANTISSA_BITS + NUM_EXPONENT_BITS) : Storage(0);
+
+  // Note: All biased/real exponent calculation are done with signed ints
+  // Use unsigned to represent data not exponent.
+  static constexpr int EXP_BIAS = detail::exponent_bias_cxx11<NUM_EXPONENT_BITS, NUM_MANTISSA_BITS>();
+  static constexpr int MAX_EXP = detail::maximum_exponent_cxx11<NUM_EXPONENT_BITS, NUM_MANTISSA_BITS, NAN_TYPE>();
+  static constexpr int MIN_EXP = detail::minimum_exponent_cxx11<NUM_EXPONENT_BITS, NUM_MANTISSA_BITS>();
+
+  // Floating-point Limits
+  static constexpr Storage MAX_POS_NORMAL_VAL = detail::max_pos_normal_value_cxx11<Storage, NUM_EXPONENT_BITS, NUM_MANTISSA_BITS, NAN_TYPE>();
+  static constexpr Storage MAX_POS_DENORMAL_VAL = detail::max_pos_denormal_value_cxx11<Storage, NUM_EXPONENT_BITS, NUM_MANTISSA_BITS, NAN_TYPE>();
+  static constexpr Storage MIN_POS_NORMAL_VAL = detail::min_pos_normal_value_cxx11<Storage, NUM_EXPONENT_BITS, NUM_MANTISSA_BITS, NAN_TYPE>();
+  static constexpr Storage MIN_POS_DENORMAL_VAL = detail::min_pos_denormal_value_cxx11<Storage, NUM_EXPONENT_BITS, NUM_MANTISSA_BITS, NAN_TYPE>();
+
+  static constexpr Storage MAX_VALUE = max_value_cxx11<Storage, NUM_EXPONENT_BITS, NUM_MANTISSA_BITS, NAN_TYPE>();
+  static constexpr Storage MIN_VALUE = min_value_cxx11<Storage, NUM_EXPONENT_BITS, NUM_MANTISSA_BITS, NAN_TYPE, IS_SIGNED>();
+
+  //
+  // C++17 Verification
+  //
+#if (CUTLASS_CXX17_OR_LATER)
+  static_assert(EXP_BIAS == detail::exponent_bias_cxx17<NUM_EXPONENT_BITS, NUM_MANTISSA_BITS>(),                "Error");
+  static_assert(MAX_EXP  == detail::maximum_exponent_cxx17<NUM_EXPONENT_BITS, NUM_MANTISSA_BITS, NAN_TYPE>(),   "Error");
+  static_assert(MIN_EXP  == detail::minimum_exponent_cxx17<NUM_EXPONENT_BITS, NUM_MANTISSA_BITS>(),             "Error");
+
+  static_assert(MAX_POS_NORMAL_VAL   == detail::max_pos_normal_value_cxx17<Storage, NUM_EXPONENT_BITS, NUM_MANTISSA_BITS, NAN_TYPE>(), "Error");
+  static_assert(MAX_POS_DENORMAL_VAL == detail::max_pos_denormal_value_cxx17<Storage, NUM_EXPONENT_BITS, NUM_MANTISSA_BITS, NAN_TYPE>(), "Error");
+  static_assert(MIN_POS_NORMAL_VAL   == detail::min_pos_normal_value_cxx17<Storage, NUM_EXPONENT_BITS, NUM_MANTISSA_BITS, NAN_TYPE>(), "Error");
+  static_assert(MIN_POS_DENORMAL_VAL == detail::min_pos_denormal_value_cxx17<Storage, NUM_EXPONENT_BITS, NUM_MANTISSA_BITS, NAN_TYPE>(), "Error");
+  static_assert(MAX_VALUE            == max_value_cxx17<Storage, NUM_EXPONENT_BITS, NUM_MANTISSA_BITS, NAN_TYPE>(), "Error");
+  static_assert(MIN_VALUE            == min_value_cxx17<Storage, NUM_EXPONENT_BITS, NUM_MANTISSA_BITS, NAN_TYPE, IS_SIGNED>(), "Error");
+#endif
+
+  // If we don't have INF defined, set the largest number. Gives us .satfinite behavior.
+  static constexpr Storage INF_MASK = (HAS_INF) ?
+      (Storage(EXPONENT_MASK) << Storage(NUM_MANTISSA_BITS)) : MAX_VALUE;
+  static constexpr Storage NAN_MASK = (Storage(EXPONENT_MASK) << Storage(NUM_MANTISSA_BITS)) | MANTISSA_MASK;
+
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 bool is_inf(Storage flt) {
+    if CUTLASS_CONSTEXPR_IF_CXX17 (!HAS_INF) {
+      return false;
+    }
+    bool exp_all_ones = (exponent_bits(flt) ^ EXPONENT_MASK) == 0;
+    bool mantissa_all_zeros = mantissa_bits(flt) == 0;
+    return exp_all_ones && mantissa_all_zeros;
+  }
+
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 bool is_canonical_nan(Storage flt) {
+    if CUTLASS_CONSTEXPR_IF_CXX17 (NAN_TYPE == NanInfEncoding::NONE) {
+      return false;
+    }
+    bool exp_all_ones = (exponent_bits(flt) ^ EXPONENT_MASK) == ZERO;
+    bool mantissa_all_ones = (mantissa_bits(flt) ^ MANTISSA_MASK) == ZERO;
+    return exp_all_ones && mantissa_all_ones;
+  }
+
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 bool is_nan(Storage flt) {
+    if CUTLASS_CONSTEXPR_IF_CXX17 (NAN_TYPE == NanInfEncoding::NONE) {
+      return false;
+    }
+
+    if CUTLASS_CONSTEXPR_IF_CXX17 (NAN_TYPE == NanInfEncoding::CANONICAL_ONLY) {
+      return is_canonical_nan(flt);
+    }
+
+    bool exp_all_ones = (exponent_bits(flt) ^ EXPONENT_MASK) == ZERO;
+    bool mantissa_has_ones = mantissa_bits(flt) > ZERO;
+    return exp_all_ones && mantissa_has_ones;
+  }
+
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 bool is_denorm(Storage flt) {
+    if CUTLASS_CONSTEXPR_IF_CXX17 (!HAS_DENORM) {
+      return false;
+    }
+    else if (exponent_bits(flt) == ZERO) {
+      // Exponent bits are all 0s
+      return true;
+    }
+    return false;
+  }
+
+  template<typename T = Storage>
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 T sign_bit(T flt) {
+    if CUTLASS_CONSTEXPR_IF_CXX17 (!IS_SIGNED) {
+      return T(0);
+    }
+    return static_cast<T>(flt >> T(SIGN_SHIFT));
+  }
+
+  template<typename T = Storage>
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 T set_sign_bit(T flt, T sign) {
+    if CUTLASS_CONSTEXPR_IF_CXX17 (!IS_SIGNED) {
+      return flt;
+    }
+    return static_cast<T>(flt | (sign << T(SIGN_SHIFT)));
+  }
+
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 Storage exponent_bits(Storage flt) {
+    if CUTLASS_CONSTEXPR_IF_CXX17 (NUM_EXPONENT_BITS == ZERO) {
+      return ZERO;
+    }
+    return (flt >> (NUM_MANTISSA_BITS)) & EXPONENT_MASK;
+  }
+
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 int exponent(Storage flt) {
+    if CUTLASS_CONSTEXPR_IF_CXX17 (NUM_EXPONENT_BITS == ZERO) {
+      return -int(EXP_BIAS);
+    }
+
+    if (HAS_DENORM && (exponent_bits(flt) == ZERO)) {
+      return 1 - int(EXP_BIAS);
+    }
+
+    return int(flt >> (NUM_MANTISSA_BITS) & EXPONENT_MASK) - int(EXP_BIAS);
+  }
+
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 Storage mantissa_bits(Storage flt) {
+    if CUTLASS_CONSTEXPR_IF_CXX17 (NUM_MANTISSA_BITS == ZERO) {
+      return ZERO;
+    }
+    return (flt & MANTISSA_MASK);
+  }
+
+  template <class FpType>
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 Storage to_bits(FpType flt) {
+    return copy_bits<FpType, Storage>(flt);
+  }
+
+  template <class DstFpBits>
+  CUTLASS_HOST_DEVICE static typename DstFpBits::Storage convert_to(
+      Storage src_val,
+      DstFpBits dst_encoding) {
+    return convert(FpBitRepresentation{}, src_val, dst_encoding);
+  }
+
+  template <class SrcFpBits>
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 Storage convert_from(
+      typename SrcFpBits::Storage src_val,
+      SrcFpBits src_encoding) {
+    return convert(src_encoding, src_val, FpBitRepresentation{});
+  }
+
+private:
+
+  template<typename T = Storage>
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 T make_fp_from_bits(T sign, T exp, T mantissa) {
+    T fp_bits = T(ZERO);
+    CUTLASS_UNUSED(sign);
+    if CUTLASS_CONSTEXPR_IF_CXX17 (IS_SIGNED) {
+      fp_bits = sign << SIGN_SHIFT;
+    }
+    fp_bits |= (exp << T(NUM_MANTISSA_BITS));
+    fp_bits |= (mantissa);
+    return fp_bits;
+  }
+
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 Storage nan_with_sign(Storage sign) {
+    Storage fp_bits = NAN_MASK;
+    return set_sign_bit(fp_bits, sign);
+  }
+
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 Storage inf_with_sign(Storage sign) {
+    if CUTLASS_CONSTEXPR_IF_CXX17 (HAS_INF) {
+      Storage fp_bits = INF_MASK;
+      return set_sign_bit(fp_bits, sign);
+    }
+    else {
+      // If INF is not defined assume satfinite behavior
+      return (sign == ZERO) ? MAX_VALUE : MIN_VALUE;
+    }
+
+    CUTLASS_GCC_UNREACHABLE;
+  }
+
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 Storage significand(Storage flt) {
+    if (is_denorm(flt)) {
+      return mantissa_bits(flt);
+    }
+    else {
+      return (ONE << Storage(NUM_MANTISSA_BITS)) | mantissa_bits(flt);
+    }
+
+    CUTLASS_GCC_UNREACHABLE;
+  }
+
+  template<typename T>
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 T significand_hidden_bits(T significand) {
+    if CUTLASS_CONSTEXPR_IF_CXX17 (NUM_MANTISSA_BITS == 0) {
+      return T(1);
+    }
+    return ((T(0b11) << T(NUM_MANTISSA_BITS)) & significand) >> T(NUM_MANTISSA_BITS);
+  }
+
+  // Current assumption round to nearest even
+  template<class T>
+  CUTLASS_HOST_DEVICE
+  static CUTLASS_CONSTEXPR_IF_CXX17 T round_significand(T src, int shift_amount) {
+    T dst_mantissa = src;
+    // If the shift amount is positive, we are shifting left
+    // Type with less mantissa bits is rounded to a type with more
+    // mantissa bits.
+    if (shift_amount > 0) {
+      dst_mantissa = (dst_mantissa << (shift_amount));
+    }
+    else {
+      // There are fewer mantissa bits in the target type
+      // we need to round the destination number up for all
+      // lower precision bits removed.
+      // We assume round-to-nearest-even here.
+      int pos_shift_amount = -shift_amount;
+
+      // Too large shift return all zeros to prevent undefined behavior for shift.
+      if (pos_shift_amount >= static_cast<int>(sizeof(T) * 8)) {
+        return T(0);
+      }
+
+      T guard_bit_mask = (T(1) << T(pos_shift_amount));            // Last bit to remain in mantissa
+      T sticky_mask    = (T(1) << T(pos_shift_amount - 1)) - T(1); // Remaining bits
+      T round_bit_mask = (T(1) << T(pos_shift_amount - 1));        // First bit removed from mantissa
+
+      bool sticky_bit = (src & sticky_mask) >= T(1);                      // ORing all sticky bits
+      bool round_bit = (src & round_bit_mask) >= T(1);
+      bool guard_bit = (src & guard_bit_mask) >= T(1);
+
+      // Shift mantissa bits to right to remove lowest precision bits
+      dst_mantissa = dst_mantissa >> pos_shift_amount;
+
+      if ((sticky_bit && round_bit) || (guard_bit && round_bit && !sticky_bit)) {
+        dst_mantissa += 1;
+      }
+    }
+    return dst_mantissa;
+  }
+
+  template <class SrcFpBits, class DstFpBits>
+  CUTLASS_HOST_DEVICE
+  static typename DstFpBits::Storage convert(
+      SrcFpBits src_encoding,
+      typename SrcFpBits::Storage src_val,
+      DstFpBits dst_encoding) {
+
+    using SrcT = typename SrcFpBits::Storage;
+    using DstT = typename DstFpBits::Storage;
+    using LargeStorage = typename cutlass::platform::conditional<(sizeof(SrcT) > sizeof(DstT)), SrcT, DstT>::type;
+
+    LargeStorage src_sign_bit = src_encoding.sign_bit(src_val);
+
+    // If the source is NaN, set the destination to NaN carrying the sign bit
+    if (src_encoding.is_nan(src_val)) {
+      return dst_encoding.nan_with_sign(DstT(src_sign_bit));
+    }
+    // If the source is INF, set the destination to INF carrying the sign bit
+    else if (src_encoding.is_inf(src_val)) {
+      return dst_encoding.set_sign_bit(DstFpBits::INF_MASK, DstT(src_sign_bit));
+    }
+    // Number is not NaN or INF: Zero and others
+
+    LargeStorage src_exp_bits = src_encoding.exponent_bits(src_val);
+    LargeStorage src_significand = src_encoding.significand(src_val);
+    int src_exp = src_encoding.exponent(src_val);
+
+    // The source value is 0. Return a signed 0.
+    if (src_exp_bits == LargeStorage(0) && src_significand == LargeStorage(0)) {
+      return dst_encoding.set_sign_bit(DstT(0), DstT(src_sign_bit));
+    }
+
+#if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+    printf("(1) src_sign: %llu src_exp_bits %llx src_exp %d src_significand %llx\n",
+      static_cast<unsigned long long>(src_sign_bit), static_cast<unsigned long long>(src_exp_bits), src_exp, static_cast<unsigned long long>(src_significand));
+#endif
+    // Normalize the number: Left shift the significand bits until hidden "1" appears.
+    // Only needed if the src value is denormal.
+    // Conditions:
+    //  If the exponent is 0, then the significand can't be 0 (src_val==0 case handled above):
+    //    there is at least one "1" bit in the significand. Loop executes.
+    //  If the exponent is not 0, then the number is normal:
+    //    significand has hidden bit set. Loop doesn't execute.
+    // Assumption: Zero is always defined for the floating point types and detected above
+
+    while (src_encoding.significand_hidden_bits(src_significand) == LargeStorage(0)) {
+      src_significand <<= LargeStorage(1);
+      src_exp--;
+    }
+
+#if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+    printf("(2) src_sign: %llu src_exp_bits %llx src_exp %d src_significand %llx\n",
+      static_cast<unsigned long long>(src_sign_bit), static_cast<unsigned long long>(src_exp_bits), src_exp, static_cast<unsigned long long>(src_significand));
+#endif
+    // The exponent exceeds DstFormat's exponent capacity
+    // Return positive/negative infinity.
+    // If no INF is defined, return positive/negative largest value.
+    if (src_exp > DstFpBits::MAX_EXP) {
+      return dst_encoding.set_sign_bit(DstFpBits::INF_MASK, DstT(src_sign_bit));
+    }
+    else if (src_exp <= DstFpBits::MAX_EXP && src_exp >= DstFpBits::MIN_EXP) {
+#if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+      printf("(3) Exp match: src_sign: %d src_exp_bits: %x src_exp: %d src_significand: %x\n",
+        src_sign_bit, src_exp_bits, src_exp, src_significand);
+#endif
+
+      int shift_amount = int(DstFpBits::NUM_MANTISSA_BITS) - int(SrcFpBits::NUM_MANTISSA_BITS);
+      int dst_exponent = src_exp + DstFpBits::EXP_BIAS;
+      LargeStorage dst_mantissa = src_significand;
+
+      // if we have an M0 case, the floating point number is always denormal.
+      // Therefore, if exponents are equal, we need to check whether it is inf
+      if (DstFpBits::NUM_EXPONENT_BITS == 0) {
+        if (dst_mantissa > DstFpBits::INF_MASK) {
+          return dst_encoding.inf_with_sign(DstT(src_sign_bit));
+        }
+      }
+
+      // Round to nearest even
+      dst_mantissa = round_significand(dst_mantissa, shift_amount);
+
+#if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+      printf("(4) after rounding src_sign: %d dst_exponent: %d dst_mantissa: %x\n",
+        src_sign_bit, dst_exponent, dst_mantissa);
+#endif
+
+      if (dst_encoding.significand_hidden_bits(dst_mantissa) > 0b1) {
+        // Significant became larger than 01.X...X. Divide significand by 2 and multiply exp by 2
+        while (dst_exponent < (DstFpBits::MAX_EXP+DstFpBits::EXP_BIAS) &&
+               dst_encoding.significand_hidden_bits(dst_mantissa) > LargeStorage(0b1)) {
+          dst_mantissa >>= LargeStorage(1);
+          dst_exponent++;
+        }
+
+#if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+        printf("(5) after rounding  max_exp: %d src_sign: %d dst_exponent: %d dst_mantissa: %x\n",
+          DstFpBits::MAX_EXP,src_sign_bit, dst_exponent, dst_mantissa);
+#endif
+
+        if (dst_encoding.significand_hidden_bits(dst_mantissa) > LargeStorage(0b1)) {
+          return dst_encoding.set_sign_bit(DstFpBits::INF_MASK, DstT(src_sign_bit));
+        }
+      }
+
+      dst_mantissa = dst_mantissa & DstFpBits::MANTISSA_MASK;
+      static_assert(sizeof(LargeStorage) >= sizeof(decltype(dst_exponent)),
+        "sizeof(LargeStorage) must be greater than or equal to sizeof(decltype(dst_exponent))");
+      LargeStorage dst_exponent_bits = static_cast<LargeStorage>(dst_exponent);
+
+      DstT final_val = static_cast<DstT>(dst_encoding.template make_fp_from_bits<LargeStorage>(src_sign_bit, dst_exponent_bits, dst_mantissa));
+
+#if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+      printf("(6) Final Value src_sign: %d dst_exp_bits: %x dst_mantissa: %x\n",
+        src_sign_bit, dst_exponent_bits, dst_mantissa);
+#endif
+
+      if (DstFpBits::is_nan(final_val)) {
+        // This NAN is generated when:
+        //  Src is not an Nan
+        //  the exp of Src == the max_exp of Dst.
+        //  The mantissa becomes all-1s after rounding.
+        // Return max value of Dst (not NAN) as it just couldn't be represented in the range of Dst.
+        return dst_encoding.set_sign_bit(DstFpBits::INF_MASK, DstT(src_sign_bit));
+      }
+      else {
+        return final_val;
+      }
+    }
+    else {
+      // Result is denormal
+#if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+      printf("(7) Denormal case src_sign: %d src_exp: %d src_significand: %x MIN_EXP: %d\n",
+        src_sign_bit, src_exp, src_significand, DstFpBits::MIN_EXP);
+#endif
+
+      int exp_diff = src_exp - DstFpBits::MIN_EXP;
+      int shift_amount = int(DstFpBits::NUM_MANTISSA_BITS) - int(SrcFpBits::NUM_MANTISSA_BITS);
+      shift_amount += exp_diff;
+      LargeStorage dst_mantissa = src_significand;
+      dst_mantissa = round_significand(dst_mantissa, shift_amount);
+
+      if (dst_encoding.significand_hidden_bits(dst_mantissa) >= LargeStorage(0b1)) {
+        if CUTLASS_CONSTEXPR_IF_CXX17 (DstFpBits::NUM_EXPONENT_BITS == 0) {
+          return dst_encoding.inf_with_sign(DstT(src_sign_bit));
+        }
+        else {
+          LargeStorage dst_exp_bits = 1;
+          dst_mantissa &= DstFpBits::MANTISSA_MASK;
+          DstT final_val = static_cast<DstT>(dst_encoding.template make_fp_from_bits<LargeStorage>(src_sign_bit, dst_exp_bits, dst_mantissa));
+          return final_val;
+        }
+      }
+#if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+      printf("(7.1) Denormal case exp_diff: %d shift_amount: %d dst_mantissa %d\n", exp_diff, shift_amount, dst_mantissa);
+#endif
+      dst_mantissa &= DstFpBits::MANTISSA_MASK;
+
+#if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+      printf("(8) Final Value src_sign: %d src_exp: %d dst_mantissa: %x\n",
+        src_sign_bit, src_exp, dst_mantissa);
+#endif
+
+      DstT final_val = static_cast<DstT>(dst_encoding.template make_fp_from_bits<LargeStorage>(src_sign_bit, LargeStorage(0), dst_mantissa));
+      return final_val;
+    }
+
+    return DstT(0);
+  }
+
+  template <class StorageType_, uint32_t NumBits_, uint32_t NumExpBits_,
+            uint32_t NumMantissaBits_, NanInfEncoding Nan_, bool IsSigned_>
+            friend struct FpBitRepresentation;
+};
+
+#if (CUTLASS_CXX17_OR_LATER)
+
+template<FpEncoding FpExMyCode>
+CUTLASS_CONSTEXPR_IF_CXX17 auto fp_encoding_selector() {
+  if CUTLASS_CONSTEXPR_IF_CXX17      (FpExMyCode == FpEncoding::E11M52) { // double
+    return cutlass::detail::FpBitRepresentation<uint64_t, 64, 11, 52, cutlass::detail::NanInfEncoding::IEEE_754>{};
+  }
+  else if CUTLASS_CONSTEXPR_IF_CXX17 (FpExMyCode == FpEncoding::E8M23)  { // float
+    return cutlass::detail::FpBitRepresentation<uint32_t, 32, 8, 23, cutlass::detail::NanInfEncoding::IEEE_754>{};
+  }
+  else if CUTLASS_CONSTEXPR_IF_CXX17 (FpExMyCode == FpEncoding::E5M2)   {   // FP8
+    return cutlass::detail::FpBitRepresentation<uint8_t, 8, 5, 2, cutlass::detail::NanInfEncoding::IEEE_754>{};
+  }
+  else if CUTLASS_CONSTEXPR_IF_CXX17 (FpExMyCode == FpEncoding::E4M3)   {   // FP8
+    return cutlass::detail::FpBitRepresentation<uint8_t, 8, 4, 3, cutlass::detail::NanInfEncoding::CANONICAL_ONLY>{};
+  }
+  
+  else if CUTLASS_CONSTEXPR_IF_CXX17 (FpExMyCode == FpEncoding::UE4M3)   {   // FP8
+    return cutlass::detail::FpBitRepresentation<uint8_t, 8, 4, 3, cutlass::detail::NanInfEncoding::CANONICAL_ONLY, false>{};
+  }
+  
+  else if CUTLASS_CONSTEXPR_IF_CXX17 (FpExMyCode == FpEncoding::UE8M0)   {   // FP8
+    return cutlass::detail::FpBitRepresentation<uint8_t, 8, 8, 0, cutlass::detail::NanInfEncoding::CANONICAL_ONLY, false>{};
+  }
+  else if CUTLASS_CONSTEXPR_IF_CXX17 (FpExMyCode == FpEncoding::E3M2)   {   // FP6
+    return cutlass::detail::FpBitRepresentation<uint8_t, 6, 3, 2, cutlass::detail::NanInfEncoding::NONE>{};
+  }
+  else if CUTLASS_CONSTEXPR_IF_CXX17 (FpExMyCode == FpEncoding::E2M3)   {   // FP6
+    return cutlass::detail::FpBitRepresentation<uint8_t, 6, 2, 3, cutlass::detail::NanInfEncoding::NONE>{};
+  }
+  else if CUTLASS_CONSTEXPR_IF_CXX17 (FpExMyCode == FpEncoding::E2M1)   {   // FP4
+    return cutlass::detail::FpBitRepresentation<uint8_t, 4, 2, 1, cutlass::detail::NanInfEncoding::NONE>{};
+  }
+  CUTLASS_GCC_UNREACHABLE;
+}
+
+#else
+//
+// Definitions for floating point encodings.
+//
+
+template <FpEncoding FpExMyCode> struct FpEncodingSelector {
+  using type = void;
+};
+
+template <> struct FpEncodingSelector<FpEncoding::E11M52> {
+  using type = cutlass::detail::FpBitRepresentation<uint64_t, 64, 11, 52, cutlass::detail::NanInfEncoding::IEEE_754>;
+};
+
+template <> struct FpEncodingSelector<FpEncoding::E8M23> {
+  using type = cutlass::detail::FpBitRepresentation<uint32_t, 32, 8, 23, cutlass::detail::NanInfEncoding::IEEE_754>;
+};
+template <> struct FpEncodingSelector<FpEncoding::E5M2> {
+  using type = cutlass::detail::FpBitRepresentation<uint8_t, 8, 5, 2, cutlass::detail::NanInfEncoding::IEEE_754>;
+};
+
+template <> struct FpEncodingSelector<FpEncoding::E4M3> {
+  using type = cutlass::detail::FpBitRepresentation<uint8_t, 8, 4, 3, cutlass::detail::NanInfEncoding::CANONICAL_ONLY>;
+};
+
+template <> struct FpEncodingSelector<FpEncoding::UE4M3> {
+  using type = cutlass::detail::FpBitRepresentation<uint8_t, 8, 4, 3, cutlass::detail::NanInfEncoding::CANONICAL_ONLY, false>;
+};
+
+template <> struct FpEncodingSelector<FpEncoding::UE8M0> {
+  using type = cutlass::detail::FpBitRepresentation<uint8_t, 8, 8, 0, cutlass::detail::NanInfEncoding::CANONICAL_ONLY, false>;
+};
+
+template <> struct FpEncodingSelector<FpEncoding::E3M2> {
+  using type = cutlass::detail::FpBitRepresentation<uint8_t, 6, 3, 2, cutlass::detail::NanInfEncoding::NONE>;
+};
+
+template <> struct FpEncodingSelector<FpEncoding::E2M3> {
+  using type = cutlass::detail::FpBitRepresentation<uint8_t, 6, 2, 3, cutlass::detail::NanInfEncoding::NONE>;
+};
+
+template <> struct FpEncodingSelector<FpEncoding::E2M1> {
+  using type = cutlass::detail::FpBitRepresentation<uint8_t, 4, 2, 1, cutlass::detail::NanInfEncoding::NONE>;
+};
+#endif
+
+} // namespace detail
+
+template <detail::FpEncoding T, class Derived>
+struct float_exmy_base
+{
+
+  static constexpr detail::FpEncoding Encoding = T;
+  using BitRepresentation =
+    #if (CUTLASS_CXX17_OR_LATER)
+      decltype(detail::fp_encoding_selector<T>())
+    #else
+      typename detail::FpEncodingSelector<T>::type
+    #endif
+      ;
+
+  using FP32BitRepresentation =
+    #if (CUTLASS_CXX17_OR_LATER)
+      decltype(cutlass::detail::fp_encoding_selector<cutlass::detail::FpEncoding::E8M23>())
+    #else
+      typename detail::FpEncodingSelector<cutlass::detail::FpEncoding::E8M23>::type
+    #endif
+      ;
+
+  using Storage = typename BitRepresentation::Storage;
+
+  //
+  // Data members
+  //
+
+  /// Data container
+  Storage storage;
+
+  /// Ctors.
+  float_exmy_base() = default;
+
+  CUTLASS_HOST_DEVICE
+  float_exmy_base(Storage s) : storage(s) {
+  }
+
+  /// Is finite implementation
+  CUTLASS_HOST_DEVICE
+  static bool isfinite(float_exmy_base flt) {
+    return !BitRepresentation::is_inf(flt.storage);
+  }
+
+  /// Is NaN implementation
+  CUTLASS_HOST_DEVICE
+  static bool isnan(float_exmy_base flt) {
+    return BitRepresentation::is_nan(flt.storage);
+  }
+
+  /// Is infinite implementation
+  CUTLASS_HOST_DEVICE
+  static bool isinf(float_exmy_base flt) {
+    return BitRepresentation::is_inf(flt.storage);
+  }
+
+  /// Is infinite implementation
+  CUTLASS_HOST_DEVICE
+  static bool isnormal(float_exmy_base flt) {
+    return !BitRepresentation::is_denorm(flt.storage);
+  }
+
+  CUTLASS_HOST_DEVICE
+  static float_exmy_base<T, Derived> bitcast(Storage x) {
+    float_exmy_base f;
+    f.storage = x;
+    return f;
+  }
+
+  CUTLASS_HOST_DEVICE
+  float_exmy_base convert_from_float(float const &flt) const {
+    FP32BitRepresentation::Storage fp32_bits = FP32BitRepresentation::to_bits(flt);
+    float_exmy_base float_exmy;
+    float_exmy.storage = BitRepresentation::convert_from(fp32_bits, FP32BitRepresentation{});
+    return float_exmy;
+  }
+
+  CUTLASS_HOST_DEVICE
+  float convert_to_float(float_exmy_base<T, Derived> const &x) const {
+    FP32BitRepresentation::Storage fp32_bits;
+    fp32_bits = BitRepresentation::convert_to(x.storage, FP32BitRepresentation{});
+    return detail::copy_bits<FP32BitRepresentation::Storage, float>(fp32_bits);
+  }
+
+  // Note: Only consider float/int conversions in this Base class
+  // Types inheriting from this class should define their own constructors and
+  // specialized type conversions
+
+  /// Floating point conversion
+  CUTLASS_HOST_DEVICE
+  explicit float_exmy_base<T, Derived>(float x) {
+    storage = static_cast<Derived*>(this)->convert_from_float(x).storage;
+  }
+
+  // Integer conversion
+  CUTLASS_HOST_DEVICE
+  explicit float_exmy_base<T, Derived>(int x) {
+    storage = static_cast<Derived*>(this)->convert_from_float(float(x)).storage;
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_exmy_base<T, Derived>(unsigned x) {
+    storage = static_cast<Derived*>(this)->convert_from_float(float(x)).storage;
+  }
+
+  /// Converts to float
+  CUTLASS_HOST_DEVICE
+  operator float() const {
+    return static_cast<const Derived*>(this)->convert_to_float(*this);
+  }
+
+  /// Converts to int
+  CUTLASS_HOST_DEVICE
+  explicit operator int() const {
+    return int(static_cast<const Derived*>(this)->convert_to_float(*this));
+  }
+
+  /// Accesses raw internal state
+  CUTLASS_HOST_DEVICE
+  Storage &raw() {
+    return storage;
+  }
+
+  /// Accesses raw internal state
+  CUTLASS_HOST_DEVICE
+  Storage raw() const {
+    return storage;
+  }
+
+  /// Returns the sign bit
+  CUTLASS_HOST_DEVICE
+  bool signbit() const {
+    return bool(BitRepresentation::sign_bit(storage));
+  }
+
+  /// Returns the biased exponent
+  CUTLASS_HOST_DEVICE
+  int exponent_biased() const {
+    return int(BitRepresentation::exponent_bits(storage));
+  }
+
+  /// Returns the unbiased exponent
+  CUTLASS_HOST_DEVICE
+  int exponent() const {
+    return int(BitRepresentation::exponent(storage));
+  }
+
+  /// Returns the mantissa
+  CUTLASS_HOST_DEVICE
+  int mantissa() const {
+    return int(BitRepresentation::mantissa_bits(storage));
+  }
+
+  ///////////////////////////////////////////////////////////////////////////////////////////////////
+  //
+  // Arithmetic operators
+  //
+  ///////////////////////////////////////////////////////////////////////////////////////////////////
+
+  // Note: Almost all data types cast to float then do the arithmetic operations
+  // Types inheriting from this class can overload them if specialized instructions are available
+  // in HW (e.g. half_t)
+
+
+  CUTLASS_HOST_DEVICE
+  friend bool operator==(float_exmy_base const &lhs, float_exmy_base const &rhs) {
+    return float(lhs) == float(rhs);
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend bool operator!=(float_exmy_base const &lhs, float_exmy_base const &rhs) {
+    return float(lhs) != float(rhs);
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend bool operator<(float_exmy_base const &lhs, float_exmy_base const &rhs) {
+    return float(lhs) < float(rhs);
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend bool operator<=(float_exmy_base const &lhs, float_exmy_base const &rhs) {
+    return float(lhs) <= float(rhs);
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend bool operator>(float_exmy_base const &lhs, float_exmy_base const &rhs) {
+    return float(lhs) > float(rhs);
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend bool operator>=(float_exmy_base const &lhs, float_exmy_base const &rhs) {
+    return float(lhs) >= float(rhs);
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend float_exmy_base operator+(float_exmy_base const &lhs, float_exmy_base const &rhs) {
+    return float_exmy_base(float(lhs) + float(rhs));
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend float_exmy_base operator-(float_exmy_base const &lhs) {
+    return float_exmy_base(-float(lhs));
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend float_exmy_base operator-(float_exmy_base const &lhs, float_exmy_base const &rhs) {
+    return float_exmy_base(float(lhs) - float(rhs));
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend float_exmy_base operator*(float_exmy_base const &lhs, float_exmy_base const &rhs) {
+    return float_exmy_base(float(lhs) * float(rhs));
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend float_exmy_base operator/(float_exmy_base const &lhs, float_exmy_base const &rhs) {
+    return float_exmy_base(float(lhs) / float(rhs));
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend float_exmy_base &operator+=(float_exmy_base &lhs, float_exmy_base const &rhs) {
+    lhs = float_exmy_base(float(lhs) + float(rhs));
+    return lhs;
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend float_exmy_base &operator-=(float_exmy_base &lhs, float_exmy_base const &rhs) {
+    lhs = float_exmy_base(float(lhs) - float(rhs));
+    return lhs;
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend float_exmy_base &operator*=(float_exmy_base &lhs, float_exmy_base const &rhs) {
+    lhs = float_exmy_base(float(lhs) * float(rhs));
+    return lhs;
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend float_exmy_base &operator/=(float_exmy_base &lhs, float_exmy_base const &rhs) {
+    lhs = float_exmy_base(float(lhs) / float(rhs));
+    return lhs;
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend float_exmy_base &operator++(float_exmy_base &lhs) {
+    float tmp(lhs);
+    ++tmp;
+    lhs = float_exmy_base(tmp);
+    return lhs;
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend float_exmy_base &operator--(float_exmy_base &lhs) {
+    float tmp(lhs);
+    --tmp;
+    lhs = float_exmy_base(tmp);
+    return lhs;
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend float_exmy_base operator++(float_exmy_base &lhs, int) {
+    float_exmy_base ret(lhs);
+    float tmp(lhs);
+    tmp++;
+    lhs = float_exmy_base(tmp);
+    return ret;
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend float_exmy_base operator--(float_exmy_base &lhs, int) {
+    float_exmy_base ret(lhs);
+    float tmp(lhs);
+    tmp--;
+    lhs = float_exmy_base(tmp);
+    return ret;
+  }
+
+};
+
+template <detail::FpEncoding T, class Derived>
+CUTLASS_HOST_DEVICE
+cutlass::float_exmy_base<T, Derived> abs(cutlass::float_exmy_base<T, Derived> const& h) {
+  using BitRepresentation = typename cutlass::float_exmy_base<T, Derived>::BitRepresentation;
+  using Storage = typename cutlass::float_exmy_base<T, Derived>::Storage;
+  return BitRepresentation::IS_SIGNED ?
+      cutlass::float_exmy_base<T, Derived>(Storage(h.raw() & Storage((1<<BitRepresentation::SIGN_SHIFT) - 1))) :
+      cutlass::float_exmy_base<T, Derived>(h.raw());
+}
+} // namespace cutlass

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/experimental/distributed/device/detail.hpp

@@ -0,0 +1,163 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Distributed gemm device layer helpers.
+*/
+
+#pragma once
+
+#include "cute/layout.hpp"
+#include "cute/tensor.hpp"
+#include "cutlass/cutlass.h"
+
+///////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass::distributed::device::detail {
+
+
+cutlass::Status check_cuda_status(cudaError_t status) {
+  if (status != cudaSuccess) {
+    auto result = cudaGetLastError();
+    CUTLASS_TRACE_HOST("  error message: " << cudaGetErrorString(result));
+    return cutlass::Status::kErrorInternal;
+  }
+  return cutlass::Status::kSuccess;                   
+}
+
+// DistGemmBufferHelper computes required buffer size and offsets for GEMM operands.
+template <
+  typename Tiler_, 
+  typename ElementA_,
+  typename ElementB_,
+  typename ElementC_,
+  typename ElementD_>
+struct DistGemmBufferHelper {
+
+  using Tiler = Tiler_;
+
+  using ElementA = ElementA_;
+  using ElementB = ElementB_;
+  using ElementC = ElementC_;
+  using ElementD = ElementD_;
+
+  static constexpr int NumBuffersA = Tiler::NumBuffersA;
+  static constexpr int NumBuffersB = Tiler::NumBuffersB;
+  static constexpr int NumBuffersC = Tiler::NumBuffersC;
+  static constexpr int NumBuffersD = Tiler::NumBuffersD;
+
+  template <typename ProblemShape>
+  static auto
+  get_buffer_size_a(ProblemShape problem_shape) {
+    auto a_buffer_layout = cute::make_layout(
+        cute::make_shape(NumBuffersA, Tiler::get_local_a_shape(problem_shape), sizeof(ElementA))
+    );
+    return size(a_buffer_layout);
+  }
+
+  template <typename ProblemShape>
+  static auto
+  get_buffer_size_b(ProblemShape problem_shape) {
+    auto b_buffer_layout = cute::make_layout(
+        cute::make_shape(NumBuffersB, Tiler::get_local_b_shape(problem_shape), sizeof(ElementB))
+    );
+    return size(b_buffer_layout);
+  }
+
+  template <typename ProblemShape>
+  static auto
+  get_buffer_size_c(ProblemShape problem_shape) {
+    auto c_buffer_layout = cute::make_layout(
+        cute::make_shape(NumBuffersC, Tiler::get_local_c_shape(problem_shape), sizeof(ElementC))
+    );
+    return size(c_buffer_layout);
+  }
+
+  template <typename ProblemShape>
+  static auto
+  get_buffer_size_d(ProblemShape problem_shape) {
+    auto d_buffer_layout = cute::make_layout(
+        cute::make_shape(NumBuffersD, Tiler::get_local_d_shape(problem_shape), sizeof(ElementD))
+    );
+    return size(d_buffer_layout);
+  }
+
+  template <typename ProblemShape>
+  static auto
+  get_buffer_size(ProblemShape problem_shape) {
+    size_t buffer_size = 0;
+
+    if constexpr (NumBuffersA > 0) {
+      buffer_size += get_buffer_size_a(problem_shape);
+    }
+    if constexpr (NumBuffersB > 0) {
+      buffer_size += get_buffer_size_b(problem_shape);
+    }
+    if constexpr (NumBuffersC > 0) {
+      buffer_size += get_buffer_size_c(problem_shape);
+    }
+    if constexpr (NumBuffersD > 0) {
+      buffer_size += get_buffer_size_d(problem_shape);
+    }
+
+    return buffer_size;
+  }
+
+  // Buffer space: |  buffer_A  |  buffer_B  |  buffer_C  |  buffer_D  |
+  // And buffer_{A,B,C,D}: |  iter 1  |  iter 2  | ... |  iter TP - 1 |
+  template <typename ProblemShape>
+  static size_t
+  get_buffer_offset_A(ProblemShape problem_shape) {
+    return 0;
+  }
+
+  template <typename ProblemShape>
+  static size_t
+  get_buffer_offset_B(ProblemShape problem_shape) {
+    return get_buffer_size_a(problem_shape);
+  }
+
+  template <typename ProblemShape>
+  static size_t
+  get_buffer_offset_C(ProblemShape problem_shape) {
+    return get_buffer_size_a(problem_shape) + get_buffer_size_b(problem_shape);
+  }
+
+  template <typename ProblemShape>
+  static size_t
+  get_buffer_offset_D(ProblemShape problem_shape) {
+    return get_buffer_size_a(problem_shape) + get_buffer_size_b(problem_shape) + get_buffer_size_c(problem_shape);
+  }
+};
+
+} // namespace cutlass::distributed::device::detail
+
+///////////////////////////////////////////////////////////////////////////////
+

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/experimental/distributed/device/dist_gemm_universal_wrapper.hpp

@@ -0,0 +1,717 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*!
+  \file Distributed GEMM Device Adapter
+
+  Sets up local GEMM stages, the cuda graph, manages buffer and barrier spaces,
+  and maps arguments to per-stage arguments.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/device_kernel.h"
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+
+#include "cutlass/experimental/distributed/device/full_barrier.hpp"
+#include "cutlass/experimental/distributed/device/detail.hpp"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass::distributed::device {
+
+template <class GemmKernel_>
+class DistributedGemmUniversalAdapter {
+public:
+  using DeviceGemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel_>;
+  using GemmKernel = GemmKernel_;
+  using TileShape = typename GemmKernel::TileShape;
+  using ElementA = typename GemmKernel::ElementA;
+  using ElementB = typename GemmKernel::ElementB;
+  using ElementC = typename GemmKernel::ElementC;
+  using ElementD = typename GemmKernel::ElementD;
+  using ElementAccumulator = typename GemmKernel::ElementAccumulator;
+  using DispatchPolicy = typename GemmKernel::DispatchPolicy;
+  using CollectiveMainloop = typename GemmKernel::CollectiveMainloop;
+  using CollectiveEpilogue = typename GemmKernel::CollectiveEpilogue;
+
+  // "Inherit" type decls and static values from device GEMM
+  using LayoutA = typename DeviceGemm::LayoutA;
+  using LayoutB = typename DeviceGemm::LayoutB;
+  using LayoutC = typename DeviceGemm::LayoutC;
+  using LayoutD = typename DeviceGemm::LayoutD;
+
+  using StrideA = typename GemmKernel::StrideA;
+  using StrideB = typename GemmKernel::StrideB;
+  using StrideC = typename GemmKernel::StrideC;
+  using StrideD = typename GemmKernel::StrideD;
+
+  static bool const kEnableCudaHostAdapter = DeviceGemm::kEnableCudaHostAdapter;
+
+  static ComplexTransform const kTransformA = DeviceGemm::kTransformA;
+  static ComplexTransform const kTransformB = DeviceGemm::kTransformB;
+
+  using MathOperator = typename DeviceGemm::MathOperator;
+  using OperatorClass = typename DeviceGemm::OperatorClass;
+  using ArchTag = typename DeviceGemm::ArchTag;
+
+  using ThreadblockSwizzle = typename DeviceGemm::ThreadblockSwizzle;
+  using ThreadblockShape = typename DeviceGemm::ThreadblockShape;
+  using ClusterShape = typename DeviceGemm::ClusterShape;
+  using InstructionShape = typename DeviceGemm::InstructionShape;
+
+  static int const kThreadCount = DeviceGemm::kThreadCount;
+  static constexpr int WarpsInMma = DeviceGemm::WarpsInMma;
+  static constexpr int WarpsInMmaM = DeviceGemm::WarpsInMmaM;
+  static constexpr int WarpsInMmaN = DeviceGemm::WarpsInMmaN;
+
+  using WarpCount = typename DeviceGemm::WarpCount;
+  using WarpShape = typename DeviceGemm::WarpShape;
+
+  static int constexpr kStages = DeviceGemm::kStages;
+
+  static int constexpr kAlignmentA = DeviceGemm::kAlignmentA;
+  static int constexpr kAlignmentB = DeviceGemm::kAlignmentB;
+  static int constexpr kAlignmentC = DeviceGemm::kAlignmentC;
+  static int constexpr kAlignmentD = DeviceGemm::kAlignmentD;
+
+  using EpilogueOutputOp = typename DeviceGemm::EpilogueOutputOp;
+
+  static int constexpr kSplitKAlignment = DeviceGemm::kSplitKAlignment;
+
+  // Distributed GEMM types and defs
+  using DistSchedule = typename GemmKernel::DistSchedule;
+  static constexpr bool HasMemcpy = DistSchedule::HasMemcpy;
+  using TP = typename DistSchedule::TP;
+  static constexpr int TP_ = TP{};
+  using ElementFlag = typename GemmKernel::ElementFlag;
+  using ElementBarrier = uint32_t;
+
+  using BufferHelper = detail::DistGemmBufferHelper<
+    DistSchedule,
+    ElementA,
+    ElementB,
+    ElementC,
+    ElementD>;
+
+  /// Argument structure
+  using Arguments = typename GemmKernel::BaseArguments;
+  using DistributedArguments = typename GemmKernel::DistributedArguments;
+  using PackedArguments = typename GemmKernel::PackedArguments;
+
+  /// Argument structure: Kernel API
+  using Params = typename GemmKernel::PackedParams;
+
+  struct DistributedGemmState {
+    int device_idx;
+
+    Params params_array[TP_];
+
+    cudaGraph_t graph;
+    cudaGraphExec_t graph_executable;
+
+    bool graph_created = false;
+    bool graph_instantiated = false;
+
+    void * memcpy_source_ptr_array[TP_];
+    void const * memcpy_remote_ptr_array[TP_];
+    size_t memcpy_bytes[TP_];
+
+    cutlass::Array<ElementBarrier*, TP_> device_barrier_ptrs;
+
+    bool is_initialized = false;
+  };
+
+private:
+
+  DistributedGemmState state_;
+
+public:
+
+  bool is_initialized() {
+    return state_.is_initialized && state_.graph_created && state_.graph_instantiated;
+  }
+
+  /// Determines whether the GEMM can execute the given problem.
+  static Status
+  can_implement(Arguments const& args) {
+    if (args.epilogue.thread.beta != 0.0 && DistSchedule::RemoteC) {
+      CUTLASS_TRACE_HOST("  CAN IMPLEMENT: Selected TP uses Remote C to communicate " <<
+          "partial results, which do not support non-zero values for beta yet " <<
+          "(epilogue must be sourceless.)\n");
+      return Status::kInvalid;
+    }
+
+    if (not DistSchedule::can_implement_global(args.problem_shape)) {
+      CUTLASS_TRACE_HOST("  CAN IMPLEMENT: Problem shape not divisible by TP.\n");
+      return Status::kInvalid;
+    }
+
+    Arguments args_copy = args;
+    args_copy.problem_shape = DistSchedule::get_local_gemm_shape(args.problem_shape);
+    for (int iteration = 0; iteration < TP_; ++iteration) {
+      if (not GemmKernel::can_implement(args_copy)) {
+        return Status::kInvalid;
+      }
+    }
+    return Status::kSuccess;
+  }
+
+  /// Gets buffer space size
+  static size_t
+  get_buffer_space_size(Arguments const& args) {
+    size_t buffer_bytes = 0;
+
+    buffer_bytes = BufferHelper::get_buffer_size(args.problem_shape);
+    buffer_bytes = round_nearest(buffer_bytes, MinWorkspaceAlignment);
+
+    return buffer_bytes;
+  }
+
+  static auto
+  get_tensor_A_for_iter(Arguments const* args_array, void** buffer_space, int device_idx, int iteration) {
+    auto args = args_array[device_idx];
+    auto tensor_A = make_tensor(args.mainloop.ptr_A, make_layout(
+          DistSchedule::get_local_a_shape(args.problem_shape),
+          args.mainloop.dA));
+
+    uint8_t* tensor_buffer = reinterpret_cast<uint8_t*>(buffer_space[device_idx]) +
+      BufferHelper::get_buffer_offset_A(args.problem_shape);
+
+    return DistSchedule::get_tensor_A(tensor_A, tensor_buffer, device_idx, iteration);
+  }
+
+  static auto
+  get_tensor_B_for_iter(Arguments const* args_array, void** buffer_space, int device_idx, int iteration) {
+    auto args = args_array[device_idx];
+    auto tensor_B = make_tensor(args.mainloop.ptr_B, make_layout(
+          DistSchedule::get_local_b_shape(args.problem_shape),
+          args.mainloop.dB));
+
+    uint8_t* tensor_buffer = reinterpret_cast<uint8_t*>(buffer_space[device_idx]) +
+      BufferHelper::get_buffer_offset_B(args.problem_shape);
+
+    return DistSchedule::get_tensor_B(tensor_B, tensor_buffer, device_idx, iteration);
+  }
+
+  static auto
+  get_tensor_C_for_iter(Arguments const* args_array, void** buffer_space, int device_idx, int iteration) {
+    auto args = args_array[device_idx];
+    auto tensor_C = make_tensor(args.epilogue.ptr_C, make_layout(
+          DistSchedule::get_local_c_shape(args.problem_shape),
+          args.epilogue.dC));
+
+    auto peer_idx_iter = DistSchedule::get_remote_peer_id(device_idx, iteration);
+    void* buffer_ptr = DistSchedule::RemoteC ? buffer_space[peer_idx_iter] : buffer_space[device_idx];
+
+    uint8_t* tensor_buffer = reinterpret_cast<uint8_t*>(buffer_ptr) +
+      BufferHelper::get_buffer_offset_C(args.problem_shape);
+
+    return DistSchedule::get_tensor_C(tensor_C, tensor_buffer, device_idx, iteration);
+  }
+
+  static auto
+  get_tensor_D_for_iter(Arguments const* args_array, void** buffer_space, int device_idx, int iteration) {
+    auto args = args_array[device_idx];
+    auto tensor_D = make_tensor(args.epilogue.ptr_D, make_layout(
+          DistSchedule::get_local_d_shape(args.problem_shape),
+          args.epilogue.dD));
+
+    // support remoteD
+    uint8_t* tensor_buffer = reinterpret_cast<uint8_t*>(buffer_space[device_idx]) +
+      BufferHelper::get_buffer_offset_D(args.problem_shape);
+
+    return DistSchedule::get_tensor_D(tensor_D, tensor_buffer, device_idx, iteration);
+  }
+
+  static size_t
+  get_workspace_size(Arguments const& args) {
+    size_t workspace_bytes = 0;
+
+    workspace_bytes = get_buffer_space_size(args);
+
+    for (int iteration = 0; iteration < TP_; ++iteration) {
+      // NOTE: assumes underlying kernels align up to alignment requirements on their own,
+      // and that the alignment requirements of the individual kernels match.
+      workspace_bytes += GemmKernel::get_workspace_size(args);
+    }
+
+    return workspace_bytes;
+  }
+
+  static size_t
+  get_barrier_bytes() {
+    return round_nearest(sizeof(ElementBarrier), 32);
+  }
+
+  static size_t
+  get_flag_bytes() {
+    return round_nearest(sizeof(ElementFlag) * TP_, 32);
+  }
+
+  static void *
+  exclusive_workspace_ptr_to_flag_ptr(void * exclusive_workspace_ptr, int iteration) {
+    return static_cast<void*>(
+        static_cast<uint8_t*>(exclusive_workspace_ptr) + 
+        get_barrier_bytes() + 
+        (sizeof(ElementFlag) * iteration));
+  }
+
+  static size_t
+  get_exclusive_workspace_size() {
+    return get_barrier_bytes() + get_flag_bytes();
+  }
+
+  /// Initializes GEMM state from arguments.
+  Status
+  initialize(
+    Arguments const* args,
+    void** workspace_ptrs,
+    void** exclusive_workspace_ptrs,
+    int device_idx,
+    cudaStream_t stream = nullptr,
+    bool launch_with_pdl = false) {
+
+    CUTLASS_TRACE_HOST("DistributedGemm::initialize() - stream: " << (stream ? "non-null" : "null"));
+
+    state_.device_idx = device_idx;
+
+    for (int device = 0; device < TP_; ++device) {
+      state_.device_barrier_ptrs[device] = reinterpret_cast<ElementBarrier*>(exclusive_workspace_ptrs[device]);
+    }
+
+    // Zero out exclusive workspace
+    zero_workspace(exclusive_workspace_ptrs[device_idx], get_exclusive_workspace_size(), stream, nullptr);
+
+    for (int iteration = 0; iteration < TP_; ++iteration) {
+
+      size_t workspace_iteration_offset = GemmKernel::get_workspace_size(args[device_idx]);
+      uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace_ptrs[device_idx]) + 
+        get_buffer_space_size(args[device_idx]) + 
+        (iteration * workspace_iteration_offset);
+
+      void * workspace_iter = reinterpret_cast<void*>(workspace_ptr);
+      void** buffer_space = workspace_ptrs;
+
+      // Set up GEMM arguments for the current stage/iteration
+      auto tensor_a_iter = get_tensor_A_for_iter(args, buffer_space, device_idx, iteration);
+      auto tensor_b_iter = get_tensor_B_for_iter(args, buffer_space, device_idx, iteration);
+      auto tensor_c_iter = get_tensor_C_for_iter(args, buffer_space, device_idx, iteration);
+      auto tensor_d_iter = get_tensor_D_for_iter(args, buffer_space, device_idx, iteration);
+
+      Arguments base_args = args[device_idx];
+      base_args.problem_shape = DistSchedule::get_local_gemm_shape(args[device_idx].problem_shape);
+      base_args.mainloop = {
+        reinterpret_cast<const ElementA*>(tensor_a_iter.data()),
+        tensor_a_iter.stride(),
+        reinterpret_cast<const ElementB*>(tensor_b_iter.data()),
+        tensor_b_iter.stride()
+      };
+      base_args.epilogue = {
+        base_args.epilogue.thread,
+        reinterpret_cast<const ElementC*>(tensor_c_iter.data()),
+        tensor_c_iter.stride(),
+        reinterpret_cast<ElementD*>(tensor_d_iter.data()),
+        tensor_d_iter.stride()
+      };
+
+      if constexpr (DistSchedule::RemoteC) {
+        if (iteration > 0) {
+          base_args.epilogue.thread.beta = 1.0;
+        }
+        else if (iteration == 0){
+          base_args.epilogue.thread.beta = 0.0;
+        }
+      }
+
+      auto [left_peer_idx, right_peer_idx] = DistSchedule::get_peers_for_device(device_idx);
+      auto flag_peer_idx = DistSchedule::KernelWritesArrivalFlag ? right_peer_idx : device_idx;
+
+      void * self_flag_ptr = exclusive_workspace_ptr_to_flag_ptr(exclusive_workspace_ptrs[device_idx], iteration);
+      void * peer_flag_ptr = exclusive_workspace_ptr_to_flag_ptr(exclusive_workspace_ptrs[flag_peer_idx], iteration);
+
+      DistributedArguments distributed_args = {
+        device_idx,
+        iteration,
+        self_flag_ptr,
+        peer_flag_ptr
+      };
+      PackedArguments args_iter = {base_args, distributed_args};
+
+      // Initialize the workspace
+      Status status = GemmKernel::initialize_workspace(args_iter, workspace_iter, stream);
+      if (status != Status::kSuccess) {
+        return status;
+      }
+
+      // Initialize the Params structure
+      state_.params_array[iteration] = GemmKernel::to_underlying_arguments(args_iter, workspace_iter);
+
+      // Set up peer buffer ptrs
+      if (iteration > 0 && HasMemcpy) {
+        auto peer_idx_iter = DistSchedule::get_remote_peer_id(device_idx, iteration);
+
+        void * local_ptr_itr = nullptr;
+        void const * remote_ptr_itr = nullptr;
+        size_t local_size = 0;
+        size_t remote_size = 0;
+
+        static_assert(not DistSchedule::HasMemcpy || (
+              DistSchedule::MemcpyA || DistSchedule::MemcpyB),
+            "Expected to either memcpy A or B when scheduler requires memcpy.");
+        if constexpr (DistSchedule::MemcpyA) {
+          local_size = cute::cosize(tensor_a_iter.layout()) * sizeof(ElementA);
+          local_ptr_itr = reinterpret_cast<void*>(tensor_a_iter.data());
+
+          // Copy peer's slice in the first iteration (direct access memcpy instead of logical ring)
+          auto remote_tensor_iter = get_tensor_A_for_iter(args, buffer_space, peer_idx_iter, 0);
+          remote_ptr_itr = reinterpret_cast<void const*>(remote_tensor_iter.data());
+          remote_size = cute::cosize(remote_tensor_iter.layout()) * sizeof(ElementA);
+        }
+        else if constexpr (DistSchedule::MemcpyB) {
+          local_size = cute::cosize(tensor_b_iter.layout()) * sizeof(ElementB);
+          local_ptr_itr = reinterpret_cast<void*>(tensor_b_iter.data());
+
+          // Copy peer's slice in the first iteration (direct access memcpy instead of logical ring)
+          auto remote_tensor_iter = get_tensor_B_for_iter(args, buffer_space, peer_idx_iter, 0);
+          remote_ptr_itr = reinterpret_cast<void const*>(remote_tensor_iter.data());
+          remote_size = cute::cosize(remote_tensor_iter.layout()) * sizeof(ElementB);
+        }
+
+        assert(local_size == remote_size && local_size > 0);
+
+        state_.memcpy_source_ptr_array[iteration] = local_ptr_itr;
+        state_.memcpy_remote_ptr_array[iteration] = remote_ptr_itr;
+        state_.memcpy_bytes[iteration] = local_size;
+      }
+    }
+
+    //
+    // Account for dynamic smem capacity if needed
+    //
+    int smem_size = GemmKernel::SharedStorageSize;
+
+    if (smem_size >= (48 << 10)) {
+      CUTLASS_TRACE_HOST("  Setting smem size to " << smem_size);
+      cudaError_t result = cudaFuncSetAttribute(
+          device_kernel<GemmKernel>,
+          cudaFuncAttributeMaxDynamicSharedMemorySize,
+          smem_size);
+      if (cudaSuccess != result) {
+        result = cudaGetLastError(); // to clear the error bit
+        CUTLASS_TRACE_HOST("  cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result));
+        return Status::kErrorInternal;
+      }
+    }
+
+    state_.is_initialized = true;
+
+    // Instantiate graph
+    Status status = construct_graph(launch_with_pdl);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    return Status::kSuccess;
+  }
+
+  Status
+  construct_graph(bool launch_with_pdl) {
+#if (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 6))
+    Status status = Status::kSuccess;
+
+    // Destroy existing graph, if created
+    if (state_.graph_created) {
+      status = detail::check_cuda_status(cudaGraphDestroy(state_.graph));
+      if (status != Status::kSuccess) {
+        return status;
+      }
+    }
+
+    state_.graph_created = true;
+
+    cudaGraphNode_t full_barrier_node;
+
+    // Create dummy stream
+    cudaStream_t stream;
+    status = detail::check_cuda_status(cudaStreamCreate(&stream));
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    // Create graph
+    status = detail::check_cuda_status(cudaGraphCreate(&state_.graph, 0));
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    // 1. Full barrier node
+    status = detail::check_cuda_status(cudaStreamBeginCaptureToGraph(
+          stream,
+          state_.graph,
+          nullptr, nullptr, 0,
+          cudaStreamCaptureModeRelaxed));
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    cutlass::Array<ElementFlag*, TP_> self_flag_ptrs;
+    for (int iteration = 0; iteration < TP_; ++iteration) {
+      self_flag_ptrs[iteration] = state_.params_array[iteration].distributed.self_flag_ptr_;
+    }
+
+    launch_full_barrier<TP_, ElementBarrier, TP_, ElementFlag>(
+        state_.device_barrier_ptrs, self_flag_ptrs, state_.device_idx, stream, launch_with_pdl);
+
+    status = detail::check_cuda_status(cudaStreamEndCapture(stream, &state_.graph));
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    size_t num_nodes;
+    status = detail::check_cuda_status(cudaGraphGetNodes(state_.graph, nullptr, &num_nodes));
+    if (status != Status::kSuccess) {
+      return status;
+    }
+    if (num_nodes != 1) {
+      CUTLASS_TRACE_HOST("  construct_graph() failure: expected a single node in the graph, got " << num_nodes << ".");
+      return Status::kErrorInternal;
+    }
+    if (status != Status::kSuccess) {
+      return status;
+    }
+    status = detail::check_cuda_status(cudaGraphGetNodes(state_.graph, &full_barrier_node, &num_nodes));
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    // 2. Optional mem copy branch
+    if constexpr (HasMemcpy) {
+
+      status = detail::check_cuda_status(cudaStreamBeginCaptureToGraph(
+            stream,
+            state_.graph,
+            &full_barrier_node,
+            /* dependencyData = */ nullptr,
+            1,
+            cudaStreamCaptureModeRelaxed));
+
+      if (status != Status::kSuccess) {
+        return status;
+      }
+
+      // No copies for first iter; we assume the data is already there.
+      for (int iteration = 1; iteration < TP_; ++iteration) {
+
+        status = detail::check_cuda_status(cudaMemcpyAsync(
+              state_.memcpy_source_ptr_array[iteration],
+              state_.memcpy_remote_ptr_array[iteration],
+              state_.memcpy_bytes[iteration],
+              cudaMemcpyDeviceToDevice, stream));
+
+        if (status != Status::kSuccess) {
+          return status;
+        }
+
+        // Set flag to non zero
+        status = detail::check_cuda_status(cudaMemsetAsync(
+              reinterpret_cast<void *>(state_.params_array[iteration].distributed.peer_flag_ptr_),
+              0b11111111,
+              sizeof(ElementFlag),
+              stream));
+
+        if (status != Status::kSuccess) {
+          return status;
+        }
+      }
+
+      status = detail::check_cuda_status(cudaStreamEndCapture(stream, &state_.graph));
+      if (status != Status::kSuccess) {
+        return status;
+      }
+    }
+
+    // 3. Run local GEMMs
+    // 3.1. Create edge between full barrier and the correct gemm stage/iteration
+    cudaGraphEdgeData barrier_to_gemm_edge = {};
+    barrier_to_gemm_edge.from_port = HasMemcpy ? cudaGraphKernelNodePortLaunchCompletion: cudaGraphKernelNodePortProgrammatic;
+    barrier_to_gemm_edge.type = cudaGraphDependencyTypeProgrammatic;
+
+    status = detail::check_cuda_status(cudaStreamBeginCaptureToGraph(
+          stream,
+          state_.graph,
+          &full_barrier_node,
+          /* dependencyData = */ &barrier_to_gemm_edge,
+          1,
+          cudaStreamCaptureModeRelaxed));
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    for (int iteration = 0; iteration < TP_; ++iteration) {
+      status = DeviceGemm::run(
+            state_.params_array[iteration],
+            stream,
+            /* cuda_adapter = */ nullptr,
+            /* launch_with_pdl = */ launch_with_pdl);
+
+      if (status != Status::kSuccess) {
+        return status;
+      }
+    }
+
+    status = detail::check_cuda_status(cudaStreamEndCapture(stream, &state_.graph));
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    // 4. Cleanup.
+    //// Destroy dummy stream
+    status = detail::check_cuda_status(cudaStreamDestroy(stream));
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    // 5. Instantiate graph
+    status = detail::check_cuda_status(cudaGraphInstantiate(
+          &state_.graph_executable,
+          state_.graph,
+          /* flags = */ 0));
+    if (status != Status::kSuccess) {
+      return status;
+    }
+    state_.graph_instantiated = true;
+
+    return Status::kSuccess;
+#else
+      CUTLASS_TRACE_HOST("  construct_graph() failure: target was compiled with an incompatible " <<
+          "version of the CUDA toolkit. Please compile Distributed GEMM with CUDA toolkit 12.4 or later.");
+      return Status::kErrorInternal;
+#endif
+  }
+
+  Status
+  update(Arguments const& args, void* workspace = nullptr) {
+    CUTLASS_TRACE_HOST("  DistributedGemm does not support updating arguments yet.");
+    return Status::kErrorInternal;
+  }
+
+  // NOTE: the interface for run() is different in Distributed Gemm:
+  //   1. launch_with_pdl is specified in `initialize`, where the cuda graph is being constructed,
+  //   2. the state of distributed gemm is an array of params for different iterations, and a
+  //      cuda graph.
+  //   3. Custom cuda adapters aren't supported for simplicity.
+  static Status
+  run(DistributedGemmState& state,
+      cudaStream_t stream = nullptr) {
+    CUTLASS_TRACE_HOST("DistributedGemm::run()");
+
+    if (not state.is_initialized) {
+      CUTLASS_TRACE_HOST("  Distributed gemm was not initialized. Did you forget to call initialize()?");
+      return Status::kErrorInternal;
+    }
+
+    if (not state.graph_instantiated) {
+      CUTLASS_TRACE_HOST("  Distributed gemm graph was not instantiated. Did you forget to call initialize()/construct_graph()?");
+      return Status::kErrorInternal;
+    }
+
+    cudaError_t result = cudaGraphLaunch(state.graph_executable, stream);
+    if (cudaSuccess != result) {
+      result = cudaGetLastError(); // to clear the error bit
+      CUTLASS_TRACE_HOST("  cudaGraphLaunch() returned error: " << cudaGetErrorString(result));
+      return Status::kErrorInternal;
+    }
+
+    return Status::kSuccess;
+  }
+
+  //
+  // Non-static launch overloads that first create and set the internal params struct of this kernel handle.
+  //
+
+  /// Overload that allows a user to re-launch the same kernel without updating internal params struct.
+  Status
+  run(
+    cudaStream_t stream = nullptr) {
+    return run(state_, stream);
+  }
+
+  /// Overload that allows a user to re-launch the same kernel without updating internal params struct.
+  Status
+  operator()(cudaStream_t stream = nullptr) {
+    return run(state_, stream);
+  }
+
+  /// Launches the kernel after first constructing Params internal state from supplied arguments.
+  Status
+  run(
+    Arguments const* args,
+    void** workspace_ptrs,
+    void** exclusive_workspace_ptrs,
+    int device_idx,
+    cudaStream_t stream = nullptr) {
+    Status status = initialize(
+        args,
+        workspace_ptrs,
+        exclusive_workspace_ptrs,
+        device_idx,
+        stream);
+
+    if (Status::kSuccess == status) {
+      status = run(stream);
+    }
+    return status;
+  }
+
+  /// Launches the kernel after first constructing Params internal state from supplied arguments.
+  Status
+  operator()(
+    Arguments const* args,
+    void** workspace_ptrs,
+    void** exclusive_workspace_ptrs,
+    int device_idx,
+    cudaStream_t stream = nullptr) {
+    return run(
+        args,
+        workspace_ptrs,
+        exclusive_workspace_ptrs,
+        device_idx,
+        stream);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace cutlass::distributed::device
+
+////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/experimental/distributed/device/full_barrier.hpp

@@ -0,0 +1,74 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Device layer interface for Distributed GEMM barrier kernel.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/experimental/distributed/kernel/full_barrier.hpp"
+
+namespace cutlass::distributed::device {
+
+template <int NP, typename IntType, int Iterations, typename FlagType>
+void launch_full_barrier(
+    cutlass::Array<IntType*, NP> device_arrival_ptrs,
+    cutlass::Array<FlagType*, Iterations> iteration_flag_ptrs,
+    IntType device_idx,
+    cudaStream_t stream,
+    bool launch_with_pdl) {
+
+#if (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 6))
+  // Legacy (kernel) launch with PDL
+  cudaLaunchAttribute attributes[1];
+  attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attributes[0].val.programmaticStreamSerializationAllowed = 1;
+
+  cudaLaunchConfig_t launch_config;
+  launch_config.gridDim = 1;
+  launch_config.blockDim = 1;
+  launch_config.dynamicSmemBytes = 0;
+  launch_config.stream = stream;
+  launch_config.attrs = attributes;
+  launch_config.numAttrs = launch_with_pdl ? 1 : 0;
+
+  cudaLaunchKernelEx(
+      &launch_config,
+      cutlass::distributed::kernel::full_barrier_kernel<NP, IntType, Iterations, FlagType>,
+      device_arrival_ptrs,
+      iteration_flag_ptrs,
+      device_idx);
+#endif
+}
+
+} // namespace cutlass::distributed::device
+

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/experimental/distributed/kernel/detail.hpp

@@ -0,0 +1,72 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Distributed gemm kernel layer helpers.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+
+///////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass::distributed::kernel::detail {
+
+// Ld with CV cache hint (don’t cache and fetch again)
+// Reference:
+// https://docs.nvidia.com/cuda/parallel-thread-execution/#cache-operators
+// Used for loading arrival counts from peer devices
+
+CUTLASS_DEVICE
+void ld_without_cache(uint64_t& val, void const * ptr) {
+  asm volatile(
+      "{\n"
+      "  ld.global.cv.u64 %0, [%1];\n"
+      "}\n"
+      : "=l"(val)
+      : "l"(ptr));
+}
+
+CUTLASS_DEVICE
+void ld_without_cache(uint32_t& val, void const * ptr) {
+  asm volatile(
+      "{\n"
+      "  ld.global.cv.u32 %0, [%1];\n"
+      "}\n"
+      : "=r"(val)
+      : "l"(ptr));
+}
+
+} // namespace cutlass::distributed::kernel::detail
+
+///////////////////////////////////////////////////////////////////////////////
+
+

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/experimental/distributed/kernel/dist_gemm_kernel_wrapper.hpp

@@ -0,0 +1,235 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*!
+  \file Distributed GEMM Kernel Wrapper
+
+  Prepends CUTLASS 3 GEMM kernels with barriers and other necessary instructions to exectue
+  a Distributed GEMM stage.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/arch/grid_dependency_control.h"
+#include "cutlass/gemm/gemm.h"
+
+#include "cutlass/experimental/distributed/kernel/detail.hpp"
+
+///////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass::distributed::kernel {
+
+namespace detail {
+
+// Allow all CUTLASS 3.X GEMM kernels
+template <typename GemmKernel_>
+struct SupportsDistributedGemm: cutlass::gemm::detail::IsCutlass3GemmKernel<GemmKernel_> {};
+
+} // namespace detail
+
+/*!
+  DistributedGemmKernelWrapper is a wrapper around a GEMM kernel.
+
+  Depending on the underlying distribution policy/schedule, it prepends the underlying local GEMM
+  kernel with a few additional instructions that gate the execution of the GEMM on buffers being
+  ready for stages/iterations > 0.
+*/
+
+template <class GemmKernel_, class DistSchedule_, class Enable = void>
+struct DistributedGemmKernelWrapper;
+
+template <class GemmKernel_, class DistSchedule_>
+struct DistributedGemmKernelWrapper<
+  GemmKernel_,
+  DistSchedule_,
+  cute::enable_if_t<detail::SupportsDistributedGemm<GemmKernel_>::value>
+  >: GemmKernel_
+{
+  using DistSchedule = DistSchedule_;
+  using TP = typename DistSchedule::TP;
+
+  static constexpr bool KernelWritesArrivalFlag = DistSchedule::KernelWritesArrivalFlag;
+
+  using BaseKernel = GemmKernel_;
+  using BaseArguments = typename BaseKernel::Arguments;
+  using BaseParams = typename BaseKernel::Params;
+
+  //static_assert(BaseKernel::ArchTag::kMinComputeCapability == 90, "DistGEMM only supports Hopper GEMMs for now.");
+  static_assert(not cute::is_same_v<typename BaseKernel::ElementC, void>, "DistributedGEMM epilogues must have a source.");
+
+  using ElementFlag = uint32_t;
+
+  // Device side arguments
+  struct DistributedArguments {
+    int device_idx = 0;
+    int iteration = 0;
+
+    void* self_flag_ptr{nullptr};
+    void* peer_flag_ptr{nullptr};
+  };
+
+  struct PackedArguments {
+    BaseArguments base{};
+    DistributedArguments distributed{};
+  };
+
+  struct DistributedParams {
+    int device_idx = 0;
+    int iteration = 0;
+
+    ElementFlag* self_flag_ptr_{nullptr};
+    ElementFlag* peer_flag_ptr_{nullptr};
+  };
+
+  // Kernel entry point API
+  struct PackedParams {
+    BaseParams base{};
+    DistributedParams distributed{};
+  };
+
+  using Params = PackedParams;
+
+  // Convert to underlying arguments. In this case, a simple copy for the aliased type.
+  static
+  PackedParams
+  to_underlying_arguments(PackedArguments const& args, void* workspace) {
+    CUTLASS_TRACE_HOST("distributed::to_underlying_arguments():");
+
+    auto kernel_params = BaseKernel::to_underlying_arguments(args.base, workspace);
+
+    DistributedParams dist_params = {
+        args.distributed.device_idx,
+        args.distributed.iteration,
+        reinterpret_cast<ElementFlag*>(args.distributed.self_flag_ptr),
+        reinterpret_cast<ElementFlag*>(args.distributed.peer_flag_ptr)
+    };
+
+    return {kernel_params, dist_params};
+  }
+
+  static bool
+  can_implement(BaseArguments const& args) {
+    return BaseKernel::can_implement(args);
+  }
+
+  static bool
+  can_implement(PackedArguments const& args) {
+    return BaseKernel::can_implement(args.base);
+  }
+
+  static size_t
+  get_workspace_size(BaseArguments const& args) {
+    return BaseKernel::get_workspace_size(args);
+  }
+
+  static size_t
+  get_workspace_size(PackedArguments const& args) {
+    return BaseKernel::get_workspace_size(args.base);
+  }
+
+  static cutlass::Status
+  initialize_workspace(BaseArguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
+    CudaHostAdapter* cuda_adapter = nullptr) {
+    return BaseKernel::initialize_workspace(args, workspace, stream, cuda_adapter);
+  }
+
+  static cutlass::Status
+  initialize_workspace(PackedArguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
+    CudaHostAdapter* cuda_adapter = nullptr) {
+    return BaseKernel::initialize_workspace(args.base, workspace, stream, cuda_adapter);
+  }
+
+  /// Computes the grid shape
+  static dim3
+  get_grid_shape(PackedParams const& params) {
+    return BaseKernel::get_grid_shape(params.base);
+  }
+  
+  static dim3
+  get_grid_shape(BaseParams const& params) {
+    return BaseKernel::get_grid_shape(params);
+  }
+
+  CUTLASS_DEVICE
+  void
+  barrier_buffer(PackedParams const& params) {
+    if (params.distributed.iteration > 0) {
+
+      ElementFlag comm_iter = 0;
+      detail::ld_without_cache(comm_iter, params.distributed.self_flag_ptr_);
+      while (comm_iter == 0) {
+        detail::ld_without_cache(comm_iter, params.distributed.self_flag_ptr_);
+        __nanosleep(40);
+      }
+
+    }
+  }
+
+  CUTLASS_DEVICE
+  void
+  maybe_signal_arrival(PackedParams const& params) {
+    if constexpr (KernelWritesArrivalFlag) {
+      if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 &&
+          threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0 &&
+          params.distributed.iteration > 0) {
+        *reinterpret_cast<ElementFlag*>(params.distributed.peer_flag_ptr_) = 1;
+      }
+    }
+  }
+
+  CUTLASS_DEVICE
+  void
+  operator()(PackedParams const& params, char* smem_buf) {
+    // Launch next grid as soon as possible
+    arch::launch_dependent_grids();
+
+    // Wait on previous kernels to flush their memory.
+    arch::wait_on_dependent_grids();
+
+    // Optionally write arrivals for the previous stage/iteration.
+    maybe_signal_arrival(params);
+
+    // Spin-wait on an arrival flag, make sure the respective buffers are ready.
+    // If the buffered operand is memcpied into, it would wait on its local flag.
+    // If it's a remote buffer that is accessed directly, it would wait on its remote flag.
+    barrier_buffer(params);
+
+    // Perform local gemm
+    BaseKernel gemm;
+    gemm(params.base, smem_buf);
+  }
+
+};
+
+} // namespace cutlass::distributed::kernel
+
+///////////////////////////////////////////////////////////////////////////////
+

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/experimental/distributed/kernel/full_barrier.hpp

@@ -0,0 +1,82 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Distributed GEMM barrier kernel.
+
+    The kernel resets the per-stage arrival flags, performs a full barrier (any-to-any),
+    and also atomically resets the local barrier arrival count.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/arch/grid_dependency_control.h"
+
+#include "cutlass/experimental/distributed/kernel/detail.hpp"
+
+namespace cutlass::distributed::kernel {
+
+template <int NP, typename IntType, int Iterations, typename FlagType>
+__global__ void full_barrier_kernel(
+    cutlass::Array<IntType*, NP> device_arrival_ptrs,
+    cutlass::Array<FlagType*, Iterations> iteration_flag_ptrs,
+    IntType device_idx) {
+
+  arch::launch_dependent_grids();
+  arch::wait_on_dependent_grids();
+
+  CUTLASS_PRAGMA_UNROLL
+  for (FlagType i = 0; i < Iterations; ++i) {
+    iteration_flag_ptrs[i][0] = static_cast<FlagType>(0);
+  }
+
+  IntType val = 1;
+  IntType max_val = static_cast<IntType>(NP - 1);
+
+  CUTLASS_PRAGMA_UNROLL
+  for (IntType d = 0; d < NP; ++d) {
+    if (d != device_idx) {
+      atomicAdd(device_arrival_ptrs[d], val);
+    }
+  }
+
+  IntType curr_val = 0;
+  detail::ld_without_cache(curr_val, device_arrival_ptrs[device_idx]);
+  while (curr_val < max_val) {
+    __nanosleep(40);
+    detail::ld_without_cache(curr_val, device_arrival_ptrs[device_idx]);
+  }
+
+  atomicSub(device_arrival_ptrs[device_idx], max_val);
+}
+
+} // namespace cutlass::distributed::kernel
+

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/experimental/distributed/schedules/dist_gemm_1d_schedules.hpp

@@ -0,0 +1,324 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*!
+  \file 1-D Distributed GEMM Schedules
+
+  NOTE: This API is __experimental__ and will change heavily over time. Particularly the use of
+  CuTe layouts as integer functions in defining iteration-to-tile mappings is over-expressive and
+  leaves plenty of room for incorrect/unexpected behavior.
+  Please proceed with caution when modifying these schedules or defining new ones.
+
+  Device/iteration mappings are defined with CuTe layouts, 
+  since they are functions from integers to integers as well.
+  
+  Each mapping is defined as a linear function of 2 variables (rank-2 layout):
+   First variable (mode) is device index, second variable (mode) is iteration.
+   A constant is also added to the final result as an offset value. This is a temporary workaround
+   so that identity ownership mappings in the final iteration can be guaranteed for the schedules
+   currently implemented.
+  How are these mappings defined?
+    Each schedule represents a unique parallel matrix multiplication algorithm, which describes how
+    matrices/tensors are distributed among TP GPUs.
+
+    Depending on the algorithm, access patterns (GPU to tile or (GPU, iteration) to tile) mappings)
+    are not necessarily going to be the identity function.
+
+  Pitfalls:
+    The current representation uses CuTe layouts as arbitrary linear functions that map
+    (GPU, iteration) to tile indices.
+    This approach is over-expressive, and therefore makes a lot of assumptions on the part of the
+    developer in how these mappings are defined. This can easily lead to incorrect implementations
+    if not handled carefully.
+
+  
+  Assumption made in all schedules: TP == number of iterations (stages)
+*/
+
+#pragma once
+
+#include "cute/layout.hpp"
+#include "cute/tensor.hpp"
+#include "cutlass/cutlass.h"
+
+#include "cutlass/experimental/distributed/schedules/dist_gemm_base_schedule.hpp"
+
+///////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass::distributed::schedules {
+
+// GEMM + Reduce Scatter
+// A and B are tiled along the K mode, which means each GPU gets an [M, K / TP]-shaped slice of A,
+// and an [N, K / TP] slice of B.
+// A is further tiled along the M mode, so that each stage/iteration computes a GEMM of shape
+// [M / TP, N, K / TP], and the epilogue will perform the reduction by reading its C tensor directly
+// from the left peer's previous D buffer.
+//
+// Below is an illustration of the tiling and iteration mappings for this pattern in the TP=4 case:
+//
+//   Rows correspond to the M mode, columns correspond to the K mode for A and B and N mode for 
+//   C and D.  Because sharding is done along K, each column of tiles is owned by one GPU.
+//   Values in the grid correspond to the iteration/stage accessing the tile.
+//   * means the same tile is accessed in all iterations/stages.
+//
+//         Tensor A                             Tensor B              
+//                                                                    
+//  GPU0  GPU1  GPU2  GPU3              GPU0  GPU1  GPU2  GPU3        
+// |-----|-----|-----|-----|           |-----|-----|-----|-----|      
+// |     |     |     |     |           |     |     |     |     |      
+// |  3  |  0  |  1  |  2  |           |     |     |     |     |      
+// |_____|_____|_____|_____|           |     |     |     |     |      
+// |     |     |     |     |           |     |     |     |     |      
+// |  2  |  3  |  0  |  1  |           |     |     |     |     |      
+// |_____|_____|_____|_____|           |  *  |  *  |  *  |  *  |      
+// |     |     |     |     |           |     |     |     |     |      
+// |  1  |  2  |  3  |  0  |           |     |     |     |     |      
+// |_____|_____|_____|_____|           |     |     |     |     |      
+// |     |     |     |     |           |     |     |     |     |      
+// |  0  |  1  |  2  |  3  |           |     |     |     |     |      
+// |_____|_____|_____|_____|           |_____|_____|_____|_____|      
+//                                                                    
+//                          M x K                               N x K 
+//
+//
+//              Tensor C                            Tensor D              
+//              (Peer's D)
+//                                         
+//                                                                        
+//      |-----------------------|           |-----------------------|     
+//      |                       |           |                       |     
+// GPU0 |         1,2,3         |      GPU0 |           *           |     
+//      |_______________________|           |_______________________|     
+//      |                       |           |                       |     
+// GPU1 |         1,2,3         |      GPU1 |           *           |     
+//      |_______________________|           |_______________________|     
+//      |                       |           |                       |     
+// GPU2 |         1,2,3         |      GPU2 |           *           |     
+//      |_______________________|           |_______________________|     
+//      |                       |           |                       |     
+// GPU3 |         1,2,3         |      GPU3 |           *           |     
+//      |_______________________|           |_______________________|     
+//                                                                        
+//                               M x N                               M x N
+//
+//
+//  Tensor A's access pattern can be expressed as follows as a function of GPU index and iteration:
+//    tile_idx = ((device_idx - 1) - iter + TP) % TP
+//  
+//  and can be expressed with the following CuTe layout:
+//    (TP, TP) : (1, -1)
+//  with ProcessorOffset = -1
+//
+//
+//  Note: Since this schedule does not expose any communication, iteration 0 has no reduction step,
+//  therefore epilogue is sourceless in iteration 0, and in the rest of the iterations the epilogue
+//  source is a remote pointer to Tensor D owned by its left peer.
+//
+//  Left peer is simply (device_idx - 1 + TP) % TP, which is expressed with the following CuTe layout:
+//    (TP, TP) : (1, 0)
+//
+template <class TP_>
+struct ReduceScatter1D_TilingA_RotatingC: BaseSchedule<
+    TP_,
+    /* ProcessorTiler_ = */ cute::Shape<_1, _1, TP_, _1>,
+    /* IterationTiler_ = */ cute::Shape<TP_, _1, _1, _1>,
+    /* PeerDeviceMapping_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _0>>,                             // (left neighbor) = (device_idx + ProcessorOffset + TP) % TP, with ProcessorOffset = -1
+    /* IterationMappingM_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _m1>>,                            // = (device_idx + ProcessorOffset - iter + TP) % TP, with ProcessorOffset = -1
+    /* IterationMappingN_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>,                             // (IterationTiler::N == 1) = 0
+    /* IterationMappingK_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>,                             // (IterationTiler::K == 1) = 0
+    /* IterationMappingL_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>,                             // (IterationTiler::L == 1) = 0
+    /* ProcessorOffset_ = */ _m1,
+    /* MemcpyA_ = */ false,
+    /* MemcpyB_ = */ false,
+    /* KernelWritesArrivalFlag_ = */ true,
+    /* NumBuffersA_ = */ 0,
+    /* NumBuffersB_ = */ 0,
+    /* NumBuffersC_ = */ 0,
+    /* NumBuffersD_  = */ TP_{} - 1> {};
+
+// This schedule is similar to ReduceScatter1D_TilingA_RotatingC, but with the second tiling
+// done along N instead of M. All other details remain unchanged.
+template <class TP_>
+struct ReduceScatter1D_TilingB_RotatingC: BaseSchedule<
+    TP_,
+    /* ProcessorTiler_ = */ cute::Shape<_1, _1, TP_, _1>,
+    /* IterationTiler_ = */ cute::Shape<_1, TP_, _1, _1>,
+    /* PeerDeviceMapping_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _0>>,                             // (left neighbor) = (device_idx + ProcessorOffset + TP) % TP, with ProcessorOffset = -1
+    /* IterationMappingM_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>,                             // (IterationTiler::N == 1) = 0
+    /* IterationMappingN_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _m1>>,                            // = (device_idx + ProcessorOffset - iter + TP) % TP, with ProcessorOffset = -1
+    /* IterationMappingK_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>,                             // (IterationTiler::K == 1) = 0
+    /* IterationMappingL_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>,                             // (IterationTiler::L == 1) = 0
+    /* ProcessorOffset_ = */ _m1,
+    /* MemcpyA_ = */ false,
+    /* MemcpyB_ = */ false,
+    /* KernelWritesArrivalFlag_ = */ true,
+    /* NumBuffersA_ = */ 0,
+    /* NumBuffersB_ = */ 0,
+    /* NumBuffersC_ = */ 0,
+    /* NumBuffersD_  = */ TP_{} - 1> {};
+
+
+// AllGather + GEMM
+// A and B are tiled along the N mode, which means each GPU allgathers A,
+// and operates with an [N / TP, K] slice of B.
+// For pipelining, A is further tiled along the M mode, so that each stage/iteration computes a
+// GEMM of shape [M / TP, N / TP, K], and concurrently we copy a peer's A slice into a local buffer
+// for the next stage/iteration.
+//
+// Below is an illustration of the tiling and iteration mappings for this pattern in the TP=4 case:
+//
+//   Rows correspond to the M mode, columns correspond to the K mode for A and B and N mode for 
+//   C and D.
+//
+//   Since this is a pipelined schedule without exposed communication, the first iteration starts
+//   off immediately and operates on local slices of A and B. In the rest of the iterations, each
+//   GPU accesses a slice of A copied from a peer GPU while it was busy with the last stage.
+//
+//   Values in the following grids correspond to the peer buffer accessed by each GPU during
+//   different iterations:
+//
+//              Tensor A                         Tensor A               
+//               iter 0                           iter 1                
+//                                                                      
+//      |-----------------------|        |-----------------------|      
+//      |                       |        |                       |      
+// GPU0 |           0           |        |           1           |      
+//      |_______________________|        |_______________________|      
+//      |                       |        |                       |      
+// GPU1 |           1           |        |           2           |      
+//      |_______________________|        |_______________________|      
+//      |                       |        |                       |      
+// GPU2 |           2           |        |           3           |      
+//      |_______________________|        |_______________________|      
+//      |                       |        |                       |      
+// GPU3 |           3           |        |           0           |      
+//      |_______________________|        |_______________________|      
+//                                                                      
+//                               M x K                            M x K 
+//
+//              Tensor A                         Tensor A               
+//               iter 2                           iter 3                
+//                                                                      
+//      |-----------------------|        |-----------------------|      
+//      |                       |        |                       |      
+// GPU0 |           2           |        |           3           |      
+//      |_______________________|        |_______________________|      
+//      |                       |        |                       |      
+// GPU1 |           3           |        |           0           |      
+//      |_______________________|        |_______________________|      
+//      |                       |        |                       |      
+// GPU2 |           0           |        |           1           |      
+//      |_______________________|        |_______________________|      
+//      |                       |        |                       |      
+// GPU3 |           1           |        |           2           |      
+//      |_______________________|        |_______________________|      
+//                                                                      
+//                               M x K                            M x K 
+//
+//   Values in the following grids correspond to the tile accessed during each iteration.
+//   * means the same tile is accessed in all iterations/stages.
+//
+//              Tensor B                             Tensor C/D               
+//                                                                          
+//                                                                          
+//      |-----------------------|            |-----|-----|-----|-----|      
+//      |                       |            |     |     |     |     |      
+// GPU0 |           *           |       GPU0 |  0  |  1  |  2  |  3  |      
+//      |_______________________|            |_____|_____|_____|_____|      
+//      |                       |            |     |     |     |     |      
+// GPU1 |           *           |       GPU1 |  3  |  0  |  1  |  2  |      
+//      |_______________________|            |_____|_____|_____|_____|      
+//      |                       |            |     |     |     |     |      
+// GPU2 |           *           |       GPU2 |  2  |  3  |  0  |  1  |      
+//      |_______________________|            |_____|_____|_____|_____|      
+//      |                       |            |     |     |     |     |      
+// GPU3 |           *           |       GPU3 |  1  |  2  |  3  |  0  |      
+//      |_______________________|            |_____|_____|_____|_____|      
+//                                                                          
+//                               N x K                                M x N 
+//
+//
+//  Tensor C/D's access pattern can be expressed as follows as a function of GPU index and iteration:
+//    tile_idx = (device_idx + iter) % TP
+//  
+//  and can be expressed with the following CuTe layout:
+//    (TP, TP) : (1, 1)
+//
+//  This schedule does not need a ProcessorOffset constant.
+//
+//  Peer devices from which A slices are copied is also expressed with the same function and CuTe
+//  layout.
+//
+template <class TP_>
+struct AllGather1D_TilingCD_RotatingA: BaseSchedule<
+    TP_,
+    /* ProcessorTiler_ = */ cute::Shape<_1, TP_, _1, _1>,
+    /* IterationTiler_ = */ cute::Shape<TP_, _1, _1, _1>,
+    /* PeerDeviceMapping_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _1>>,                             // = device_idx + iter
+    /* IterationMappingM_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _1>>,                             // = device_idx + iter
+    /* IterationMappingN_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>,                             // (IterationTiler::N == 1) = 0
+    /* IterationMappingK_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>,                             // (IterationTiler::K == 1) = 0
+    /* IterationMappingL_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>,                             // (IterationTiler::L == 1) = 0
+    /* ProcessorOffset_ = */ _0,
+    /* MemcpyA_ = */ true,
+    /* MemcpyB_ = */ false,
+    /* KernelWritesArrivalFlag_ = */ false,
+    /* NumBuffersA_ = */ TP_{} - 1,
+    /* NumBuffersB_ = */ 0,
+    /* NumBuffersC_ = */ 0,
+    /* NumBuffersD_ = */ 0>{};
+
+// This schedule is similar to AllGather1D_TilingCD_RotatingA, but with the order of tiling
+// swapped from N then M to M then N. This means slices of B are rotated around GPUs instead of
+// slices of A. All other details remain unchanged.
+template <class TP_>
+struct AllGather1D_TilingCD_RotatingB: BaseSchedule<
+    TP_,
+    /* ProcessorTiler_ = */ cute::Shape<TP_, _1, _1, _1>,
+    /* IterationTiler_ = */ cute::Shape<_1, TP_, _1, _1>,
+    /* PeerDeviceMapping_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _1>>,                             // = device_idx + iter
+    /* IterationMappingM_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>,                             // (IterationTiler::M == 1) = 0
+    /* IterationMappingN_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _1>>,                             // = device_idx + iter
+    /* IterationMappingK_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>,                             // (IterationTiler::K == 1) = 0
+    /* IterationMappingL_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>,                             // (IterationTiler::L == 1) = 0
+    /* ProcessorOffset_ = */ _0,
+    /* MemcpyA_ = */ false,
+    /* MemcpyB_ = */ true,
+    /* KernelWritesArrivalFlag_ = */ false,
+    /* NumBuffersA_ = */ 0,
+    /* NumBuffersB_ = */ TP_{} - 1,
+    /* NumBuffersC_ = */ 0,
+    /* NumBuffersD_ = */ 0>{};
+
+
+} // namespace cutlass::distributed::schedules
+
+///////////////////////////////////////////////////////////////////////////////
+

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/experimental/distributed/schedules/dist_gemm_base_schedule.hpp

@@ -0,0 +1,538 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*!
+  \file Base Schedule for Distributed GEMM
+
+  Templates Distributed GEMM schedules so that they can be expressed as a set of CuTe primitives and
+  other static values.
+
+  NOTE: This API is __experimental__ and will change heavily over time. Particularly the use of
+  CuTe layouts as integer functions in defining iteration-to-tile mappings is over-expressive and
+  leaves plenty of room for incorrect/unexpected behavior.
+  Please proceed with caution when modifying these schedules or defining new ones.
+*/
+
+#pragma once
+
+#include "cute/layout.hpp"
+#include "cute/tensor.hpp"
+#include "cutlass/cutlass.h"
+
+
+///////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass::distributed::schedules {
+
+/*
+ * Distributed GEMM schedules define exactly how operand tensors are tiled and sliced across 
+ * processors (GPUs) and stages/iterations.
+ *
+ * BaseSchedule's role is to ease the implementation of arbitrary Distributed GEMM schedules
+ * and reduce code repetition, simply by reducing the implementation to CuTe primitives and a few
+ * other static values (buffer sizes, whether tensors are rotated using memcpies or not, and the
+ * like.)
+ */
+template <
+  class TP_,                      // CuTe constant defining the number of processors / GPUs / TP value
+  class ProcessorTiler_,          // CuTe tiler defining how fully materialized tensors are sharded across devices
+  class IterationTiler_,          // CuTe tiler defining how local tensors are tiled across stages/iterations
+  class PeerDeviceMapping_,       // CuTe layout mapping device index and stage/iteration to the device's peer index for that stage/iteration
+  class IterationMappingM_,       // CuTe layout mapping device index and stage/iteration to M tile index
+  class IterationMappingN_,       // CuTe layout mapping device index and stage/iteration to N tile index
+  class IterationMappingK_,       // CuTe layout mapping device index and stage/iteration to K tile index
+  class IterationMappingL_,       // CuTe layout mapping device index and stage/iteration to L tile index
+  class ProcessorOffset_,         // Constant offset for processor / GPU index in iteration mapping
+  bool MemcpyA_,                  // Whether tensor A is memcpied
+  bool MemcpyB_,                  // Whether tensor B is memcpied
+  bool KernelWritesArrivalFlag_,  // Whether the kernel writes arrival flags (when tensors are directly accessed from peer and not memcpied)
+  int NumBuffersA_,               // Number of buffers required for tensor A
+  int NumBuffersB_,               // Number of buffers required for tensor B
+  int NumBuffersC_,               // Number of buffers required for tensor C
+  int NumBuffersD_>               // Number of buffers required for tensor D
+struct BaseSchedule {
+
+  using TP = TP_;
+
+  static_assert(
+      cute::is_static<TP>::value && cute::is_integral<TP>::value && cute::rank(TP{}) == 1 && cute::depth(TP{}) == 0,
+      "Only integers allowed for TP at this time.");
+
+  static_assert(cute::rank(ProcessorTiler_{}) == 4, "Expected rank-4 processor tiler.");
+  static_assert(cute::rank(IterationTiler_{}) == 4, "Expected rank-4 iteration tiler.");
+
+  static_assert(cute::rank(PeerDeviceMapping_{}) == 2, 
+      "PeerDeviceMapping must be rank-2 (device_idx, iter)");
+
+  static_assert(cute::rank(IterationMappingM_{}) == 2, 
+      "IterationMappingM must be rank-2 (device_idx, iter).");
+  static_assert(cute::rank(IterationMappingN_{}) == 2, 
+      "IterationMappingN must be rank-2 (device_idx, iter).");
+  static_assert(cute::rank(IterationMappingK_{}) == 2, 
+      "IterationMappingK must be rank-2 (device_idx, iter).");
+  static_assert(cute::rank(IterationMappingL_{}) == 2, 
+      "IterationMappingL must be rank-2 (device_idx, iter).");
+
+  using ProcessorTiler = ProcessorTiler_;
+  using IterationTiler = IterationTiler_;
+
+  using PeerDeviceMapping = PeerDeviceMapping_;
+  using IterationMappingM = IterationMappingM_;
+  using IterationMappingN = IterationMappingN_;
+  using IterationMappingK = IterationMappingK_;
+  using IterationMappingL = IterationMappingL_;
+
+  using ProcessorOffset = ProcessorOffset_;
+
+  static constexpr bool KernelWritesArrivalFlag = KernelWritesArrivalFlag_;
+  static constexpr bool MemcpyA = MemcpyA_;
+  static constexpr bool MemcpyB = MemcpyB_;
+  static constexpr bool HasMemcpy = MemcpyA || MemcpyB;
+
+  static constexpr int NumBuffersA = NumBuffersA_;
+  static constexpr int NumBuffersB = NumBuffersB_;
+  static constexpr int NumBuffersC = NumBuffersC_;
+  static constexpr int NumBuffersD = NumBuffersD_;
+
+  static_assert(
+      NumBuffersA > 0 ^ 
+      NumBuffersB > 0 ^ 
+      NumBuffersC > 0 ^ 
+      NumBuffersD > 0,
+      "Only one of the ABCD tensors can be buffered!");
+
+  static constexpr bool BufferedOutput = NumBuffersC > 0 || NumBuffersD > 0;
+  static constexpr bool RemoteC = NumBuffersC == 0 && NumBuffersD > 0;
+  static constexpr bool RemoteD = NumBuffersD == 0 && NumBuffersC > 0;
+
+  static_assert(not RemoteD, "Remote D is not supported yet.");
+
+  // Host-side API: can_implement based on the GLOBAL problem shape
+  template <typename ProblemShape>
+  static bool
+  can_implement_global(ProblemShape const& global_problem_shape) {
+    auto [M, N, K, L] = append<4>(global_problem_shape, 1);
+
+    auto [ptileM, ptileN, ptileK, ptileL] = ProcessorTiler{};
+    auto [itileM, itileN, itileK, itileL] = IterationTiler{};
+
+    auto tileM = ptileM * itileM;
+    auto tileN = ptileN * itileN;
+    auto tileK = ptileK * itileK;
+    auto tileL = ptileL * itileL;
+
+    return M % tileM == 0 && N % tileN == 0 && K % tileK == 0 && L % tileL == 0;
+  }
+
+  template <typename ProblemShape>
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_local_gemm_shape(ProblemShape const& global_problem_shape) {
+    auto problem_shape_MNKL = append<4>(global_problem_shape, 1);
+
+    return shape_div(
+        shape_div(
+          problem_shape_MNKL,
+          ProcessorTiler{}),
+        IterationTiler{});
+  }
+
+  // Host-side API: determine peers
+  static auto
+  get_peers_for_device(int device_idx) {
+    auto left_peer_id = device_idx > 0 ? device_idx - 1 : TP{} - 1;
+    auto right_peer_id = device_idx < TP{} - 1 ? device_idx + 1 : 0;
+
+    return cute::make_tuple(left_peer_id, right_peer_id);
+  }
+
+  // Determines peer given device index and iteration
+  static int
+  get_remote_peer_id(int device_idx, int iteration) {
+    auto device_iter_to_peer_idx = PeerDeviceMapping{};
+    auto peer_idx = (
+      device_iter_to_peer_idx(device_idx + ProcessorOffset{}, iteration) + TP{}
+    ) % TP{};
+    return peer_idx;
+  }
+
+  // Construct tilers and index mappers for sharding across processors
+  template <typename Tensor>
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_processor_tiler_a(Tensor tensor) {
+    if constexpr (NumBuffersA > 0) {
+      return shape_div(tensor.shape(), select<0,2,3>(IterationTiler{}));
+    } else {
+      return shape_div(tensor.shape(), select<0,2,3>(ProcessorTiler{}));
+    }
+  }
+
+  template <typename Tensor>
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_processor_tiler_b(Tensor tensor) {
+    if constexpr (NumBuffersB > 0) {
+      return shape_div(tensor.shape(), select<1,2,3>(IterationTiler{}));
+    } else {
+      return shape_div(tensor.shape(), select<1,2,3>(ProcessorTiler{}));
+    }
+  }
+
+  template <typename Tensor>
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_processor_tiler_c(Tensor tensor) {
+    if constexpr (BufferedOutput) {
+      return shape_div(tensor.shape(), select<0,1,3>(IterationTiler{}));
+    } else {
+      return shape_div(tensor.shape(), select<0,1,3>(ProcessorTiler{}));
+    }
+  }
+
+  template <typename Tensor>
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_processor_tiler_d(Tensor tensor) {
+    return get_processor_tiler_c(tensor);
+  }
+
+  // Construct tilers and index mappers for tiling and iterating on device
+  template <typename Tensor>
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_device_tiler_a(Tensor tensor) {
+    static_assert(NumBuffersA == 0, "Buffered tensors don't have device tilers!");
+    return shape_div(tensor.shape(), select<0,2,3>(IterationTiler{}));
+  }
+
+  template <typename Tensor>
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_device_tiler_b(Tensor tensor) {
+    static_assert(NumBuffersB == 0, "Buffered tensors don't have device tilers!");
+    return shape_div(tensor.shape(), select<1,2,3>(IterationTiler{}));
+  }
+
+  template <typename Tensor>
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_device_tiler_c(Tensor tensor) {
+    static_assert(NumBuffersC == 0 && NumBuffersD == 0, "Buffered tensors don't have device tilers!");
+    return shape_div(tensor.shape(), select<0,1,3>(IterationTiler{}));
+  }
+
+  template <typename Tensor>
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_device_tiler_d(Tensor tensor) {
+    static_assert(NumBuffersC == 0 && NumBuffersD == 0, "Buffered tensors don't have device tilers!");
+    return shape_div(tensor.shape(), select<0,1,3>(IterationTiler{}));
+  }
+
+  // Map device index and iteration to tile coordinate
+  // Must be implemented by children for now.
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_device_tile_idx_a(int device_idx, int iteration) {
+    auto mapping_m = IterationMappingM{};
+    auto mapping_k = IterationMappingK{};
+    auto mapping_l = IterationMappingL{};
+    auto crd_m = (mapping_m(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
+    auto crd_k = (mapping_k(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
+    auto crd_l = (mapping_l(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
+    return make_coord(crd_m, crd_k, crd_l);
+  }
+
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_device_tile_idx_b(int device_idx, int iteration) {
+    auto mapping_n = IterationMappingN{};
+    auto mapping_k = IterationMappingK{};
+    auto mapping_l = IterationMappingL{};
+    auto crd_n = (mapping_n(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
+    auto crd_k = (mapping_k(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
+    auto crd_l = (mapping_l(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
+    return make_coord(crd_n, crd_k, crd_l);
+  }
+
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_device_tile_idx_c(int device_idx, int iteration) {
+    auto mapping_m = IterationMappingM{};
+    auto mapping_n = IterationMappingN{};
+    auto mapping_l = IterationMappingL{};
+    auto crd_m = (mapping_m(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
+    auto crd_n = (mapping_n(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
+    auto crd_l = (mapping_l(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
+    return make_coord(crd_m, crd_n, crd_l);
+  }
+
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_device_tile_idx_d(int device_idx, int iteration) {
+    auto mapping_m = IterationMappingM{};
+    auto mapping_n = IterationMappingN{};
+    auto mapping_l = IterationMappingL{};
+    auto crd_m = (mapping_m(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
+    auto crd_n = (mapping_n(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
+    auto crd_l = (mapping_l(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
+    return make_coord(crd_m, crd_n, crd_l);
+  }
+
+  // Device Partitioners: partition non-buffered processor-resident operands.
+  // Processor-resident operands fall into two categories: buffered, and not buffered.
+  // Those buffered aren't expected to be further partitioned, and those 
+  template <typename Tensor>
+  static auto
+  get_tensor_A(Tensor original_tensor, void * tensor_buffer_ptr, int device_idx, int iteration) {
+    static_assert(rank(original_tensor) == 3);
+
+    using Element = typename Tensor::value_type;
+    // Recreate tensor without constness. This is to ensure return types match.
+    Element* ptr = const_cast<Element*>(original_tensor.data());
+    auto shape = original_tensor.shape();
+    auto layout = original_tensor.layout();
+    auto tensor = make_tensor(ptr, layout);
+
+    if constexpr (NumBuffersA  == 0) {
+      auto tiler = get_device_tiler_a(tensor);
+      auto idx = get_device_tile_idx_a(device_idx, iteration);
+      return inner_partition(tensor, tiler, idx);
+    } else {
+      Element* ptr_buffer = reinterpret_cast<Element*>(tensor_buffer_ptr);
+      if (iteration == 0) {
+        return tensor;
+      }
+      ptr_buffer += size(shape) * (iteration - 1);
+
+      return make_tensor(ptr_buffer, layout);
+    }
+  }
+
+  template <typename Tensor>
+  static auto
+  get_tensor_B(Tensor original_tensor, void * tensor_buffer_ptr, int device_idx, int iteration) {
+    static_assert(rank(original_tensor) == 3);
+
+    using Element = typename Tensor::value_type;
+    // Recreate tensor without constness. This is to ensure return types match.
+    Element * ptr = const_cast<Element *>(original_tensor.data());
+    auto shape = original_tensor.shape();
+    auto layout = original_tensor.layout();
+    auto tensor = make_tensor(ptr, layout);
+
+    if constexpr (NumBuffersB  == 0) {
+      auto tiler = get_device_tiler_b(tensor);
+      auto idx = get_device_tile_idx_b(device_idx, iteration);
+      return inner_partition(tensor, tiler, idx);
+    } else {
+      Element * ptr_buffer = reinterpret_cast<Element *>(tensor_buffer_ptr);
+      if (iteration == 0) {
+        return tensor;
+      }
+      ptr_buffer += size(shape) * (iteration - 1);
+
+      return make_tensor(ptr_buffer, layout);
+    }
+  }
+
+  template <typename Tensor>
+  static auto
+  get_tensor_C(Tensor original_tensor, void * tensor_buffer_ptr, int device_idx, int iteration) {
+    static_assert(rank(original_tensor) == 3);
+
+    using Element = typename Tensor::value_type;
+    // Recreate tensor without constness. This is to ensure return types match.
+    Element * ptr = const_cast<Element *>(original_tensor.data());
+    auto shape = original_tensor.shape();
+    auto layout = original_tensor.layout();
+    auto tensor = make_tensor(ptr, layout);
+
+    if constexpr (not BufferedOutput) {
+      auto tiler = get_device_tiler_c(tensor);
+      auto idx = get_device_tile_idx_c(device_idx, iteration);
+      return inner_partition(tensor, tiler, idx);
+    } else {
+      // implement Remote D
+      static_assert(RemoteC, "");
+
+      Element * ptr_buffer = reinterpret_cast<Element *>(tensor_buffer_ptr);
+      if (iteration == 0) {
+        return tensor;
+      }
+      ptr_buffer += size(shape) * (iteration - 1);
+
+      return make_tensor(ptr_buffer, layout);
+    }
+  }
+
+  template <typename Tensor>
+  static auto
+  get_tensor_D(Tensor original_tensor, void * tensor_buffer_ptr, int device_idx, int iteration) {
+    static_assert(rank(original_tensor) == 3);
+
+    using Element = typename Tensor::value_type;
+    // Recreate tensor without constness. This is to ensure return types match.
+    Element * ptr = const_cast<Element *>(original_tensor.data());
+    auto shape = original_tensor.shape();
+    auto layout = original_tensor.layout();
+    auto tensor = make_tensor(ptr, layout);
+
+    if constexpr (not BufferedOutput) {
+      auto tiler = get_device_tiler_d(tensor);
+      auto idx = get_device_tile_idx_d(device_idx, iteration);
+      return inner_partition(tensor, tiler, idx);
+    } else {
+      // implement Remote D
+      static_assert(RemoteC, "");
+
+      Element * ptr_buffer = reinterpret_cast<Element *>(tensor_buffer_ptr);
+      // last iteration is the local tensor, the rest are buffers
+      if (iteration == TP{} - 1) {
+        return tensor;
+      }
+      ptr_buffer += size(shape) * iteration; // note: iteration, not iteration - 1
+
+      return make_tensor(ptr_buffer, layout);
+    }
+  }
+
+  template <typename ProblemShape>
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_local_a_shape(ProblemShape problem_shape) {
+    auto problem_shape_MNKL = append<4>(problem_shape, 1);
+    if constexpr (NumBuffersA == 0) {
+      return shape_div(
+            select<0,2,3>(problem_shape_MNKL),
+            select<0,2,3>(ProcessorTiler{}));
+    } else {
+      return shape_div(
+          shape_div(
+            select<0,2,3>(problem_shape_MNKL),
+            select<0,2,3>(ProcessorTiler{})),
+          select<0,2,3>(IterationTiler{}));
+    }
+  }
+
+  template <typename ProblemShape>
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_local_b_shape(ProblemShape problem_shape) {
+    auto problem_shape_MNKL = append<4>(problem_shape, 1);
+    if constexpr (NumBuffersB == 0) {
+      return shape_div(
+            select<1,2,3>(problem_shape_MNKL),
+            select<1,2,3>(ProcessorTiler{}));
+    } else {
+      return shape_div(
+          shape_div(
+            select<1,2,3>(problem_shape_MNKL),
+            select<1,2,3>(ProcessorTiler{})),
+          select<1,2,3>(IterationTiler{}));
+    }
+  }
+
+  template <typename ProblemShape>
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_local_c_shape(ProblemShape problem_shape) {
+    auto problem_shape_MNKL = append<4>(problem_shape, 1);
+    if constexpr (not BufferedOutput) {
+      return shape_div(
+            select<0,1,3>(problem_shape_MNKL),
+            select<0,1,3>(ProcessorTiler{}));
+    } else {
+      return shape_div(
+          shape_div(
+            select<0,1,3>(problem_shape_MNKL),
+            select<0,1,3>(ProcessorTiler{})),
+          select<0,1,3>(IterationTiler{}));
+    }
+  }
+
+  template <typename ProblemShape>
+  CUTLASS_HOST_DEVICE
+  static auto
+  get_local_d_shape(ProblemShape problem_shape) {
+    auto problem_shape_MNKL = append<4>(problem_shape, 1);
+    if constexpr (not BufferedOutput) {
+      return shape_div(
+            select<0,1,3>(problem_shape_MNKL),
+            select<0,1,3>(ProcessorTiler{}));
+    } else {
+      return shape_div(
+          shape_div(
+            select<0,1,3>(problem_shape_MNKL),
+            select<0,1,3>(ProcessorTiler{})),
+          select<0,1,3>(IterationTiler{}));
+    }
+  }
+
+  // Host-side APIs: get_device_slice_{A,B,C,D}
+  // Slice off a view of the GLOBAL tensor that corresponds to the shard that 
+  // is going to be owned by a specific device. This helps with the initial 
+  // distribution of the GLOBAL operands among devices.
+  template <typename Tensor>
+  static auto
+  get_device_slice_A(Tensor tensor, int device_idx) {
+    auto tiler = get_processor_tiler_a(tensor);
+    return inner_partition(tensor, tiler, device_idx);
+  }
+
+  template <typename Tensor>
+  static auto
+  get_device_slice_B(Tensor tensor, int device_idx) {
+    auto tiler = get_processor_tiler_b(tensor);
+    return inner_partition(tensor, tiler, device_idx);
+  }
+
+  template <typename Tensor>
+  static auto
+  get_device_slice_C(Tensor tensor, int device_idx) {
+    auto tiler = get_processor_tiler_c(tensor);
+    return inner_partition(tensor, tiler, device_idx);
+  }
+
+  template <typename Tensor>
+  static auto
+  get_device_slice_D(Tensor tensor, int device_idx) {
+    auto tiler = get_processor_tiler_d(tensor);
+    return inner_partition(tensor, tiler, device_idx);
+  }
+};
+
+
+
+} // namespace cutlass::gemm::distributed
+
+///////////////////////////////////////////////////////////////////////////////
+

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/fast_math.h

@@ -0,0 +1,1085 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+#include "cutlass/cutlass.h"
+#if defined(__CUDACC_RTC__)
+#include CUDA_STD_HEADER(cstdint)
+#else
+#include <cstdint>
+#include <cmath>
+#include <type_traits>
+#endif
+#if !defined(__QNX__)
+#include CUDA_STD_HEADER(utility)
+#endif
+#include "cutlass/array.h"
+#include "cutlass/uint128.h"
+#include "cutlass/coord.h"
+#include "cutlass/half.h"
+
+/**
+ * \file
+ * \brief Math utilities
+ */
+
+namespace cutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+#if !defined(__QNX__)
+using ::cuda::std::swap;
+#else
+template <typename T>
+CUTLASS_HOST_DEVICE void swap(T &lhs, T &rhs) {
+  T tmp = lhs;
+  lhs = rhs;
+  rhs = tmp;
+}
+#endif
+
+/******************************************************************************
+ * Static math utilities
+ ******************************************************************************/
+
+/// Mixed precision dot product
+template <typename Index, typename LongIndex, int N>
+CUTLASS_HOST_DEVICE LongIndex dot(
+  Coord<N, Index> const &coord,
+  Coord<N, LongIndex> const &stride,
+  LongIndex acc = LongIndex()) {
+
+  CUTLASS_PRAGMA_UNROLL
+  for (int n = 0; n < N; ++n) {
+    acc += LongIndex(coord[n]) * stride[n];
+  }
+  return acc;
+}
+
+/**
+ * Statically determine if N is a power-of-two
+ */
+template <int N>
+struct is_pow2 {
+  static bool const value = ((N & (N - 1)) == 0);
+};
+
+/**
+ * Statically determine log2(N), rounded down
+ */
+template <int N, int CurrentVal = N, int Count = 0>
+struct log2_down {
+  /// Static logarithm value
+  enum { value = log2_down<N, (CurrentVal >> 1), Count + 1>::value };
+};
+
+// Base case
+template <int N, int Count>
+struct log2_down<N, 1, Count> {
+  enum { value = Count };
+};
+
+/**
+ * Statically determine log2(N), rounded up
+ */
+template <int N, int CurrentVal = N, int Count = 0>
+struct log2_up {
+  /// Static logarithm value
+  enum { value = log2_up<N, (CurrentVal >> 1), Count + 1>::value };
+};
+
+// Base case
+template <int N, int Count>
+struct log2_up<N, 1, Count> {
+  enum { value = ((1 << Count) < N) ? Count + 1 : Count };
+};
+
+/**
+ * Statically estimate sqrt(N) to the nearest power-of-two
+ */
+template <int N>
+struct sqrt_est {
+  enum { value = 1 << (log2_up<N>::value / 2) };
+};
+
+/**
+ * For performing a constant-division with a compile-time assertion that the
+ * Divisor evenly-divides the Dividend.
+ */
+template <int Dividend, int Divisor>
+struct divide_assert {
+  enum { value = Dividend / Divisor };
+
+  static_assert((Dividend % Divisor == 0), "Not an even multiple");
+};
+
+/******************************************************************************
+ * Rounding
+ ******************************************************************************/
+
+/**
+ * Round dividend up to the nearest multiple of divisor
+ */
+template <typename dividend_t, typename divisor_t>
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17
+dividend_t round_nearest(dividend_t dividend, divisor_t divisor) {
+  return ((dividend + divisor - 1) / divisor) * divisor;
+}
+
+template <typename value_t>
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17
+value_t abs_for_integer(value_t a) {
+  return ((a > value_t{0}) ? a : -a);
+}
+/**
+ * Greatest common divisor
+ */
+template <typename value_t>
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17
+value_t gcd(value_t a, value_t b) {
+  for (;;) {
+    if (a == value_t{0}) return cutlass::abs_for_integer(b);
+    b %= a;
+    if (b == value_t{0}) return cutlass::abs_for_integer(a);
+    a %= b;
+  }
+}
+
+/**
+ * Least common multiple
+ */
+template <typename value_t>
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17
+value_t lcm(value_t a, value_t b) {
+  value_t temp = cutlass::gcd(a, b);
+  return (temp != value_t{0}) ? value_t(cutlass::abs_for_integer(a) / temp * cutlass::abs_for_integer(b)) : value_t{};
+}
+
+/**
+ * Greatest common divisor
+ */
+template <typename value_t>
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17
+value_t gcd_cxx11(value_t a, value_t b) {
+  return (a == value_t{0} || b == value_t{0}) ? cutlass::abs_for_integer(a | b) : cutlass::gcd_cxx11(b, a % b);
+}
+
+/**
+ * Least common multiple
+ */
+template <typename value_t>
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17
+value_t lcm_cxx11(value_t a, value_t b) {
+  return cutlass::gcd_cxx11(a, b) ? (cutlass::abs_for_integer(a) / cutlass::gcd_cxx11(a, b) *
+                                    cutlass::abs_for_integer(b))
+                                  : value_t{};
+}
+
+/// Returns the smallest value in the half-open range [a, a+b) that is a multiple of b
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17
+int round_up(int a, int b) {
+  return ((a + b - 1) / b) * b;
+}
+
+/// Returns the ceiling of (a / b)
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17
+int ceil_div(int a, int b) {
+  return (a + b - 1) / b;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/**
+ * log2 computation, what's the
+ * difference between the below codes and
+ * log2_up/down codes?
+ */
+template <typename value_t>
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17
+value_t clz(value_t x) {
+  for (int i = 31; i >= 0; --i) {
+    if ((1 << i) & x)
+      return value_t(31 - i);
+  }
+  return value_t(32);
+}
+
+template <typename value_t>
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17
+value_t find_log2(value_t x) {
+  int a = int(31 - clz(x));
+  a += (x & (x - 1)) != 0;  // Round up, add 1 if not a power of 2.
+  return a;
+}
+
+
+/**
+ * Find divisor, using find_log2
+ */
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17
+void find_divisor(unsigned int& mul, unsigned int& shr, unsigned int denom) {
+  if (denom == 1) {
+    mul = 0;
+    shr = 0;
+  } else {
+    unsigned int p = 31 + find_log2(denom);
+    unsigned m = unsigned(((1ull << p) + unsigned(denom) - 1) / unsigned(denom));
+
+    mul = m;
+    shr = p - 32;
+  }
+}
+
+/**
+ * Find quotient and remainder using device-side intrinsics
+ */
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17
+void fast_divmod(int& quo, int& rem, int src, int div, unsigned int mul, unsigned int shr) {
+
+  #if defined(__CUDA_ARCH__)
+  // Use IMUL.HI if div != 1, else simply copy the source.
+  quo = (div != 1) ? __umulhi(src, mul) >> shr : src;
+  #else
+  quo = int((div != 1) ? int(((int64_t)src * mul) >> 32) >> shr : src);
+  #endif
+
+  // The remainder.
+  rem = src - (quo * div);
+}
+
+// For long int input
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17
+void fast_divmod(int& quo, int64_t& rem, int64_t src, int div, unsigned int mul, unsigned int shr) {
+
+  #if defined(__CUDA_ARCH__)
+  // Use IMUL.HI if div != 1, else simply copy the source.
+  quo = (div != 1) ? __umulhi(src, mul) >> shr : src;
+  #else
+  quo = int((div != 1) ? ((src * mul) >> 32) >> shr : src);
+  #endif
+  // The remainder.
+  rem = src - (quo * div);
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Object to encapsulate the fast division+modulus operation.
+///
+/// This object precomputes two values used to accelerate the computation and is best used
+/// when the divisor is a grid-invariant. In this case, it may be computed in host code and
+/// marshalled along other kernel arguments using the 'Params' pattern.
+///
+/// Example:
+///
+///
+///   int quotient, remainder, dividend, divisor;
+///
+///   FastDivmod divmod(divisor);
+///
+///   divmod(quotient, remainder, dividend);
+///
+///   // quotient = (dividend / divisor)
+///   // remainder = (dividend % divisor)
+///
+struct FastDivmod {
+  using value_div_type = int;
+  using value_mod_type = int64_t;
+  int32_t divisor = 1;
+  uint32_t multiplier = 0u;
+  uint32_t shift_right = 0u;
+
+  // Find quotient and remainder using device-side intrinsics
+  CUTLASS_HOST_DEVICE
+  void fast_divmod(int& quotient, int& remainder, int dividend) const {
+
+#if defined(__CUDA_ARCH__)
+    // Use IMUL.HI if divisor != 1, else simply copy the source.
+    quotient = (divisor != 1) ? __umulhi(dividend, multiplier) >> shift_right : dividend;
+#else
+    quotient = int((divisor != 1) ? int(((int64_t)dividend * multiplier) >> 32) >> shift_right : dividend);
+#endif
+
+    // The remainder.
+    remainder = dividend - (quotient * divisor);
+  }
+
+  /// For long int input
+  CUTLASS_HOST_DEVICE
+  void fast_divmod(int& quotient, int64_t& remainder, int64_t dividend) const {
+
+#if defined(__CUDA_ARCH__)
+    // Use IMUL.HI if divisor != 1, else simply copy the source.
+    quotient = (divisor != 1) ? __umulhi(dividend, multiplier) >> shift_right : dividend;
+#else
+    quotient = int((divisor != 1) ? ((dividend * multiplier) >> 32) >> shift_right : dividend);
+#endif
+    // The remainder.
+    remainder = dividend - (quotient * divisor);
+  }
+
+
+  /// Construct the FastDivmod object, in host code ideally.
+  ///
+  /// This precomputes some values based on the divisor and is computationally expensive.
+
+  constexpr FastDivmod() = default;
+
+  CUTLASS_HOST_DEVICE
+  FastDivmod(int divisor_): divisor(divisor_) {
+    assert(divisor_ >= 0);
+    if (divisor != 1) {
+      unsigned int p = 31 + find_log2(divisor);
+      unsigned m = unsigned(((1ull << p) + unsigned(divisor) - 1) / unsigned(divisor));
+
+      multiplier = m;
+      shift_right = p - 32;
+    }
+  }
+
+  /// Computes integer division and modulus using precomputed values. This is computationally
+  /// inexpensive.
+  CUTLASS_HOST_DEVICE
+  void operator()(int &quotient, int &remainder, int dividend) const {
+    fast_divmod(quotient, remainder, dividend);
+  }
+
+  /// Computes integer division using precomputed values. This is computationally
+  /// inexpensive.
+  CUTLASS_HOST_DEVICE
+  int div(int dividend) const {
+    int quotient, remainder;
+    fast_divmod(quotient, remainder, dividend);
+    return quotient;
+  }
+
+  /// Alias for `div` to match the interface of FastDivmodU64
+  CUTLASS_HOST_DEVICE
+  int divide(int dividend) const {
+    return div(dividend);
+  }
+
+  /// Computes integer division remainder using precomputed values.
+  CUTLASS_HOST_DEVICE
+  int rem(int dividend) const {
+    int quotient, remainder;
+    fast_divmod(quotient, remainder, dividend);
+    return remainder;
+  }
+
+  /// Alias for `rem`
+  CUTLASS_HOST_DEVICE
+  int remainder(int dividend) const {
+    return rem(dividend);
+  }
+
+  /// Computes integer division and modulus using precomputed values. This is computationally
+  /// inexpensive.
+  ///
+  /// Simply returns the quotient
+  CUTLASS_HOST_DEVICE
+  int divmod(int &remainder, int dividend) const {
+    int quotient;
+    fast_divmod(quotient, remainder, dividend);
+    return quotient;
+  }
+
+  /// Computes integer division and modulus using precomputed values. This is computationally
+  /// inexpensive.
+  CUTLASS_HOST_DEVICE
+  void operator()(int &quotient, int64_t &remainder, int64_t dividend) const {
+    fast_divmod(quotient, remainder, dividend);
+  }
+
+  /// Computes integer division and modulus using precomputed values. This is computationally
+  /// inexpensive.
+  CUTLASS_HOST_DEVICE
+  int divmod(int64_t &remainder, int64_t dividend) const {
+    int quotient;
+    fast_divmod(quotient, remainder, dividend);
+    return quotient;
+  }
+
+  /// Returns the divisor when cast to integer
+  CUTLASS_HOST_DEVICE
+  operator int() const { return divisor; }
+
+};
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Object to encapsulate the fast division+modulus operation for 64b integer division.
+///
+/// This object precomputes two values used to accelerate the computation and is best used
+/// when the divisor is a grid-invariant. In this case, it may be computed in host code and
+/// marshalled along other kernel arguments using the 'Params' pattern.
+///
+/// Example:
+///
+///
+///   uint64_t quotient, remainder, dividend, divisor;
+///
+///   FastDivmodU64 divmod(divisor);
+///
+///   divmod(quotient, remainder, dividend);
+///
+///   // quotient = (dividend / divisor)
+///   // remainder = (dividend % divisor)
+///
+struct FastDivmodU64 {
+
+  uint64_t divisor;
+  uint64_t multiplier;
+  unsigned int shift_right;
+  unsigned int round_up;
+
+  //
+  // Static methods
+  //
+
+  /// Computes b, where 2^b is the greatest power of two that is less than or equal to x
+  CUTLASS_HOST_DEVICE
+  static uint32_t integer_log2(uint64_t x) {
+    uint32_t n = 0;
+    while (x >>= 1) {
+      ++n;
+    }
+    return n;
+  }
+
+  /// Default ctor
+  CUTLASS_HOST_DEVICE
+  FastDivmodU64(): divisor(0), multiplier(0), shift_right(0), round_up(0) { }
+
+  /// Construct the FastDivmod object, in host code ideally.
+  ///
+  /// This precomputes some values based on the divisor and is computationally expensive.
+  CUTLASS_HOST_DEVICE
+  FastDivmodU64(uint64_t divisor_): divisor(divisor_), multiplier(1), shift_right(0), round_up(0) {
+
+    if (divisor) {
+      shift_right = integer_log2(divisor);
+
+      if ((divisor & (divisor - 1)) == 0) {
+        multiplier = 0;
+      }
+      else {
+        uint64_t power_of_two = (uint64_t(1) << shift_right);
+        uint64_t multiplier_lo = uint128_t(0, power_of_two) / divisor;
+        multiplier = uint128_t(power_of_two, power_of_two) / divisor;
+        round_up = (multiplier_lo == multiplier ? 1 : 0);
+      }
+    }
+  }
+
+  /// Returns the quotient of floor(dividend / divisor)
+  CUTLASS_HOST_DEVICE
+  uint64_t divide(uint64_t dividend) const {
+    uint64_t quotient = 0;
+
+    #ifdef __CUDA_ARCH__
+      uint64_t x = dividend;
+      if (multiplier) {
+        x = __umul64hi(dividend + round_up, multiplier);
+      }
+      quotient = (x >> shift_right);
+    #else
+      quotient = dividend / divisor;
+    #endif
+
+    return quotient;
+  }
+
+  /// Computes the remainder given a computed quotient and dividend
+  CUTLASS_HOST_DEVICE
+  uint64_t modulus(uint64_t quotient, uint64_t dividend) const {
+    return dividend - quotient * divisor;
+  }
+
+  /// Returns the quotient of floor(dividend / divisor) and computes the remainder
+  CUTLASS_HOST_DEVICE
+  uint64_t divmod(uint64_t &remainder, uint64_t dividend) const {
+    uint64_t quotient = divide(dividend);
+    remainder = modulus(quotient, dividend);
+    return quotient;
+  }
+
+  /// Computes integer division and modulus using precomputed values. This is computationally
+  /// inexpensive.
+  CUTLASS_HOST_DEVICE
+  void operator()(uint64_t &quotient, uint64_t &remainder, uint64_t dividend) const {
+    quotient = divmod(remainder, dividend);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Object to encapsulate the fast division+modulus operation for 64b integer division
+/// in which the divisor is a power of two.
+struct FastDivmodU64Pow2 {
+
+  uint64_t divisor;
+  unsigned int shift_right;
+
+  /// Default ctor
+  CUTLASS_HOST_DEVICE
+  FastDivmodU64Pow2(): divisor(0), shift_right(0) { }
+
+  /// Construct the FastDivmod object, in host code ideally.
+  ///
+  /// This precomputes some values based on the divisor and is computationally expensive.
+  CUTLASS_HOST_DEVICE
+  FastDivmodU64Pow2(uint64_t divisor_): divisor(divisor_), shift_right(FastDivmodU64::integer_log2(divisor_)) { }
+
+  /// Returns the quotient of floor(dividend / divisor)
+  CUTLASS_HOST_DEVICE
+  uint64_t divide(uint64_t dividend) const {
+    return dividend >> shift_right;
+  }
+
+  /// Computes the remainder given a computed quotient and dividend
+  CUTLASS_HOST_DEVICE
+  uint64_t modulus(uint64_t dividend) const {
+    // See https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html#division-modulo-operations
+    return dividend & (divisor - 1);
+  }
+
+  /// Returns the quotient of floor(dividend / divisor) and computes the remainder
+  CUTLASS_HOST_DEVICE
+  uint64_t divmod(uint64_t &remainder, uint64_t dividend) const {
+    uint64_t quotient = divide(dividend);
+    remainder = modulus(dividend);
+    return quotient;
+  }
+
+  /// Computes integer division and modulus using precomputed values. This is computationally
+  /// inexpensive.
+  CUTLASS_HOST_DEVICE
+  void operator()(uint64_t &quotient, uint64_t &remainder, uint64_t dividend) const {
+    quotient = divmod(remainder, dividend);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes the coordinate decomposition from a linear index (64-bit linear index => coord<int32_t>)
+///
+/// This decomposition is accelerated by the FastDivmodU64 object. It is assumed that
+/// a coordinate of <Rank> indices can be decomposed by <Rank - 1> div/mod operations.
+/// Note, is assumed that element divmod[0] divides by extent[1].
+///
+/// For example, assume 4-D coordinate (n, p, q, c) is mapped to a linear index `npqc`. This
+/// can be decomposed via three divide and modulus operations:
+///
+///      c = npqc % C;         |  divmod[2] = FastDivmodU64(C)
+///    npq = npqc / C;         |   coord[3] = c
+///
+///      q =  npq % Q;         |  divmod[1] = FastDivmodU64(Q)
+///     np =  npq / Q;         |   coord[2] = q
+///
+///      p =   np % P;         |  divmod[0] = FastDivmodU64(P)
+///      n =   np / P;         |   coord[1] = p
+///
+///                            |   coord[0] = n
+///
+template <int Rank>
+CUTLASS_HOST_DEVICE Coord<Rank> CoordinateDecomposition(
+  uint64_t linear_idx,                    ///< Linear index to decompose
+  FastDivmodU64 const *divmod) {          ///< Pointer to array of Rank-1 FastDivmodU64 objects
+
+  static_assert(Rank > 0, "CoordinateDecomposition requires Rank=1 or greater.");
+
+  Coord<Rank> coord;
+
+  CUTLASS_PRAGMA_UNROLL
+  for (int i = Rank; i > 1; --i) {
+    uint64_t remainder;
+    linear_idx = divmod[i - 2].divmod(remainder, linear_idx);
+    coord[i - 1] = int(remainder);
+  }
+
+  coord[0] = int(linear_idx);
+
+  return coord;
+}
+
+/// Computes the coordinate decomposition from a linear index (32-bit linear index => coord<int32_t>)
+template <int Rank>
+CUTLASS_HOST_DEVICE Coord<Rank> CoordinateDecomposition(
+  int linear_idx,                    ///< Linear index to decompose
+  FastDivmod const *divmod) {          ///< Pointer to array of Rank-1 FastDivmodU64 objects
+
+  static_assert(Rank > 0, "CoordinateDecomposition requires Rank=1 or greater.");
+
+  Coord<Rank> coord;
+
+  CUTLASS_PRAGMA_UNROLL
+  for (int i = Rank; i > 1; --i) {
+    int remainder;
+    linear_idx = divmod[i - 2].divmod(remainder, linear_idx);
+    coord[i - 1] = int(remainder);
+  }
+
+  coord[0] = int(linear_idx);
+
+  return coord;
+}
+
+template <int Rank>
+CUTLASS_HOST_DEVICE Coord<Rank> CoordinateDecompositionLittleEndian(
+  uint64_t linear_idx,                    ///< Linear index to decompose
+  FastDivmodU64 const *divmod) {          ///< Pointer to array of Rank-1 FastDivmodU64 objects
+
+  static_assert(Rank > 0, "CoordinateDecomposition requires Rank=1 or greater.");
+
+  Coord<Rank> coord;
+
+  CUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < Rank - 1; ++i) {
+    uint64_t remainder;
+    linear_idx = divmod[i].divmod(remainder, linear_idx);
+    coord[i] = int(remainder);
+  }
+
+  coord[Rank - 1] = int(linear_idx);
+
+  return coord;
+}
+
+/// Computes the coordinate decomposition from a linear index (32-bit linear index => coord<int32_t>)
+template <int Rank>
+CUTLASS_HOST_DEVICE Coord<Rank> CoordinateDecompositionLittleEndian(
+  int linear_idx,                    ///< Linear index to decompose
+  FastDivmod const *divmod) {          ///< Pointer to array of Rank-1 FastDivmodU64 objects
+
+  static_assert(Rank > 0, "CoordinateDecomposition requires Rank=1 or greater.");
+
+  Coord<Rank> coord;
+
+  CUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < Rank - 1; ++i) {
+    int remainder;
+    linear_idx = divmod[i].divmod(remainder, linear_idx);
+    coord[i] = int(remainder);
+  }
+
+  coord[Rank - 1] = int(linear_idx);
+
+  return coord;
+}
+
+/// Safely computes the offset of a linear index in bytes for all types
+template <typename Element>
+CUTLASS_HOST_DEVICE int64_t OffsetBytes(int64_t index) {
+
+  static_assert(
+    (sizeof_bits<Element>::value >= 8 && !(sizeof_bits<Element>::value % 8)) ||
+    (sizeof_bits<Element>::value <  8 && !(8 % sizeof_bits<Element>::value)),
+    "Size of numeric type in bits must either be divisible by 8 bits, or 8 bits must be divisible by the size.");
+
+  if (sizeof_bits<Element>::value >= 8) {
+    return index * (sizeof_bits<Element>::value / 8);
+  }
+  else {
+    int const kElementsPerByte = ((8 / sizeof_bits<Element>::value) + ((sizeof_bits<Element>::value >= 8) ? 1 : 0));
+    return index / kElementsPerByte;
+  }
+}
+
+CUTLASS_HOST_DEVICE int64_t OffsetBytes(int64_t index, int64_t element_sizeof_bits) {
+  if (element_sizeof_bits >= 8) {
+    return index * (element_sizeof_bits / 8);
+  }
+  else {
+    int64_t const kElementsPerByte = ((8 / element_sizeof_bits) + ((element_sizeof_bits >= 8) ? 1 : 0));
+    return index / kElementsPerByte;
+  }
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// Min/Max
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int A, int B>
+struct Min {
+  static int const kValue = (A < B) ? A : B;
+};
+
+template <int A, int B>
+struct Max {
+  static int const kValue = (A > B) ? A : B;
+};
+
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17 int const_min(int a, int b) {
+    return (b < a ? b : a);
+}
+
+CUTLASS_HOST_DEVICE
+CUTLASS_CONSTEXPR_IF_CXX17 int const_max(int a, int b) {
+    return (b > a ? b : a);
+}
+
+template <typename T>
+CUTLASS_HOST_DEVICE
+T fast_min(T a, T b) {
+  return (b < a ? b : a);
+}
+
+template <>
+CUTLASS_HOST_DEVICE
+float fast_min(float a, float b) {
+  return fminf(a, b);
+}
+
+template <typename T>
+CUTLASS_HOST_DEVICE
+T fast_max(T a, T b) {
+  return (a < b ? b : a);
+}
+
+template <>
+CUTLASS_HOST_DEVICE
+float fast_max(float a, float b) {
+  return fmaxf(a, b);
+}
+
+CUTLASS_HOST_DEVICE
+float fast_cos(float theta) {
+  #if defined(__CUDA_ARCH__)
+  return ::cosf(theta);
+  #else
+  return std::cos(theta);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+double fast_cos(double theta) {
+  #if defined(__CUDA_ARCH__)
+  return ::cos(theta);
+  #else
+  return std::cos(theta);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+float fast_sin(float theta) {
+  #if defined(__CUDA_ARCH__)
+  return ::sinf(theta);
+  #else
+  return std::sin(theta);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+double fast_sin(double theta) {
+  #if defined(__CUDA_ARCH__)
+  return ::sin(theta);
+  #else
+  return std::sin(theta);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+float fast_acos(float theta) {
+  #if defined(__CUDA_ARCH__)
+  return ::acosf(theta);
+  #else
+  return std::acos(theta);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+double fast_acos(double theta) {
+  #if defined(__CUDA_ARCH__)
+  return ::acos(theta);
+  #else
+  return std::acos(theta);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+float fast_asin(float theta) {
+  #if defined(__CUDA_ARCH__)
+  return ::asinf(theta);
+  #else
+  return std::asin(theta);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+double fast_asin(double theta) {
+  #if defined(__CUDA_ARCH__)
+  return ::asin(theta);
+  #else
+  return std::asin(theta);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+float fast_sqrt(float theta) {
+  #if defined(__CUDA_ARCH__)
+  return ::sqrtf(theta);
+  #else
+  return std::sqrt(theta);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+double fast_sqrt(double theta) {
+  #if defined(__CUDA_ARCH__)
+  return ::sqrt(theta);
+  #else
+  return std::sqrt(theta);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+float fast_exp(float x) {
+  #if defined(__CUDA_ARCH__)
+  return ::expf(x);
+  #else
+  return std::exp(x);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+double fast_exp(double x) {
+  #if defined(__CUDA_ARCH__)
+  return ::exp(x);
+  #else
+  return std::exp(x);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+half_t fast_exp(half_t x) {
+  #if defined(__CUDA_ARCH__) && (__CUDACC_VER_MAJOR__ >= 10) && (__CUDA_ARCH__ >= 750)
+      return (half_t)(::hexp(x.to_half()));
+  #else
+      return (half_t)(fast_exp(float(x)));
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+float fast_log(float x) {
+  #if defined(__CUDA_ARCH__)
+  return ::logf(x);
+  #else
+  return std::log(x);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+double fast_log(double x) {
+  #if defined(__CUDA_ARCH__)
+  return ::log(x);
+  #else
+  return std::log(x);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+float fast_tanh(float x) {
+  #if defined(__CUDA_ARCH__)
+    #if (__CUDACC_VER_MAJOR__ >= 11) && (__CUDA_ARCH__ >= 750)
+      float y;
+      asm volatile ( "tanh.approx.f32 %0, %1; " : "=f"(y) : "f"(x));
+      return y;
+    #else
+      return ::tanhf(x);
+    #endif
+  #else
+  return std::tanh(x);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+double fast_tanh(double x) {
+  #if defined(__CUDA_ARCH__)
+  return ::tanh(x);
+  #else
+  return std::tanh(x);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+half_t fast_tanh(half_t x) {
+  #if defined(__CUDA_ARCH__) && (__CUDACC_VER_MAJOR__ >= 11) && (__CUDA_ARCH__ >= 750)
+
+  asm volatile ( "tanh.approx.f16 %0, %1;" : "=h"(x.raw()) : "h"(x.raw()));
+  return x;
+
+  #else
+  return half_t(fast_tanh(float(x)));
+  #endif
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+struct fast_exp_op {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &rhs) const {
+    return fast_exp(rhs);
+  }
+};
+
+#if defined(__CUDA_ARCH__) && (__CUDACC_VER_MAJOR__ >= 10) && (__CUDA_ARCH__ >= 750)
+template <int N>
+struct fast_exp_op<Array<half_t, N>> {
+  CUTLASS_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const &rhs) const {
+
+    Array<half_t, N> result;
+
+    // use x2 specialization
+    __half2 const *in  = reinterpret_cast<__half2 const *>(&rhs);
+    __half2 *out = reinterpret_cast<__half2 *>(&result);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      out[i] = ::h2exp(in[i]);
+    }
+
+    // residual
+    if (N % 2) {
+      half_t last = rhs[N - 1];
+      result[N - 1] = half_t(::hexp(last.to_half()));
+    }
+
+    return result;
+  }
+};
+#endif // #if defined(__CUDA_ARCH__)
+
+template <typename T, int N>
+struct fast_exp_op<Array<T, N>> {
+  CUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &rhs) const {
+
+    fast_exp_op<T> fast_op;
+    Array<T, N> y;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      y[i] = fast_op(rhs[i]);
+    }
+
+    return y;
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+struct fast_tanh_op {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &rhs) const {
+    return fast_tanh(rhs);
+  }
+};
+
+#if defined(__CUDA_ARCH__) && (__CUDACC_VER_MAJOR__ >= 11) && (__CUDA_ARCH__ >= 750)
+template <int N>
+struct fast_tanh_op<Array<half_t, N>> {
+  CUTLASS_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const &rhs) const {
+
+    Array<half_t, N> result;
+
+    // use x2 specialization
+    uint32_t const *in  = reinterpret_cast<uint32_t const *>(&rhs);
+    uint32_t *out = reinterpret_cast<uint32_t *>(&result);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      asm volatile ("tanh.approx.f16x2 %0, %1;" : "=r"(out[i]) : "r"(in[i]));
+    }
+
+    // residual
+    if (N % 2) {
+      uint16_t const *in = reinterpret_cast<uint16_t const *>(&rhs);
+      uint16_t *out = reinterpret_cast<uint16_t *>(&result);
+      asm volatile ("tanh.approx.f16 %0, %1;" : "=h"(out[N - 1]) : "h"(in[N - 1]));
+    }
+
+    return result;
+  }
+};
+#endif // #if defined(__CUDA_ARCH__)
+
+template <typename T, int N>
+struct fast_tanh_op<Array<T, N>> {
+  CUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &rhs) const {
+
+    fast_tanh_op<T> fast_op;
+    Array<T, N> y;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      y[i] = fast_op(rhs[i]);
+    }
+
+    return y;
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Absolute value function
+template <typename T>
+CUTLASS_HOST_DEVICE
+T absolute_value(T x) {
+  if (x < T()) {
+    return -x;
+  }
+  return x;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/float8.h

@@ -0,0 +1,1685 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*!
+    \file
+    \brief Defines a class for using IEEE half-precision floating-point types in host or
+      device code.
+*/
+
+#pragma once
+
+
+#include "cutlass/arch/config.h"
+
+
+// FP8 types are available starting CUDA 11.8+
+#if (__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 8))
+#define CUDA_FP8_ENABLED 1
+#endif
+
+#if defined(__CUDA_ARCH__)
+#  if (__CUDA_ARCH__ >= 900)
+#    if (__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 8))
+#      define CUDA_PTX_FP8_CVT_ENABLED 1
+#    endif // (__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 8))
+#  elif (__CUDA_ARCH__ == 890)
+#    if (__CUDACC_VER_MAJOR__ > 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ >= 1))
+#      define CUDA_PTX_FP8_CVT_ENABLED 1
+#    endif // (__CUDACC_VER_MAJOR__ > 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ >= 1))
+#  endif // (__CUDA_ARCH__ >= 900)
+#endif // defined(__CUDA_ARCH__)
+
+
+#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED) ||\
+     defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM110A_ENABLED) ||\
+     defined(CUTLASS_ARCH_MMA_SM120A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121A_ENABLED))
+#  define CUDA_PTX_UE8M0_CVT_ENABLED 1
+#endif
+
+#if (defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101F_ENABLED) ||\
+     defined(CUTLASS_ARCH_MMA_SM103F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM110F_ENABLED) ||\
+     defined(CUTLASS_ARCH_MMA_SM120F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121F_ENABLED))
+#  define CUDA_PTX_UE8M0_CVT_ENABLED 1
+#endif
+
+#ifdef __GNUC__
+// Ignore checks on reinterpret-casts that are being used for bitcasts.
+#pragma GCC diagnostic ignored "-Wstrict-aliasing"
+#endif
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+#if defined(__CUDACC_RTC__)
+
+#include "cutlass/floating_point_nvrtc.h"
+
+#else
+//
+// Standard Library headers belong here to avoid conflicts with NVRTC.
+//
+#include <cmath>
+#include <limits>
+#include <cstdint>
+#include <cstring>
+#endif
+
+#ifdef CUDA_FP8_ENABLED
+#include <cuda_fp8.h>
+#endif
+#include <cuda_fp16.h>
+
+#include "cutlass/cutlass.h"
+
+#include "cutlass/exmy_base.h"
+
+#include "cute/util/type_traits.hpp"
+
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+//  FP8 Has 2 encodings possible : E4M3 and E5M2
+//
+//  E4M3 : 7  |  6 5 4 3  |  2 1 0
+//  E5M2 : 7  |  6 5 4 3 2  |  1 0
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+enum class FloatEncoding {
+    E4M3,
+    E5M2
+};
+
+template<FloatEncoding T>
+struct alignas(1) float8_base {
+
+    static constexpr bool IS_E4M3 = (T == FloatEncoding::E4M3);
+    static constexpr bool IS_E5M2 = (T == FloatEncoding::E5M2);
+
+    // Number of Bits representing mantissa and exponents
+    static constexpr int FP32_NUM_BITS = 32;
+    static constexpr int FP32_NUM_EXPONENT_BITS = 8;
+    static constexpr int FP32_NUM_MANTISSA_BITS = 23;
+    static constexpr uint32_t FP32_NAN = 0x7fffffff;
+    static constexpr uint32_t FP32_INFINITY_MASK = 0x7f800000;
+    static constexpr int FP32_MAX_EXPONENT  =  127;
+    static constexpr int FP32_MIN_EXPONENT  = -126;
+    static constexpr int FP32_EXPONENT_BIAS =  127;
+
+    static constexpr int FP16_NUM_BITS = 16;
+    static constexpr int FP16_NUM_EXPONENT_BITS = 5;
+    static constexpr int FP16_NUM_MANTISSA_BITS = 10;
+    static constexpr uint16_t FP16_NAN = 0x7fff;
+    static constexpr uint16_t FP16_INFINITY_MASK = 0x7c00;
+    static constexpr int FP16_MAX_EXPONENT  = 15;
+    static constexpr int FP16_MIN_EXPONENT  = -14;
+    static constexpr int FP16_EXPONENT_BIAS = 15;
+
+    static constexpr int FP8_NUM_BITS = 8;
+    static constexpr int FP8_NUM_EXPONENT_BITS = IS_E4M3 ? 4 : 5;
+    static constexpr int FP8_NUM_MANTISSA_BITS = IS_E4M3 ? 3 : 2;
+    static constexpr uint8_t  FP8_NAN = 0x7f; // Also F8_INF
+    static constexpr uint8_t  FP8_INFINITY_MASK = IS_E4M3 ? 0x78 : 0x7c;
+    static constexpr int FP8_MAX_EXPONENT  = IS_E4M3 ?  7 :  15;
+    static constexpr int FP8_MIN_EXPONENT  = IS_E4M3 ? -6 : -14;
+    static constexpr int FP8_EXPONENT_BIAS = IS_E4M3 ?  7 :  15;
+
+    static constexpr uint8_t  FP8_EXPONENT_MASK = (1 << FP8_NUM_EXPONENT_BITS) - 1;
+    static constexpr uint8_t  FP8_MANTISSA_MASK = (1 << FP8_NUM_MANTISSA_BITS) - 1;
+
+    static constexpr uint8_t FP8_MAX_FLT = (IS_E4M3 ? 0x7e : 0x7b);
+
+    // 256 in float
+    static constexpr uint32_t FP8_SAT_VAL_FP32 = 0x43800000;
+
+    //
+    // Data members
+    //
+
+    /// Data container
+    uint8_t storage;
+
+    /// Ctors.
+    CUTLASS_HOST_DEVICE
+    float8_base() : storage(0) { }
+
+    /// Is finite implementation
+    CUTLASS_HOST_DEVICE
+    static bool isfinite(float flt) {
+        uint32_t s;
+
+        #if defined(__CUDA_ARCH__)
+        s = reinterpret_cast<uint32_t const &>(flt);
+        #else
+        std::memcpy(&s, &flt, sizeof(s));
+        #endif
+
+        return (s & 0x7f800000) < 0x7f800000;
+    }
+
+    /// Is NaN implementation
+    CUTLASS_HOST_DEVICE
+    static bool isnan(float flt) {
+        uint32_t s;
+
+        #if defined(__CUDA_ARCH__)
+        s = reinterpret_cast<uint32_t const &>(flt);
+        #else
+        std::memcpy(&s, &flt, sizeof(s));
+        #endif
+
+        return (s & 0x7fffffff) > 0x7f800000;
+    }
+
+    /// Is infinite implementation
+    CUTLASS_HOST_DEVICE
+    static bool isinf(float flt) {
+        uint32_t s;
+
+        #if defined(__CUDA_ARCH__)
+        s = reinterpret_cast<uint32_t const &>(flt);
+        #else
+        std::memcpy(&s, &flt, sizeof(s));
+        #endif
+
+        // Sign = 0 for +inf, 1 for -inf
+        // Exponent = all ones
+        // Mantissa = all zeros
+        return (s == 0x7f800000) || (s == 0xff800000);
+    }
+
+    /// FP32 -> FP8 conversion - rounds to nearest even
+    CUTLASS_HOST_DEVICE
+    static uint8_t convert_float_to_fp8(float const& flt) {
+
+        // software implementation rounds toward nearest even
+        uint32_t s;
+
+        #if defined(__CUDA_ARCH__)
+        s = reinterpret_cast<uint32_t const &>(flt);
+        #else
+        std::memcpy(&s, &flt, sizeof(s));
+        #endif
+
+        // Extract the bits in the FP32 type
+        uint8_t sign = uint8_t((s >> 24 & 0x80));
+        int32_t exp = int32_t((s >> FP32_NUM_MANTISSA_BITS) & 0xff) - FP32_EXPONENT_BIAS;
+        int mantissa = s & 0x7fffff;
+        uint8_t u = 0;
+
+        uint8_t const kF8_NaN = 0x7f;
+
+        // NaN => NaN
+        if (isnan(flt)) {
+            return kF8_NaN;
+        }
+
+        // Inf => MAX_FLT (satfinite)
+        if (isinf(flt)) {
+            return sign | FP8_MAX_FLT;
+        }
+
+        // Special handling
+        if (exp == -128) {
+            // int8 range is from -128 to 127
+            // So 255(inf) - 127(bias) = 128 - will show up as -128
+
+            // satfinite
+            return (sign | FP8_MAX_FLT);
+        }
+
+        int sticky_bit = 0;
+
+        bool skip_sign = false;
+        bool may_be_nan = false;
+
+        if ( (exp >= FP8_MIN_EXPONENT) && (exp <= FP8_MAX_EXPONENT) ) {
+            // normal fp32 to normal fp8
+            exp = exp + FP8_EXPONENT_BIAS;
+            u = uint8_t((uint32_t(exp) & FP8_EXPONENT_MASK) << FP8_NUM_MANTISSA_BITS);
+            u = uint8_t(u | (mantissa >> (FP32_NUM_MANTISSA_BITS - FP8_NUM_MANTISSA_BITS)));
+        } else if(exp < FP8_MIN_EXPONENT) {
+            // normal single-precision to subnormal float8-precision representation
+            int rshift = (FP8_MIN_EXPONENT - exp);
+            if (rshift < FP32_NUM_BITS) {
+                mantissa |= (1 << FP32_NUM_MANTISSA_BITS);
+
+                sticky_bit = ((mantissa & ((1 << rshift) - 1)) != 0);
+
+                mantissa = (mantissa >> rshift);
+                u = (uint8_t(mantissa >> (FP32_NUM_MANTISSA_BITS- FP8_NUM_MANTISSA_BITS)) & FP8_MANTISSA_MASK);
+            } else {
+                mantissa = 0;
+                u = 0;
+            }
+        // Exponent > FP8_MAX_EXPONENT - this is a special case done to match HW
+        // 0x4380_0000 to 0x43e0_0000 - maps from 256 to 448, and does not saturate / inf.
+        } else {
+            if( exp == (FP8_MAX_EXPONENT + 1) ) {
+                uint8_t mantissa_tmp = uint8_t(mantissa >> (FP32_NUM_MANTISSA_BITS - FP8_NUM_MANTISSA_BITS));
+                if( mantissa_tmp < FP8_MANTISSA_MASK) {
+                    exp = exp + FP8_EXPONENT_BIAS;
+                    u = uint8_t(uint32_t(exp) << FP8_NUM_MANTISSA_BITS) | mantissa_tmp;
+                    may_be_nan =  (mantissa_tmp == (FP8_MANTISSA_MASK-1));
+                } else {
+                    // satfinite
+                    return (sign | FP8_MAX_FLT);
+                }
+            } else{
+                // satfinite
+                return (sign | FP8_MAX_FLT);
+            }
+        }
+
+        // round to nearest even
+        int NUM_BITS_SHIFT = FP32_NUM_MANTISSA_BITS - (FP8_NUM_MANTISSA_BITS + 1);
+        int round_bit = ((mantissa >> NUM_BITS_SHIFT) & 1);
+        sticky_bit |= ((mantissa & ((1 << NUM_BITS_SHIFT) - 1)) != 0);
+
+        if ((round_bit && sticky_bit) || (round_bit && (u & 1))) {
+            u = uint8_t(u + 1);
+            if( may_be_nan ) {
+                skip_sign = true;
+            }
+        }
+
+        if (u > FP8_MAX_FLT) {
+            // satfinite
+            u = (sign | FP8_MAX_FLT);
+        }
+
+        if( ! skip_sign ) {
+            u |= sign;
+        }
+
+        return u;
+    }
+
+
+    /// Converts a fp8 value stored as a uint8_t to a float
+    CUTLASS_HOST_DEVICE
+    static float convert_fp8_to_float(uint8_t const& x) {
+
+        uint32_t constexpr kF32_NaN = 0x7fffffff;
+
+        uint8_t const &f8 = x;
+        uint32_t sign = (f8 >> (FP8_NUM_BITS - 1)) & 1;
+        uint32_t exp = (f8 >> FP8_NUM_MANTISSA_BITS) & FP8_EXPONENT_MASK;
+        uint32_t mantissa = f8 & FP8_MANTISSA_MASK;
+        unsigned f = (sign << (FP32_NUM_BITS-1));
+
+        if (IS_E4M3 && exp == 15 && mantissa == 0x7) {
+            f = kF32_NaN;
+        }
+        else if (exp > 0 && (IS_E4M3 || exp < (FP8_MAX_EXPONENT + FP8_EXPONENT_BIAS + 1))) {
+            // normal
+            exp += (FP32_EXPONENT_BIAS - FP8_EXPONENT_BIAS);
+            f = f |
+                (exp << FP32_NUM_MANTISSA_BITS) |
+                (mantissa << (FP32_NUM_MANTISSA_BITS-FP8_NUM_MANTISSA_BITS));
+        } else if (exp == 0) {
+            if (mantissa) {
+                // subnormal
+                exp += (FP32_EXPONENT_BIAS - FP8_EXPONENT_BIAS) + 1;
+                while ((mantissa & (1 << FP8_NUM_MANTISSA_BITS)) == 0) {
+                    mantissa <<= 1;
+                    exp--;
+                }
+                mantissa &= FP8_MANTISSA_MASK;
+                f = f |
+                    (exp << FP32_NUM_MANTISSA_BITS) |
+                    (mantissa << (FP32_NUM_MANTISSA_BITS-FP8_NUM_MANTISSA_BITS));
+            } else {
+                // sign-preserving zero
+            }
+        } else {
+            if(mantissa == 0){
+                // Sign-preserving infinity
+                f = (f | 0x7f800000);
+            } else {
+                // Canonical NaN
+                f = kF32_NaN;
+            }
+        }
+
+        #if defined(__CUDA_ARCH__)
+        return reinterpret_cast<float const&>(f);
+        #else
+        float flt;
+        std::memcpy(&flt, &f, sizeof(flt));
+        return flt;
+        #endif
+    }
+};
+
+
+// Forward declaration of float_e5m2_t to define float_e4m3_t <=> float_e5m2_t
+// conversions in class float_e4m3_t
+struct float_e5m2_t;
+
+
+///////////////////////////////////////////////////////////////
+///
+/// floating-point 8 type : E4M3
+///
+///////////////////////////////////////////////////////////////
+struct alignas(1) float_e4m3_t : float8_base<FloatEncoding::E4M3> {
+
+    using Base = float8_base<FloatEncoding::E4M3>;
+
+    static constexpr int MAX_EXPONENT = Base::FP8_MAX_EXPONENT;
+
+    //
+    // Static conversion operators
+    //
+
+    /// Constructs from an uint8_t
+    CUTLASS_HOST_DEVICE
+    static float_e4m3_t bitcast(uint8_t x) {
+        float_e4m3_t f;
+        f.storage = x;
+        return f;
+    }
+
+    /// FP32 -> FP8 conversion - rounds to nearest even
+    CUTLASS_HOST_DEVICE
+    static float_e4m3_t from_float(float const& flt) {
+    #if defined(CUDA_PTX_FP8_CVT_ENABLED)
+        uint16_t tmp;
+        float y = float();
+        asm volatile("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(tmp) : "f"(y), "f"(flt));
+
+        return *reinterpret_cast<float_e4m3_t *>(&tmp);
+    #else
+        return bitcast(Base::convert_float_to_fp8(flt));
+    #endif
+    }
+
+    /// FP16 -> E5M2 conversion - rounds to nearest even
+    CUTLASS_HOST_DEVICE
+    static float_e4m3_t from_half(half const& flt) {
+    #if defined(CUDA_PTX_FP8_CVT_ENABLED)
+        uint16_t tmp = 0;
+        uint32_t bits = reinterpret_cast<uint16_t const &>(flt);
+        asm volatile("cvt.rn.satfinite.e4m3x2.f16x2 %0, %1;" : "=h"(tmp) : "r"(bits));
+
+        return *reinterpret_cast<float_e4m3_t *>(&tmp);
+    #else
+        return bitcast(Base::convert_float_to_fp8(__half2float(flt)));
+    #endif
+    }
+
+    // E4M3 -> half
+    CUTLASS_HOST_DEVICE
+    static half to_half(float_e4m3_t const& x) {
+    #if defined(CUDA_PTX_FP8_CVT_ENABLED)
+        uint16_t bits = x.storage;
+        uint32_t packed;
+        asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;\n" : "=r"(packed) : "h"(bits));
+
+        return reinterpret_cast<half2 const &>(packed).x;
+    #else
+        return __float2half(Base::convert_fp8_to_float(x.storage));
+    #endif
+    }
+
+    // E4M3 -> Float
+    CUTLASS_HOST_DEVICE
+    static float to_float(float_e4m3_t const& x) {
+    #if defined(CUDA_PTX_FP8_CVT_ENABLED)
+        uint16_t bits = x.storage;
+        uint32_t packed;
+        asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;\n" : "=r"(packed) : "h"(bits));
+
+        return __half2float(reinterpret_cast<half2 const &>(packed).x);
+    #else
+        return Base::convert_fp8_to_float(x.storage);
+    #endif
+    }
+
+    //
+    // Methods
+    //
+
+    /// Constructor inheritance
+    using Base::Base;
+
+    /// Default constructor
+    float_e4m3_t() = default;
+
+#ifdef CUDA_FP8_ENABLED
+    /// Conversion from CUDA's FP8 type
+    CUTLASS_HOST_DEVICE
+    explicit float_e4m3_t(__nv_fp8_e4m3 x) {
+        storage = x.__x;
+    }
+#endif
+
+    /// Floating point conversion
+    CUTLASS_HOST_DEVICE
+    explicit float_e4m3_t(float x) {
+        storage = from_float(x).storage;
+    }
+
+    CUTLASS_HOST_DEVICE
+    explicit float_e4m3_t(half x) {
+        storage = from_half(x).storage;
+    }
+
+    /// Floating point conversion
+    CUTLASS_HOST_DEVICE
+    explicit float_e4m3_t(double x): float_e4m3_t(float(x)) {
+    }
+
+    /// Integer conversion
+    CUTLASS_HOST_DEVICE
+    explicit float_e4m3_t(int x): float_e4m3_t(float(x)) {
+    }
+
+    CUTLASS_HOST_DEVICE
+    explicit float_e4m3_t(unsigned x): float_e4m3_t(float(x)) {
+    }
+
+    /// E5M2 conversion. Defined after float_e5m2_t is defined.
+    CUTLASS_HOST_DEVICE
+    explicit float_e4m3_t(float_e5m2_t x);
+
+#ifdef CUDA_FP8_ENABLED
+    /// Assignment from CUDA's FP8 type
+    CUTLASS_HOST_DEVICE
+    float_e4m3_t & operator=(__nv_fp8_e4m3 x) {
+        storage = x.__x;
+        return *this;
+    }
+#endif
+
+    /// Converts to float
+    CUTLASS_HOST_DEVICE
+    operator float() const {
+        return to_float(*this);
+    }
+
+    /// Converts to half
+    CUTLASS_HOST_DEVICE
+    operator half() const {
+        return to_half(*this);
+    }
+
+    /// Converts to float
+    CUTLASS_HOST_DEVICE
+    explicit operator double() const {
+        return double(to_float(*this));
+    }
+
+    /// Converts to int
+    CUTLASS_HOST_DEVICE
+    explicit operator int() const {
+    #if defined(__CUDA_ARCH__)
+        return __half2int_rn(to_half(*this));
+    #else
+        return int(to_float(*this));
+    #endif
+    }
+
+    /// Casts to bool
+    CUTLASS_HOST_DEVICE
+    explicit operator bool() const {
+    #if defined(__CUDA_ARCH__)
+        return bool(__half2int_rn(to_half(*this)));
+    #else
+        return bool(int(to_float(*this)));
+    #endif
+    }
+
+    /// Accesses raw internal state
+    CUTLASS_HOST_DEVICE
+    uint8_t& raw() {
+        return storage;
+    }
+
+    /// Accesses raw internal state
+    CUTLASS_HOST_DEVICE
+    uint8_t raw() const {
+        return storage;
+    }
+
+    /// Returns the sign bit
+    CUTLASS_HOST_DEVICE
+    bool signbit() const {
+        return ((storage & (1 << (Base::FP8_NUM_BITS - 1))) != 0);
+    }
+
+    /// Returns the biased exponent
+    CUTLASS_HOST_DEVICE
+    int exponent_biased() const {
+        return int((storage >> FP8_NUM_MANTISSA_BITS) & Base::FP8_EXPONENT_MASK);
+    }
+
+    /// Returns the unbiased exponent
+    CUTLASS_HOST_DEVICE
+    int exponent() const {
+        return exponent_biased() - 15;
+    }
+
+    /// Returns the mantissa
+    CUTLASS_HOST_DEVICE
+    int mantissa() const {
+        return int(storage & Base::FP8_MANTISSA_MASK);
+    }
+
+    CUTLASS_HOST_DEVICE
+    friend bool isnan(float_e4m3_t const& x) {
+      return x.storage == uint8_t(0x7f);
+    }
+
+};
+///////////////////////////////////////////////////////////////
+///
+/// floating-point 8 type : E5M2
+///
+///////////////////////////////////////////////////////////////
+struct alignas(1) float_e5m2_t : float8_base<FloatEncoding::E5M2> {
+
+    using Base = float8_base<FloatEncoding::E5M2>;
+
+    static constexpr int MAX_EXPONENT = Base::FP8_MAX_EXPONENT;
+
+    //
+    // Static conversion operators
+    //
+
+    /// Constructs from an uint8_t
+    CUTLASS_HOST_DEVICE
+    static float_e5m2_t bitcast(uint8_t x) {
+        float_e5m2_t f;
+        f.storage = x;
+        return f;
+    }
+
+    /// FP32 -> FP8 conversion - rounds to nearest even
+    CUTLASS_HOST_DEVICE
+    static float_e5m2_t from_float(float const& flt) {
+    #if defined(CUDA_PTX_FP8_CVT_ENABLED)
+        uint16_t tmp;
+        float y = float();
+        asm volatile("cvt.rn.satfinite.e5m2x2.f32 %0, %1, %2;" : "=h"(tmp) : "f"(y), "f"(flt));
+
+        return *reinterpret_cast<float_e5m2_t *>(&tmp);
+    #else
+        return bitcast(Base::convert_float_to_fp8(flt));
+    #endif
+    }
+
+    /// FP16 -> E5M2 conversion - rounds to nearest even
+    CUTLASS_HOST_DEVICE
+    static float_e5m2_t from_half(half const& flt) {
+    #if defined(CUDA_PTX_FP8_CVT_ENABLED)
+        uint16_t tmp = 0;
+        uint32_t bits = reinterpret_cast<uint16_t const &>(flt);
+        asm volatile("cvt.rn.satfinite.e5m2x2.f16x2 %0, %1;" : "=h"(tmp) : "r"(bits));
+
+        return *reinterpret_cast<float_e5m2_t *>(&tmp);
+    #else
+        return bitcast(Base::convert_float_to_fp8(__half2float(flt)));
+    #endif
+    }
+
+    // E5M2 -> half
+    CUTLASS_HOST_DEVICE
+    static half to_half(float_e5m2_t const& x) {
+    #if defined(CUDA_PTX_FP8_CVT_ENABLED)
+        uint16_t bits = x.storage;
+        uint32_t packed;
+        asm volatile("cvt.rn.f16x2.e5m2x2 %0, %1;\n" : "=r"(packed) : "h"(bits));
+
+        return reinterpret_cast<half2 const &>(packed).x;
+    #else
+        return __float2half(Base::convert_fp8_to_float(x.storage));
+    #endif
+    }
+
+    // E5M2 -> Float
+    CUTLASS_HOST_DEVICE
+    static float to_float(float_e5m2_t const& x) {
+    #if defined(CUDA_PTX_FP8_CVT_ENABLED)
+        uint16_t bits = x.storage;
+        uint32_t packed;
+        asm volatile("cvt.rn.f16x2.e5m2x2 %0, %1;\n" : "=r"(packed) : "h"(bits));
+
+        return __half2float(reinterpret_cast<half2 const &>(packed).x);
+    #else
+        return Base::convert_fp8_to_float(x.storage);
+    #endif
+    }
+
+    //
+    // Methods
+    //
+
+    /// Constructor inheritance
+    using Base::Base;
+
+    /// Default constructor
+    float_e5m2_t() = default;
+
+#ifdef CUDA_FP8_ENABLED
+    /// Conversion from CUDA's FP8 type
+    CUTLASS_HOST_DEVICE
+    explicit float_e5m2_t(__nv_fp8_e5m2 x) {
+        storage = x.__x;
+    }
+#endif
+
+    /// Floating point conversion
+    CUTLASS_HOST_DEVICE
+    explicit float_e5m2_t(float x) {
+        storage = from_float(x).storage;
+    }
+
+    CUTLASS_HOST_DEVICE
+    explicit float_e5m2_t(half x) {
+      storage = from_half(x).storage;
+    }
+
+    /// Floating point conversion
+    CUTLASS_HOST_DEVICE
+    explicit float_e5m2_t(double x): float_e5m2_t(float(x)) {
+    }
+
+    /// Integer conversion
+    CUTLASS_HOST_DEVICE
+    explicit float_e5m2_t(int x): float_e5m2_t(float(x)) {
+    }
+
+    CUTLASS_HOST_DEVICE
+    explicit float_e5m2_t(unsigned x): float_e5m2_t(float(x)) {
+    }
+
+    /// E4M3 conversion
+    CUTLASS_HOST_DEVICE
+    explicit float_e5m2_t(float_e4m3_t x);
+
+#ifdef CUDA_FP8_ENABLED
+    /// Assignment from CUDA's FP8 type
+    CUTLASS_HOST_DEVICE
+    float_e5m2_t & operator=(__nv_fp8_e5m2 x) {
+        storage = x.__x;
+        return *this;
+    }
+#endif
+
+    /// Converts to float
+    CUTLASS_HOST_DEVICE
+    operator float() const {
+        return to_float(*this);
+    }
+
+    /// Converts to half
+    CUTLASS_HOST_DEVICE
+    operator half() const {
+      return to_half(*this);
+    }
+
+    /// Converts to float
+    CUTLASS_HOST_DEVICE
+    explicit operator double() const {
+        return double(to_float(*this));
+    }
+
+    /// Converts to int
+    CUTLASS_HOST_DEVICE
+    explicit operator int() const {
+    #if defined(__CUDA_ARCH__)
+        return __half2int_rn(to_half(*this));
+    #else
+        return int(to_float(*this));
+    #endif
+    }
+
+    /// Casts to bool
+    CUTLASS_HOST_DEVICE
+    explicit operator bool() const {
+    #if defined(__CUDA_ARCH__)
+        return bool(__half2int_rn(to_half(*this)));
+    #else
+        return bool(int(to_float(*this)));
+    #endif
+    }
+
+    /// Accesses raw internal state
+    CUTLASS_HOST_DEVICE
+    uint8_t& raw() {
+        return storage;
+    }
+
+    /// Accesses raw internal state
+    CUTLASS_HOST_DEVICE
+    uint8_t raw() const {
+        return storage;
+    }
+
+    /// Returns the sign bit
+    CUTLASS_HOST_DEVICE
+    bool signbit() const {
+        return ((storage & (1 << (Base::FP8_NUM_BITS - 1))) != 0);
+    }
+
+    /// Returns the biased exponent
+    CUTLASS_HOST_DEVICE
+    int exponent_biased() const {
+        return int((storage >> FP8_NUM_MANTISSA_BITS) & Base::FP8_EXPONENT_MASK);
+    }
+
+    /// Returns the unbiased exponent
+    CUTLASS_HOST_DEVICE
+    int exponent() const {
+        return exponent_biased() - 15;
+    }
+
+    /// Returns the mantissa
+    CUTLASS_HOST_DEVICE
+    int mantissa() const {
+        return int(storage & Base::FP8_MANTISSA_MASK);
+    }
+    
+    CUTLASS_HOST_DEVICE
+    friend bool isnan(float_e5m2_t const& x) {
+      return x.storage == uint8_t(0x7f);
+    }
+
+};
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Arithmetic operators
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+CUTLASS_HOST_DEVICE
+bool operator==(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float(lhs) == float(rhs);
+}
+
+CUTLASS_HOST_DEVICE
+bool operator!=(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float(lhs) != float(rhs);
+}
+
+CUTLASS_HOST_DEVICE
+bool operator<(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float(lhs) < float(rhs);
+}
+
+CUTLASS_HOST_DEVICE
+bool operator<=(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float(lhs) <= float(rhs);
+}
+
+CUTLASS_HOST_DEVICE
+bool operator>(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float(lhs) > float(rhs);
+}
+
+CUTLASS_HOST_DEVICE
+bool operator>=(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float(lhs) >= float(rhs);
+}
+
+CUTLASS_HOST_DEVICE
+float_e4m3_t operator+(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float_e4m3_t(float(lhs) + float(rhs));
+}
+
+CUTLASS_HOST_DEVICE
+float_e4m3_t operator-(float_e4m3_t const& lhs) {
+    return float_e4m3_t(-float(lhs));
+}
+
+CUTLASS_HOST_DEVICE
+float_e4m3_t operator-(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float_e4m3_t(float(lhs) - float(rhs));
+}
+
+CUTLASS_HOST_DEVICE
+float_e4m3_t operator*(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float_e4m3_t(float(lhs) * float(rhs));
+}
+
+CUTLASS_HOST_DEVICE
+float_e4m3_t operator/(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float_e4m3_t(float(lhs) / float(rhs));
+}
+
+CUTLASS_HOST_DEVICE
+float_e4m3_t& operator+=(float_e4m3_t & lhs, float_e4m3_t const& rhs) {
+    lhs = float_e4m3_t(float(lhs) + float(rhs));
+    return lhs;
+}
+
+CUTLASS_HOST_DEVICE
+float_e4m3_t& operator-=(float_e4m3_t & lhs, float_e4m3_t const& rhs) {
+    lhs = float_e4m3_t(float(lhs) - float(rhs));
+    return lhs;
+}
+
+CUTLASS_HOST_DEVICE
+float_e4m3_t& operator*=(float_e4m3_t & lhs, float_e4m3_t const& rhs) {
+    lhs = float_e4m3_t(float(lhs) * float(rhs));
+    return lhs;
+}
+
+CUTLASS_HOST_DEVICE
+float_e4m3_t& operator/=(float_e4m3_t & lhs, float_e4m3_t const& rhs) {
+    lhs = float_e4m3_t(float(lhs) / float(rhs));
+    return lhs;
+}
+
+CUTLASS_HOST_DEVICE
+float_e4m3_t& operator++(float_e4m3_t & lhs) {
+    float tmp(lhs);
+    ++tmp;
+    lhs = float_e4m3_t(tmp);
+    return lhs;
+}
+
+CUTLASS_HOST_DEVICE
+float_e4m3_t& operator--(float_e4m3_t & lhs) {
+    float tmp(lhs);
+    --tmp;
+    lhs = float_e4m3_t(tmp);
+    return lhs;
+}
+
+CUTLASS_HOST_DEVICE
+float_e4m3_t operator++(float_e4m3_t & lhs, int) {
+    float_e4m3_t ret(lhs);
+    float tmp(lhs);
+    tmp++;
+    lhs = float_e4m3_t(tmp);
+    return ret;
+}
+
+CUTLASS_HOST_DEVICE
+float_e4m3_t operator--(float_e4m3_t & lhs, int) {
+    float_e4m3_t ret(lhs);
+    float tmp(lhs);
+    tmp--;
+    lhs = float_e4m3_t(tmp);
+    return ret;
+}
+
+CUTLASS_HOST_DEVICE
+bool operator==(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float(lhs) == float(rhs);
+}
+
+CUTLASS_HOST_DEVICE
+bool operator!=(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float(lhs) != float(rhs);
+}
+
+CUTLASS_HOST_DEVICE
+bool operator<(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float(lhs) < float(rhs);
+}
+
+CUTLASS_HOST_DEVICE
+bool operator<=(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float(lhs) <= float(rhs);
+}
+
+CUTLASS_HOST_DEVICE
+bool operator>(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float(lhs) > float(rhs);
+}
+
+CUTLASS_HOST_DEVICE
+bool operator>=(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float(lhs) >= float(rhs);
+}
+
+CUTLASS_HOST_DEVICE
+float_e5m2_t operator+(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float_e5m2_t(float(lhs) + float(rhs));
+}
+
+CUTLASS_HOST_DEVICE
+float_e5m2_t operator-(float_e5m2_t const& lhs) {
+    return float_e5m2_t(-float(lhs));
+}
+
+CUTLASS_HOST_DEVICE
+float_e5m2_t operator-(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float_e5m2_t(float(lhs) - float(rhs));
+}
+
+CUTLASS_HOST_DEVICE
+float_e5m2_t operator*(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float_e5m2_t(float(lhs) * float(rhs));
+}
+
+CUTLASS_HOST_DEVICE
+float_e5m2_t operator/(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float_e5m2_t(float(lhs) / float(rhs));
+}
+
+CUTLASS_HOST_DEVICE
+float_e5m2_t& operator+=(float_e5m2_t & lhs, float_e5m2_t const& rhs) {
+    lhs = float_e5m2_t(float(lhs) + float(rhs));
+    return lhs;
+}
+
+CUTLASS_HOST_DEVICE
+float_e5m2_t& operator-=(float_e5m2_t & lhs, float_e5m2_t const& rhs) {
+    lhs = float_e5m2_t(float(lhs) - float(rhs));
+    return lhs;
+}
+
+CUTLASS_HOST_DEVICE
+float_e5m2_t& operator*=(float_e5m2_t & lhs, float_e5m2_t const& rhs) {
+    lhs = float_e5m2_t(float(lhs) * float(rhs));
+    return lhs;
+}
+
+CUTLASS_HOST_DEVICE
+float_e5m2_t& operator/=(float_e5m2_t & lhs, float_e5m2_t const& rhs) {
+    lhs = float_e5m2_t(float(lhs) / float(rhs));
+    return lhs;
+}
+
+CUTLASS_HOST_DEVICE
+float_e5m2_t& operator++(float_e5m2_t & lhs) {
+    float tmp(lhs);
+    ++tmp;
+    lhs = float_e5m2_t(tmp);
+    return lhs;
+}
+
+CUTLASS_HOST_DEVICE
+float_e5m2_t& operator--(float_e5m2_t & lhs) {
+    float tmp(lhs);
+    --tmp;
+    lhs = float_e5m2_t(tmp);
+    return lhs;
+}
+
+CUTLASS_HOST_DEVICE
+float_e5m2_t operator++(float_e5m2_t & lhs, int) {
+    float_e5m2_t ret(lhs);
+    float tmp(lhs);
+    tmp++;
+    lhs = float_e5m2_t(tmp);
+    return ret;
+}
+
+CUTLASS_HOST_DEVICE
+float_e5m2_t operator--(float_e5m2_t & lhs, int) {
+    float_e5m2_t ret(lhs);
+    float tmp(lhs);
+    tmp--;
+    lhs = float_e5m2_t(tmp);
+    return ret;
+}
+
+
+///////////////////////////////////////////////////////////////
+///
+/// floating-point 8 type : UE4M3
+///
+///////////////////////////////////////////////////////////////
+// UE4M3:
+//   4 Exponent bits, 3 Mantissa bits
+//   Range: [0:448]
+//   has_inf: false
+//   has_NaN: true
+//   has_denorm: true
+//   Exponent bias (exp_bias): 7
+struct float_ue4m3_t : public float_exmy_base<cutlass::detail::FpEncoding::UE4M3, float_ue4m3_t> {
+  using Base = float_exmy_base<cutlass::detail::FpEncoding::UE4M3, float_ue4m3_t>;
+
+  float_ue4m3_t() = default;
+
+  CUTLASS_HOST_DEVICE
+  float_ue4m3_t convert_from_float(float const &flt) const {
+    #if defined(CUDA_PTX_FP8_CVT_ENABLED)
+      uint16_t tmp;
+      float y = float();
+      asm volatile("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(tmp) : "f"(y), "f"(flt));
+      return bitcast(*reinterpret_cast<uint8_t *>(&tmp));
+    #else 
+      Base::FP32BitRepresentation::Storage fp32_bits = Base::FP32BitRepresentation::to_bits(flt);
+      return bitcast(BitRepresentation::convert_from(fp32_bits, Base::FP32BitRepresentation{}));
+    #endif
+  }
+
+  CUTLASS_HOST_DEVICE
+  float convert_to_float(float_ue4m3_t const &x) const {
+    #if defined(CUDA_PTX_FP8_CVT_ENABLED)
+      uint16_t bits = x.storage;
+      uint32_t packed;
+      asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;\n" : "=r"(packed) : "h"(bits));
+      return __half2float(reinterpret_cast<half2 const &>(packed).x);
+    #else 
+      Base::FP32BitRepresentation::Storage fp32_bits;
+      fp32_bits = Base::BitRepresentation::convert_to(x.storage, Base::FP32BitRepresentation{});
+      return detail::copy_bits<Base::FP32BitRepresentation::Storage, float>(fp32_bits);
+    #endif
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_ue4m3_t(double x) : Base(float(x)) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_ue4m3_t(float x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_ue4m3_t(int x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_ue4m3_t(unsigned x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  float_ue4m3_t(Base x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend bool isnan(float_ue4m3_t const& x) {
+    return x.storage == uint8_t(0x7f);
+  }
+
+};
+
+/// Defines the size of an element in bits - specialized for float_ue4m3_t
+template <>
+struct sizeof_bits<float_ue4m3_t> {
+  static constexpr int value = sizeof_bits<float_exmy_base<cutlass::detail::FpEncoding::UE4M3, float_ue4m3_t>>::value;
+};
+
+
+
+///////////////////////////////////////////////////////////////
+///
+/// floating-point 8 type : UE8M0
+///
+///////////////////////////////////////////////////////////////
+// UE8M0:
+//   8 Exponent bits, 0 Mantissa bits
+//   Range: [2^-127:2^127]
+//   has_inf: false
+//   has_NaN: true (11111111)
+//   has_denorm: true
+//   Exponent bias (exp_bias): 8
+
+struct float_ue8m0_t : public float_exmy_base<cutlass::detail::FpEncoding::UE8M0, float_ue8m0_t> {
+  using Base = float_exmy_base<cutlass::detail::FpEncoding::UE8M0, float_ue8m0_t>;
+  using FP32Bits = typename Base::FP32BitRepresentation;
+
+  float_ue8m0_t() = default;
+
+  CUTLASS_HOST_DEVICE
+  float_ue8m0_t convert_from_float(float const &flt) const {
+  #if defined(CUDA_PTX_UE8M0_CVT_ENABLED)
+    uint16_t out;
+    asm volatile(
+        "{ cvt.rp.satfinite.ue8m0x2.f32 %0, 0.0, %1; }"
+        : "=h"(out) : "f"(flt));      
+    return bitcast(*reinterpret_cast<uint8_t *>(&out));
+  #else
+    if (CUTLASS_CMATH_NAMESPACE::isnan(flt) || CUTLASS_CMATH_NAMESPACE::isinf(flt)) {
+      return bitcast(0xFF);
+    }
+    uint32_t flt_uint32 = cutlass::detail::copy_bits<float, uint32_t>(flt);
+    uint8_t exp = (flt_uint32 >> 23) & 0xff;  // Extract the 8 bit exponent
+    uint32_t mant = flt_uint32 & 0x7fffff;    // Extract the 23 bit mantissa
+    // Do the round up
+    // Deals w/ satfinite all at once
+    if ((mant > 0) && (exp != 0xFE) && !(exp == 0 && mant <= 0x00400000)) {
+      exp++;
+    }
+    return bitcast(exp);
+  #endif
+  }
+
+  CUTLASS_HOST_DEVICE
+  float convert_to_float(float_ue8m0_t const &x) const {
+    //////////////////////////////////////////////////////////////
+    // The conversion of UE8M0 to FP32 scale can be done simply
+    // with a left shift (No rounding necessary)
+    // Note: The base class implements ue8m0 to FP32 based on the rules of float math conversions.
+    //       The result of current implementation and base class are aligned.
+    //////////////////////////////////////////////////////////////
+    #if defined(CUDA_PTX_UE8M0_CVT_ENABLED)
+      uint16_t bits = x.storage;
+      uint32_t bf16x2_val;
+      // E8 -> BF16
+      asm volatile(
+        "{\n"
+        "cvt.rn.bf16x2.ue8m0x2 %0, %1;\n"
+        "}\n" : "=r"(bf16x2_val): "h"(bits));
+      // BF16 -> FP32
+      float f1;
+      asm(
+        "{\n"
+        "prmt.b32 %0, %1, %2, %3;\n"
+        "}\n"
+          : "=f"(f1) 
+          : "r"(0), "r"(bf16x2_val), "r"(0x5410));
+      return f1;
+    #else
+      using FP32Bits = cutlass::detail::FpBitRepresentation<uint32_t, 32, 8, 23, cutlass::detail::NanInfEncoding::IEEE_754>;
+      if (x.storage == 0x00) {
+        return cutlass::detail::copy_bits<FP32Bits::Storage, float>(0x00400000);
+      }
+      else if (x.storage == 0xFF) {
+        return cutlass::detail::copy_bits<FP32Bits::Storage, float>(0x7fffffff);
+      }
+      else {
+        auto f8 = static_cast<FP32Bits::Storage>(x.storage);
+        FP32Bits::Storage f = (f8 << FP32Bits::NUM_MANTISSA_BITS);
+        return cutlass::detail::copy_bits<FP32Bits::Storage, float>(f);
+      }
+    #endif
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_ue8m0_t(double x) : Base(float(x)) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_ue8m0_t(float x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_ue8m0_t(int x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_ue8m0_t(unsigned x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  float_ue8m0_t(Base x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  friend bool isnan(float_ue8m0_t const& x) {
+    return x.storage == uint8_t(0xff);
+  }
+
+};
+
+/// Defines the size of an element in bits - specialized for float_ue8m0_t
+template <>
+struct sizeof_bits<float_ue8m0_t> {
+  static constexpr int value = sizeof_bits<float_exmy_base<cutlass::detail::FpEncoding::UE8M0, float_ue8m0_t>>::value;
+};
+
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// float_e4m3_t <=> float_e5m2_t conversions
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// float_e4m3_t <= float_e5m2_t
+CUTLASS_HOST_DEVICE
+float_e4m3_t::float_e4m3_t(float_e5m2_t x) {
+    storage = from_float(float_e5m2_t::to_float(x)).storage;
+}
+
+/// float_e5m2_t <= float_e4m3_t
+CUTLASS_HOST_DEVICE
+float_e5m2_t::float_e5m2_t(float_e4m3_t x) {
+    storage = from_float(float_e4m3_t::to_float(x)).storage;
+}
+
+///////////////////////////////////////////////////////////////
+///
+/// Umbrella floating-point 8-bit data type : type_erased_dynamic_float8_t
+/// This umbrella datatype can be enabled when a user provides a specific
+/// datatype in runtime argument list.
+///
+/// Currently supported runtime datatypes compatible with type_erased_dynamic_float8_t:
+///   MXF8F6F4Format::E5M2
+///   MXF8F6F4Format::E4M3
+///
+///////////////////////////////////////////////////////////////
+
+union type_erased_dynamic_float8_t {
+  uint8_t data;
+  cutlass::float_e5m2_t e5m2;
+  cutlass::float_e4m3_t e4m3;
+  CUTLASS_HOST_DEVICE
+  explicit operator cutlass::float_e5m2_t() const {
+    return e5m2;
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit operator cutlass::float_e4m3_t() const {
+    return e4m3;
+  }
+
+};
+
+
+
+///////////////////////////////////////////////////////////////
+/// MX type for float8
+/// Intended to be used in builders
+///////////////////////////////////////////////////////////////
+
+template <class F8Type>
+struct mx_float8_t {
+  static_assert(cute::is_same_v<F8Type,cutlass::float_e5m2_t>
+                || cute::is_same_v<F8Type,cutlass::float_e4m3_t>
+                || cute::is_same_v<F8Type,type_erased_dynamic_float8_t>
+                , "Only float_e5m2_t, float_e4m3_t can have scale factors for MXFP8");
+  using ScaleFactorType = cutlass::float_ue8m0_t;
+  using DataType = F8Type;
+};
+
+using type_erased_dynamic_mx_float8_t = mx_float8_t<type_erased_dynamic_float8_t>;
+
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace cutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Standard Library operations and definitions
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+#if !defined(__CUDACC_RTC__)
+namespace std {
+
+/// Numeric limits common to all float8 types
+template <typename T>
+struct float8_base_numeric_limits {
+private:
+  using F8Type = T;
+public:
+  static bool const is_specialized = true;
+  static bool const is_signed = true;
+  static bool const is_integer = false;
+  static bool const is_exact = false;
+  static bool const has_quiet_NaN = true;
+  static bool const has_signaling_NaN = false;
+  static std::float_denorm_style const has_denorm = std::denorm_present;
+  static bool const has_denorm_loss = true;
+  static std::float_round_style const round_style = std::round_to_nearest;
+  static bool const is_iec559 = false;
+  static bool const is_bounded = true;
+  static bool const is_modulo = false;
+  static int const digits = F8Type::FP8_NUM_MANTISSA_BITS;
+
+  /// Least positive value
+  CUTLASS_HOST_DEVICE
+  static F8Type min() { return F8Type::bitcast(0x01); }
+
+  /// Maximum finite value
+  CUTLASS_HOST_DEVICE
+  static F8Type max() { return F8Type::bitcast(F8Type::FP8_MAX_FLT); }
+
+  /// Returns maximum rounding error
+  CUTLASS_HOST_DEVICE
+  static F8Type round_error() { return F8Type(0.5f); }
+
+  /// Returns positive infinity value
+  CUTLASS_HOST_DEVICE
+  static F8Type infinity() { return F8Type::bitcast(F8Type::FP8_INFINITY_MASK); }
+
+  /// Returns quiet NaN value
+  CUTLASS_HOST_DEVICE
+  static F8Type quiet_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); }
+
+  /// Returns signaling NaN value
+  CUTLASS_HOST_DEVICE
+  static F8Type signaling_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); }
+
+  /// Returns smallest positive subnormal value
+  CUTLASS_HOST_DEVICE
+  static F8Type denorm_min() { return F8Type::bitcast(0x01); }
+};
+
+/// Numeric limits for float_e4m3_t
+template <>
+struct numeric_limits<cutlass::float_e4m3_t> :
+    public float8_base_numeric_limits<cutlass::float_e4m3_t> {
+  static bool const has_infinity = false;
+
+  /// Minimum finite value
+  static cutlass::float_e4m3_t lowest() { return cutlass::float_e4m3_t::bitcast(0xfe); }
+
+  /// Machine epsilon, that is, the difference between 1.0 and the next representable value
+  static cutlass::float_e4m3_t epsilon() { return cutlass::float_e4m3_t::bitcast(0x20); }
+};
+
+/// Numeric limits for float_e5m2_t
+template <>
+struct numeric_limits<cutlass::float_e5m2_t>  :
+    public float8_base_numeric_limits<cutlass::float_e5m2_t> {
+  static bool const has_infinity = true;
+
+  /// Minimum finite value
+  static cutlass::float_e5m2_t lowest() { return cutlass::float_e5m2_t::bitcast(0xfb); }
+
+  /// Machine epsilon, that is, the difference between 1.0 and the next representable value
+  static cutlass::float_e5m2_t epsilon() { return cutlass::float_e5m2_t::bitcast(0x34); }
+};
+
+
+template <typename T>
+struct float8_exmy_numeric_limits
+{
+private:
+  using type = T;
+
+public:
+  static bool const is_specialized = true;
+  static bool const is_signed = true;
+  static bool const is_integer = false;
+  static bool const is_exact = false;
+  static bool const has_quiet_NaN = true;
+  static bool const has_signaling_NaN = false;
+  static bool const has_denorm_loss = true;
+  static cutlass::platform::float_denorm_style const has_denorm = cutlass::platform::denorm_present;
+  static cutlass::platform::float_round_style const round_style = cutlass::platform::round_to_nearest;
+  static bool const is_iec559 = false;
+  static bool const is_bounded = true;
+  static bool const is_modulo = false;
+  static int const digits = type::Base::BitRepresentation::NUM_MANTISSA_BITS;
+  static bool const has_infinity = false;
+
+  /// Least positive value
+  CUTLASS_HOST_DEVICE
+  static type min() { return type::bitcast(0x01); }
+
+  /// Maximum finite value
+  CUTLASS_HOST_DEVICE
+  static type max() { return type::bitcast(type::Base::BitRepresentation::MAX_VALUE); }
+
+  /// Returns maximum rounding error
+  CUTLASS_HOST_DEVICE
+  static type round_error() { return type(0.5f); }
+
+  /// Returns positive infinity value
+  CUTLASS_HOST_DEVICE
+  static type infinity() { return type::bitcast(type::Base::BitRepresentation::INF_MASK); }
+
+  /// Returns quiet NaN value
+  CUTLASS_HOST_DEVICE
+  static type quiet_NaN() { return type::bitcast(type::Base::BitRepresentation::INF_MASK); }
+
+  /// Returns signaling NaN value
+  CUTLASS_HOST_DEVICE
+  static type signaling_NaN() { return type::bitcast(type::Base::BitRepresentation::INF_MASK); }
+
+  /// Returns smallest positive subnormal value
+  CUTLASS_HOST_DEVICE
+  static type denorm_min() { return type::bitcast(0x01); }
+};
+
+/// Numeric limits for float_ue8m0_t
+template <>
+struct numeric_limits<cutlass::float_ue8m0_t> :
+    public float8_exmy_numeric_limits<cutlass::float_ue8m0_t> {
+  static bool const has_infinity = false;
+  static bool const is_signed = false;
+
+  /// Minimum finite value
+  static cutlass::float_ue8m0_t lowest() { return cutlass::float_ue8m0_t::bitcast(0xfe); }
+
+  /// Machine epsilon, that is, the difference between 1.0 and the next representable value (2^0)
+  static cutlass::float_ue8m0_t epsilon() { return cutlass::float_ue8m0_t::bitcast(0x7f); }
+};
+
+
+}  // namespace std
+#endif
+
+namespace cutlass {
+namespace platform {
+
+/// Numeric limits common to all float8 types
+template <typename T>
+struct float8_base_numeric_limits {
+private:
+  using F8Type = T;
+public:
+  static bool const is_specialized = true;
+  static bool const is_signed = true;
+  static bool const is_integer = false;
+  static bool const is_exact = false;
+  static bool const has_quiet_NaN = true;
+  static bool const has_signaling_NaN = false;
+#if !defined(__CUDACC_RTC__)
+  static std::float_denorm_style const has_denorm = std::denorm_present;
+#endif
+  static bool const has_denorm_loss = true;
+#if !defined(__CUDACC_RTC__)
+  static std::float_round_style const round_style = std::round_to_nearest;
+#endif
+  static bool const is_iec559 = false;
+  static bool const is_bounded = true;
+  static bool const is_modulo = false;
+  static int const digits = F8Type::FP8_NUM_MANTISSA_BITS;
+
+  /// Least positive value
+  CUTLASS_HOST_DEVICE
+  static F8Type min() { return F8Type::bitcast(0x01); }
+
+  /// Maximum finite value
+  CUTLASS_HOST_DEVICE
+  static F8Type max() { return F8Type::bitcast(F8Type::FP8_MAX_FLT); }
+
+  /// Returns maximum rounding error
+  CUTLASS_HOST_DEVICE
+  static F8Type round_error() { return F8Type(0.5f); }
+
+  /// Returns positive infinity value
+  CUTLASS_HOST_DEVICE
+  static F8Type infinity() { return F8Type::bitcast(F8Type::FP8_INFINITY_MASK); }
+
+  /// Returns quiet NaN value
+  CUTLASS_HOST_DEVICE
+  static F8Type quiet_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); }
+
+  /// Returns signaling NaN value
+  CUTLASS_HOST_DEVICE
+  static F8Type signaling_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); }
+
+  /// Returns smallest positive subnormal value
+  CUTLASS_HOST_DEVICE
+  static F8Type denorm_min() { return F8Type::bitcast(0x01); }
+};
+
+/// Forward Declaration
+template <class T>
+struct numeric_limits;
+
+/// Numeric limits for float_e4m3_t
+template <>
+struct numeric_limits<cutlass::float_e4m3_t> :
+    public float8_base_numeric_limits<cutlass::float_e4m3_t> {
+  static bool const has_infinity = false;
+
+  /// Minimum finite value
+  static cutlass::float_e4m3_t lowest() { return cutlass::float_e4m3_t::bitcast(0xfe); }
+
+  /// Machine epsilon, that is, the difference between 1.0 and the next representable value
+  static cutlass::float_e4m3_t epsilon() { return cutlass::float_e4m3_t::bitcast(0x20); }
+};
+
+/// Numeric limits for float_e5m2_t
+template <>
+struct numeric_limits<cutlass::float_e5m2_t>  :
+    public float8_base_numeric_limits<cutlass::float_e5m2_t> {
+  static bool const has_infinity = true;
+
+  /// Minimum finite value
+  static cutlass::float_e5m2_t lowest() { return cutlass::float_e5m2_t::bitcast(0xfb); }
+
+  /// Machine epsilon, that is, the difference between 1.0 and the next representable value
+  static cutlass::float_e5m2_t epsilon() { return cutlass::float_e5m2_t::bitcast(0x34); }
+};
+
+
+template <typename T>
+struct float8_exmy_numeric_limits
+{
+private:
+  using type = T;
+
+public:
+  static bool const is_specialized = true;
+  static bool const is_signed = true;
+  static bool const is_integer = false;
+  static bool const is_exact = false;
+  static bool const has_quiet_NaN = true;
+  static bool const has_signaling_NaN = false;
+  static bool const has_denorm_loss = true;
+  static cutlass::platform::float_denorm_style const has_denorm = cutlass::platform::denorm_present;
+  static cutlass::platform::float_round_style const round_style = cutlass::platform::round_to_nearest;
+  static bool const is_iec559 = false;
+  static bool const is_bounded = true;
+  static bool const is_modulo = false;
+  static int const digits = type::Base::BitRepresentation::NUM_MANTISSA_BITS;
+  static bool const has_infinity = false;
+
+  /// Least positive value
+  CUTLASS_HOST_DEVICE
+  static type min() { return type::bitcast(0x01); }
+
+  /// Maximum finite value
+  CUTLASS_HOST_DEVICE
+  static type max() { return type::bitcast(type::Base::BitRepresentation::MAX_VALUE); }
+
+  /// Returns maximum rounding error
+  CUTLASS_HOST_DEVICE
+  static type round_error() { return type(0.5f); }
+
+  /// Returns positive infinity value
+  CUTLASS_HOST_DEVICE
+  static type infinity() { return type::bitcast(type::Base::BitRepresentation::INF_MASK); }
+
+  /// Returns quiet NaN value
+  CUTLASS_HOST_DEVICE
+  static type quiet_NaN() { return type::bitcast(type::Base::BitRepresentation::INF_MASK); }
+
+  /// Returns signaling NaN value
+  CUTLASS_HOST_DEVICE
+  static type signaling_NaN() { return type::bitcast(type::Base::BitRepresentation::INF_MASK); }
+
+  /// Returns smallest positive subnormal value
+  CUTLASS_HOST_DEVICE
+  static type denorm_min() { return type::bitcast(0x01); }
+};
+
+/// Numeric limits for float_ue8m0_t
+template <>
+struct numeric_limits<cutlass::float_ue8m0_t> :
+    public float8_exmy_numeric_limits<cutlass::float_ue8m0_t> {
+  static bool const has_infinity = false;
+  static bool const is_signed = false;
+
+  /// Minimum finite value
+  static cutlass::float_ue8m0_t lowest() { return cutlass::float_ue8m0_t::bitcast(0xfe); }
+
+  /// Machine epsilon, that is, the difference between 1.0 and the next representable value (2^0)
+  static cutlass::float_ue8m0_t epsilon() { return cutlass::float_ue8m0_t::bitcast(0x7f); }
+};
+
+
+}  // namespace platform
+
+}  // namespace cutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// User-defined literals
+//
+
+CUTLASS_HOST_DEVICE
+cutlass::float_e4m3_t operator "" _fe4m3(long double x) {
+  return cutlass::float_e4m3_t(float(x));
+}
+
+CUTLASS_HOST_DEVICE
+cutlass::float_e4m3_t operator "" _fe4m3(unsigned long long int x) {
+  return cutlass::float_e4m3_t(int(x));
+}
+
+
+CUTLASS_HOST_DEVICE
+cutlass::float_ue4m3_t operator "" _fue4m3(long double x) {
+  return cutlass::float_ue4m3_t(float(x));
+}
+
+CUTLASS_HOST_DEVICE
+cutlass::float_ue4m3_t operator "" _fue4m3(unsigned long long int x) {
+  return cutlass::float_ue4m3_t(int(x));
+}
+
+
+CUTLASS_HOST_DEVICE
+cutlass::float_e5m2_t operator "" _fe5m2(long double x) {
+  return cutlass::float_e5m2_t(float(x));
+}
+
+CUTLASS_HOST_DEVICE
+cutlass::float_e5m2_t operator "" _fe5m2(unsigned long long int x) {
+  return cutlass::float_e5m2_t(int(x));
+}
+
+
+CUTLASS_HOST_DEVICE
+cutlass::float_ue8m0_t operator "" _fue8m0(long double x)
+{
+  return cutlass::float_ue8m0_t(float(x));
+}
+
+CUTLASS_HOST_DEVICE
+cutlass::float_ue8m0_t operator "" _fue8m0(unsigned long long int x)
+{
+  return cutlass::float_ue8m0_t(int(x));
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/float_subbyte.h

@@ -0,0 +1,797 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+
+/*!
+  \file
+  \brief Defines classes for FP4/FP6 datatypes
+*/
+#pragma once
+
+#include "cutlass/arch/config.h"
+#include "cutlass/float8.h"
+
+// FP4 types are available starting CUDA 12+
+#if (__CUDACC_VER_MAJOR__ >= 12)
+#define CUDA_FP4_ENABLED 1
+#endif
+
+#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED) ||\
+     defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM110A_ENABLED) ||\
+     defined(CUTLASS_ARCH_MMA_SM120A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121A_ENABLED))
+#  define CUDA_PTX_FP4FP6_CVT_ENABLED 1
+#endif
+
+#if (defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101F_ENABLED) ||\
+     defined(CUTLASS_ARCH_MMA_SM103F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM110F_ENABLED) ||\
+     defined(CUTLASS_ARCH_MMA_SM120F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121F_ENABLED))
+#  define CUDA_PTX_FP4FP6_CVT_ENABLED 1
+#endif
+
+#include "cutlass/cutlass.h"
+#include "cutlass/exmy_base.h"
+
+#include "cute/util/type_traits.hpp"
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+
+// FP4 and FP6 types
+struct float_e2m1_t;
+struct float_e3m2_t;
+// E2M1:
+//   2 Exponent bits with 1 Mantissa bit
+//   Range: +-[0,0.5,1,1.5,2,3,4,5,6]
+//   has_Inf: false
+//   has_NaN: false
+//   has_denorm: true
+//   Exponent bias (exp_bias): 1
+
+struct float_e2m1_t : public float_exmy_base<cutlass::detail::FpEncoding::E2M1, float_e2m1_t> {
+  
+  using Base = float_exmy_base<cutlass::detail::FpEncoding::E2M1, float_e2m1_t>;
+
+  float_e2m1_t() = default;
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m1_t(double x) : Base(float(x)) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m1_t(float x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m1_t(int x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  float_e2m1_t(Base x) : Base(x) {
+  }
+};
+
+namespace detail {
+
+// This new type is used to select correct MMA type and TMA type.
+struct float_e2m1_unpacksmem_t : public float_exmy_base<cutlass::detail::FpEncoding::E2M1, float_e2m1_t> {
+
+  using Base = float_exmy_base<cutlass::detail::FpEncoding::E2M1, float_e2m1_t>;
+
+  float_e2m1_unpacksmem_t() = default;
+
+  CUTLASS_HOST_DEVICE
+  float_e2m1_unpacksmem_t(float_e2m1_unpacksmem_t const& x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m1_unpacksmem_t(double x) : Base(float(x)) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m1_unpacksmem_t(float x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m1_unpacksmem_t(int x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  float_e2m1_unpacksmem_t(Base x) : Base(x) {
+  }
+};
+
+} // namespace detail
+
+/// Defines the size of an element in bits - specialized for float_e2m1_t
+template <>
+struct sizeof_bits<float_e2m1_t> {
+  static constexpr int value = 4;
+};
+
+template <>
+struct sizeof_bits<detail::float_e2m1_unpacksmem_t> {
+  static constexpr int value = 4;
+};
+
+CUTLASS_HOST_DEVICE
+float_e2m1_t abs(float_e2m1_t const& val) {
+  using BaseType = typename float_e2m1_t::Base;
+  return float_e2m1_t(abs(BaseType{val.raw()}));
+}
+
+
+// E2M3:
+//   2 Exponent bits with 3 Mantissa bit
+//   Range: [-7.5,+7.5]
+//   has_Inf: false
+//   has_NaN: false
+//   has_denorm: true
+//   Exponent bias (exp_bias): 1
+
+struct float_e2m3_t : public float_exmy_base<cutlass::detail::FpEncoding::E2M3, float_e2m3_t> {
+
+  using Base = float_exmy_base<cutlass::detail::FpEncoding::E2M3, float_e2m3_t>;
+
+  float_e2m3_t() = default;
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m3_t(double x) : Base(float(x)) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m3_t(float x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m3_t(int x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  float_e2m3_t(Base x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m3_t(float_e3m2_t x);
+};
+
+namespace detail {
+
+struct float_e2m3_unpack8bits_t: public float_exmy_base<cutlass::detail::FpEncoding::E2M3, float_e2m3_unpack8bits_t> {
+  // Used in register.
+  using Base = float_exmy_base<cutlass::detail::FpEncoding::E2M3, float_e2m3_unpack8bits_t>;
+
+  float_e2m3_unpack8bits_t() = default;
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m3_unpack8bits_t(double x) : Base(float(x)) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m3_unpack8bits_t(float x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m3_unpack8bits_t(int x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  float_e2m3_unpack8bits_t(Base x) : Base(x) {
+  }
+};
+
+// This new type is used to select correct MMA type and TMA type.
+struct float_e2m3_unpacksmem_t : public float_exmy_base<cutlass::detail::FpEncoding::E2M3, float_e2m3_t> {
+
+  using Base = float_exmy_base<cutlass::detail::FpEncoding::E2M3, float_e2m3_t>;
+
+  float_e2m3_unpacksmem_t() = default;
+
+  CUTLASS_HOST_DEVICE
+  float_e2m3_unpacksmem_t(float_e2m3_unpacksmem_t const& x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m3_unpacksmem_t(double x) : Base(float(x)) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m3_unpacksmem_t(float x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e2m3_unpacksmem_t(int x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  float_e2m3_unpacksmem_t(Base x) : Base(x) {
+  }
+};
+
+} // namespace detail
+
+/// Defines the size of an element in bits - specialized for float_e2m3_t
+template <>
+struct sizeof_bits<float_e2m3_t> {
+  static constexpr int value = 6;
+};
+
+/// Defines the size of an element in bits - specialized for float_e2m3_unpacksmem_t
+template <>
+struct sizeof_bits<detail::float_e2m3_unpacksmem_t> {
+  static constexpr int value = 6;
+};
+
+CUTLASS_HOST_DEVICE
+float_e2m3_t abs(float_e2m3_t const& val) {
+  using BaseType = typename float_e2m3_t::Base;
+  return float_e2m3_t(abs(BaseType{val.raw()}));
+}
+
+// E3M2:
+//   3 Exponent bits, 2 Mantissa bits
+//   Range: [-28:+28]
+//   has_inf: false
+//   has_NaN: false
+//   has_denorm: true
+//   Exponent bias (exp_bias): 3
+
+struct float_e3m2_t : public float_exmy_base<cutlass::detail::FpEncoding::E3M2, float_e3m2_t> {
+
+  using Base = float_exmy_base<cutlass::detail::FpEncoding::E3M2, float_e3m2_t>;
+
+  float_e3m2_t() = default;
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e3m2_t(double x) : Base(float(x)) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e3m2_t(float x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e3m2_t(int x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  float_e3m2_t(Base x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e3m2_t(float_e2m3_t x);
+};
+
+namespace detail {
+
+struct float_e3m2_unpack8bits_t : public float_exmy_base<cutlass::detail::FpEncoding::E3M2, float_e3m2_unpack8bits_t> {
+
+  using Base = float_exmy_base<cutlass::detail::FpEncoding::E3M2, float_e3m2_unpack8bits_t>;
+
+  float_e3m2_unpack8bits_t() = default;
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e3m2_unpack8bits_t(double x) : Base(float(x)) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e3m2_unpack8bits_t(float x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e3m2_unpack8bits_t(int x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  float_e3m2_unpack8bits_t(Base x) : Base(x) {
+  }
+};
+
+// This new type is used to select correct MMA type and TMA type.
+struct float_e3m2_unpacksmem_t : public float_exmy_base<cutlass::detail::FpEncoding::E3M2, float_e3m2_t> {
+
+  using Base = float_exmy_base<cutlass::detail::FpEncoding::E3M2, float_e3m2_t>;
+
+  float_e3m2_unpacksmem_t() = default;
+
+  CUTLASS_HOST_DEVICE
+  float_e3m2_unpacksmem_t(float_e3m2_unpacksmem_t const& x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e3m2_unpacksmem_t(double x) : Base(float(x)) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e3m2_unpacksmem_t(float x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit float_e3m2_unpacksmem_t(int x) : Base(x) {
+  }
+
+  CUTLASS_HOST_DEVICE
+  float_e3m2_unpacksmem_t(Base x) : Base(x) {
+  }
+};
+
+} // namespace detail
+
+/// Defines the size of an element in bits - specialized for float_e3m2_t
+template <>
+struct sizeof_bits<float_e3m2_t> {
+  static constexpr int value = 6;
+};
+
+/// Defines the size of an element in bits - specialized for float_e3m2_unpacksmem_t
+template <>
+struct sizeof_bits<detail::float_e3m2_unpacksmem_t> {
+  static constexpr int value = 6;
+};
+
+CUTLASS_HOST_DEVICE
+float_e3m2_t abs(float_e3m2_t const& val) {
+  using BaseType = typename float_e3m2_t::Base;
+  return float_e3m2_t(abs(BaseType{val.raw()}));
+}
+
+/// Defines the size of an element in bits - specialized for float_e3m2_unpack8bits_t
+template <>
+struct sizeof_bits<detail::float_e3m2_unpack8bits_t> {
+  static constexpr int value = 8;
+};
+
+/// Defines the size of an element in bits - specialized for float_e2m3_unpack8bits_t
+template <>
+struct sizeof_bits<detail::float_e2m3_unpack8bits_t> {
+  static constexpr int value = 8;
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Get the register type used in kernel
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template<typename T>
+struct get_unpacked_element_type;
+
+template <>
+struct get_unpacked_element_type<float_e2m3_t> {
+  using type = detail::float_e2m3_unpack8bits_t;
+};
+
+template <>
+struct get_unpacked_element_type<float_e3m2_t> {
+  using type = detail::float_e3m2_unpack8bits_t;
+};
+} // namespace detail
+// ///////////////////////////////////////////////////////////////////////////////////////////////////
+// //
+// // float_e2m3_t <=> float_e3m2_t conversions
+// //
+// ///////////////////////////////////////////////////////////////////////////////////////////////////
+
+CUTLASS_HOST_DEVICE
+float_e2m3_t::float_e2m3_t(float_e3m2_t x)
+{
+  storage = convert_from_float(float(x)).storage;
+}
+
+CUTLASS_HOST_DEVICE
+float_e3m2_t::float_e3m2_t(float_e2m3_t x)
+{
+  storage = convert_from_float(float(x)).storage;
+}
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+///////////////////////////////////////////////////////////////
+///
+/// Umbrella floating-point 6-bit data type : type_erased_dynamic_float6_t
+/// This umbrella datatype can be enabled when a user provides a specific
+/// datatype in runtime argument list.
+/// 
+/// Currently supported runtime datatypes compatible with type_erased_dynamic_float6_t:
+///   MXF8F6F4Format::E2M3
+///   MXF8F6F4Format::E3M2
+///
+///////////////////////////////////////////////////////////////
+
+union type_erased_dynamic_float6_t {
+  cutlass::float_e2m3_t e2m3;
+  cutlass::float_e3m2_t e3m2;
+
+  CUTLASS_HOST_DEVICE
+  explicit operator cutlass::float_e2m3_t() const { 
+    return e2m3;
+  }
+
+  CUTLASS_HOST_DEVICE
+  explicit operator cutlass::float_e3m2_t() const { 
+    return e3m2;
+  }
+};
+
+template <>
+struct sizeof_bits<type_erased_dynamic_float6_t> {
+  static constexpr int value = 6;
+};
+
+///////////////////////////////////////////////////////////////
+///
+/// Umbrella floating-point 4-bit data type : type_erased_dynamic_float4_t
+/// This umbrella datatype can be enabled when a user provides a specific
+/// datatype in runtime argument list.
+/// 
+/// Currently supported runtime datatypes compatible with type_erased_dynamic_float4_t:
+///   MXF8F6F4Format::E2M1
+///
+///////////////////////////////////////////////////////////////
+
+union type_erased_dynamic_float4_t {
+  cutlass::float_e2m1_t e2m1;
+  CUTLASS_HOST_DEVICE
+  explicit operator cutlass::float_e2m1_t() const { 
+    return e2m1;
+  }
+};
+
+template <>
+struct sizeof_bits<type_erased_dynamic_float4_t> {
+  static constexpr int value = 4;
+};
+
+
+///////////////////////////////////////////////////////////////
+/// MX/NV types for float6 and float4
+/// Intended to be used in builders
+///////////////////////////////////////////////////////////////
+
+template <class F6Type>
+struct mx_float6_t
+{
+  static_assert(cute::is_same_v<F6Type,cutlass::float_e2m3_t>
+                || cute::is_same_v<F6Type,cutlass::float_e3m2_t>
+                || cute::is_same_v<F6Type,type_erased_dynamic_float6_t>
+                , "Only float_e2m3_t, float_e3m2_t can have scale factors for MXFP6");
+  using ScaleFactorType = cutlass::float_ue8m0_t;
+  using DataType = F6Type;
+};
+
+using type_erased_dynamic_mx_float6_t = mx_float6_t<type_erased_dynamic_float6_t>;
+
+template <class F4Type>
+struct mx_float4_t
+{
+  static_assert(cute::is_same_v<F4Type,cutlass::float_e2m1_t>
+                || cute::is_same_v<F4Type,type_erased_dynamic_float4_t>
+                , "Only float_e2m1_t type_erased_dynamic_float4_t can have scale factors for MXFP4");
+  using ScaleFactorType = cutlass::float_ue8m0_t;
+  using DataType = F4Type;
+};
+
+using type_erased_dynamic_mx_float4_t = mx_float4_t<type_erased_dynamic_float4_t>;
+
+template <class F4Type>
+struct nv_float4_t
+{
+  static_assert(cute::is_same_v<F4Type,cutlass::float_e2m1_t>
+                || cute::is_same_v<F4Type,type_erased_dynamic_float4_t>
+                , "Only float_e2m1_t type_erased_dynamic_float4_t can have scale factors for NVFP4");
+  using ScaleFactorType = cutlass::float_ue4m3_t;
+  using DataType = F4Type;
+};
+
+using type_erased_dynamic_nv_float4_t = nv_float4_t<type_erased_dynamic_float4_t>;
+
+
+namespace detail {
+
+union type_erased_dynamic_float6_unpacksmem_t {
+  cutlass::detail::float_e2m3_unpacksmem_t e2m3_unpacksmem;
+  cutlass::detail::float_e3m2_unpacksmem_t e3m2_unpacksmem;
+
+  CUTLASS_HOST_DEVICE
+  explicit operator cutlass::detail::float_e2m3_unpacksmem_t() const { 
+    return e2m3_unpacksmem;
+  }
+  
+  CUTLASS_HOST_DEVICE
+  explicit operator cutlass::detail::float_e3m2_unpacksmem_t() const { 
+    return e3m2_unpacksmem;
+  }
+};
+
+union type_erased_dynamic_float4_unpacksmem_t {
+  cutlass::detail::float_e2m1_unpacksmem_t e2m1_unpacksmem;
+
+  CUTLASS_HOST_DEVICE
+  explicit operator cutlass::detail::float_e2m1_unpacksmem_t() const { 
+    return e2m1_unpacksmem;
+  }
+};
+
+};
+
+template <>
+struct sizeof_bits<detail::type_erased_dynamic_float6_unpacksmem_t> {
+  static constexpr int value = 6;
+};
+
+
+template <>
+struct sizeof_bits<detail::type_erased_dynamic_float4_unpacksmem_t> {
+  static constexpr int value = 4;
+};
+
+} // namespace cutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Standard Library operations and definitions
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+#if !defined(__CUDACC_RTC__)
+namespace std {
+/// Numeric limits common to all float4 types
+template <typename T>
+struct float_subbyte_base_numeric_limits
+{
+private:
+  using type = T;
+
+public:
+  static bool const is_specialized = true;
+  static bool const is_signed = true;
+  static bool const is_integer = false;
+  static bool const is_exact = false;
+  static bool const has_quiet_NaN = false;
+  static bool const has_signaling_NaN = false;
+  static bool const has_denorm_loss = true;
+  static cutlass::platform::float_denorm_style const has_denorm = cutlass::platform::denorm_present;
+  static cutlass::platform::float_round_style const round_style = cutlass::platform::round_to_nearest;
+  static bool const is_iec559 = false;
+  static bool const is_bounded = true;
+  static bool const is_modulo = false;
+  static int const digits = type::Base::BitRepresentation::NUM_MANTISSA_BITS;
+  static bool const has_infinity = false;
+
+  /// Least positive value
+  static type min() { return type::bitcast(0x01); }
+
+  /// Maximum finite value
+  static type max() { return type::bitcast(type::Base::BitRepresentation::MAX_VALUE); }
+
+  /// Returns maximum rounding error
+  static type round_error() { return type(0.5f); }
+
+  /// Returns positive infinity value
+  static type infinity() { return type::bitcast(type::Base::BitRepresentation::INF_MASK); }
+
+  /// Returns quiet NaN value
+  static type quiet_NaN() { return type::bitcast(type::Base::BitRepresentation::INF_MASK); }
+
+  /// Returns signaling NaN value
+  static type signaling_NaN() { return type::bitcast(type::Base::BitRepresentation::INF_MASK); }
+
+  /// Returns smallest positive subnormal value
+  static type denorm_min() { return type::bitcast(0x01); }
+};
+/// Numeric limits for float_e2m1_t
+template <>
+struct numeric_limits<cutlass::float_e2m1_t> : public float_subbyte_base_numeric_limits<cutlass::float_e2m1_t>
+{
+  /// Minimum finite value
+  static cutlass::float_e2m1_t lowest() { return cutlass::float_e2m1_t::bitcast(0xf); }
+
+  /// Returns machine epsilon, that is, the difference between 1.0 and the next value representable by the floating-point
+  static cutlass::float_e2m1_t epsilon() { return cutlass::float_e2m1_t::bitcast(0x1); }
+};
+
+/// Numeric limits for float_e2m3_t
+template <>
+struct numeric_limits<cutlass::float_e2m3_t> : public float_subbyte_base_numeric_limits<cutlass::float_e2m3_t>
+{
+  /// Minimum finite value
+  static cutlass::float_e2m3_t lowest() { return cutlass::float_e2m3_t::bitcast(0x2f); }
+
+  /// Returns machine epsilon, that is, the difference between 1.0 and the next value representable by the floating-point
+  static cutlass::float_e2m3_t epsilon() { return cutlass::float_e2m3_t::bitcast(0x1); }   
+};
+
+/// Numeric limits for float_e3m2_t
+
+template <>
+struct numeric_limits<cutlass::float_e3m2_t> : public float_subbyte_base_numeric_limits<cutlass::float_e3m2_t>
+{
+  /// Minimum finite value
+  static cutlass::float_e3m2_t lowest() { return cutlass::float_e3m2_t::bitcast(0x2f); }
+
+  /// Returns machine epsilon, that is, the difference between 1.0 and the next value representable by the floating-point
+  static cutlass::float_e3m2_t epsilon() { return cutlass::float_e3m2_t::bitcast(0x4); }
+};
+} // namespace std
+#endif
+
+namespace cutlass {
+namespace platform {
+
+/// Numeric limits common to all float4 types
+template <typename T>
+struct float_subbyte_base_numeric_limits
+{
+private:
+  using type = T;
+
+public:
+  static bool const is_specialized = true;
+  static bool const is_signed = true;
+  static bool const is_integer = false;
+  static bool const is_exact = false;
+  static bool const has_quiet_NaN = false;
+  static bool const has_signaling_NaN = false;
+  static bool const has_denorm_loss = true;
+  static cutlass::platform::float_denorm_style const has_denorm = cutlass::platform::denorm_present;
+  static cutlass::platform::float_round_style const round_style = cutlass::platform::round_to_nearest;
+  static bool const is_iec559 = false;
+  static bool const is_bounded = true;
+  static bool const is_modulo = false;
+  static int const digits = type::Base::BitRepresentation::NUM_MANTISSA_BITS;
+  static bool const has_infinity = false;
+
+  /// Least positive value
+  static type min() { return type::bitcast(0x01); }
+
+  /// Maximum finite value
+  CUTLASS_HOST_DEVICE static type max() { return type::bitcast(type::Base::BitRepresentation::MAX_VALUE); }
+
+  /// Returns maximum rounding error
+  static type round_error() { return type(0.5f); }
+
+  /// Returns positive infinity value
+  static type infinity() { return type::bitcast(type::Base::BitRepresentation::INF_MASK); }
+
+  /// Returns quiet NaN value
+  static type quiet_NaN() { return type::bitcast(type::Base::BitRepresentation::INF_MASK); }
+
+  /// Returns signaling NaN value
+  static type signaling_NaN() { return type::bitcast(type::Base::BitRepresentation::INF_MASK); }
+
+  /// Returns smallest positive subnormal value
+  static type denorm_min() { return type::bitcast(0x01); }
+};
+
+/// Forward Declaration
+template <class T>
+struct numeric_limits;
+/// Numeric limits for float_e2m1_t
+template <>
+struct numeric_limits<cutlass::float_e2m1_t> : public float_subbyte_base_numeric_limits<cutlass::float_e2m1_t>
+{
+  /// Minimum finite value
+  static cutlass::float_e2m1_t lowest() { return cutlass::float_e2m1_t::bitcast(0xf); }
+
+  /// Returns machine epsilon, that is, the difference between 1.0 and the next value representable by the floating-point
+  static cutlass::float_e2m1_t epsilon() { return cutlass::float_e2m1_t::bitcast(0x1); }
+};
+
+/// Numeric limits for float_e2m3_t
+template <>
+struct numeric_limits<cutlass::float_e2m3_t> : public float_subbyte_base_numeric_limits<cutlass::float_e2m3_t>
+{
+  /// Minimum finite value
+  static cutlass::float_e2m3_t lowest() { return cutlass::float_e2m3_t::bitcast(0x2f); }
+
+  /// Returns machine epsilon, that is, the difference between 1.0 and the next value representable by the floating-point
+  static cutlass::float_e2m3_t epsilon() { return cutlass::float_e2m3_t::bitcast(0x1); }   
+};
+
+/// Numeric limits for float_e3m2_t
+
+template <>
+struct numeric_limits<cutlass::float_e3m2_t> : public float_subbyte_base_numeric_limits<cutlass::float_e3m2_t>
+{
+  /// Minimum finite value
+  static cutlass::float_e3m2_t lowest() { return cutlass::float_e3m2_t::bitcast(0x2f); }
+
+  /// Returns machine epsilon, that is, the difference between 1.0 and the next value representable by the floating-point
+  static cutlass::float_e3m2_t epsilon() { return cutlass::float_e3m2_t::bitcast(0x4); }
+};
+
+/// Numeric limits for float_e2m3_unpack8bits_t
+template <>
+struct numeric_limits<cutlass::detail::float_e2m3_unpack8bits_t> : public float_subbyte_base_numeric_limits<cutlass::detail::float_e2m3_unpack8bits_t>
+{
+  /// Minimum finite value
+  static cutlass::detail::float_e2m3_unpack8bits_t lowest() { return cutlass::detail::float_e2m3_unpack8bits_t::bitcast(0x2f); }
+
+  /// Returns machine epsilon, that is, the difference between 1.0 and the next value representable by the floating-point
+  static cutlass::detail::float_e2m3_unpack8bits_t epsilon() { return cutlass::detail::float_e2m3_unpack8bits_t::bitcast(0x1); }   
+};
+
+/// Numeric limits for float_e3m2_unpack8bits_t
+
+template <>
+struct numeric_limits<cutlass::detail::float_e3m2_unpack8bits_t> : public float_subbyte_base_numeric_limits<cutlass::detail::float_e3m2_unpack8bits_t>
+{
+  /// Minimum finite value
+  static cutlass::detail::float_e3m2_unpack8bits_t lowest() { return cutlass::detail::float_e3m2_unpack8bits_t::bitcast(0x2f); }
+
+  /// Returns machine epsilon, that is, the difference between 1.0 and the next value representable by the floating-point
+  static cutlass::detail::float_e3m2_unpack8bits_t epsilon() { return cutlass::detail::float_e3m2_unpack8bits_t::bitcast(0x4); }
+};
+} // namespace platform
+
+} // namespace cutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// User-defined literals
+//
+CUTLASS_HOST_DEVICE
+cutlass::float_e2m1_t operator"" _fe2m1(long double x)
+{
+  return cutlass::float_e2m1_t(float(x));
+}
+
+CUTLASS_HOST_DEVICE
+cutlass::float_e2m1_t operator"" _fe2m1(unsigned long long int x)
+{
+  return cutlass::float_e2m1_t(int(x));
+}
+CUTLASS_HOST_DEVICE
+cutlass::float_e2m3_t operator"" _fe2m3(long double x)
+{
+  return cutlass::float_e2m3_t(float(x));
+}
+
+CUTLASS_HOST_DEVICE
+cutlass::float_e2m3_t operator"" _fe2m3(unsigned long long int x)
+{
+  return cutlass::float_e2m3_t(int(x));
+}
+
+CUTLASS_HOST_DEVICE
+cutlass::float_e3m2_t operator"" _fe3m2(long double x)
+{
+  return cutlass::float_e3m2_t(float(x));
+}
+
+CUTLASS_HOST_DEVICE
+cutlass::float_e3m2_t operator"" _fe3m2(unsigned long long int x)
+{
+  return cutlass::float_e3m2_t(int(x));
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/floating_point_nvrtc.h

@@ -0,0 +1,104 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*!
+    \file
+    \brief Defines categories for floating point numbers for use in NVRTC-compiled code
+*/
+
+#pragma once
+
+#include <cutlass/detail/helper_macros.hpp> // CUTLASS_HOST_DEVICE
+#include <cutlass/platform/platform.h> // uint32_t
+#if !defined(__CUDACC_RTC__)
+#include <cstring> // std::memcpy
+#endif
+
+namespace cutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+// All floating-point numbers can be put in one of these categories.
+enum  {
+    FP_NAN =
+# define FP_NAN 0
+      FP_NAN,
+    FP_INFINITE =
+# define FP_INFINITE 1
+      FP_INFINITE,
+    FP_ZERO =
+# define FP_ZERO 2
+      FP_ZERO,
+    FP_SUBNORMAL =
+# define FP_SUBNORMAL 3
+      FP_SUBNORMAL,
+    FP_NORMAL =
+# define FP_NORMAL 4
+      FP_NORMAL
+};
+
+CUTLASS_HOST_DEVICE
+int fpclassify(float const& f) {
+
+  uint32_t s;
+
+  #if defined(__CUDA_ARCH__)
+  s = reinterpret_cast<uint32_t const &>(f);
+  #else
+  std::memcpy(&s, &f, sizeof(s));
+  #endif
+
+  uint32_t exp      = s & 0x7f800000;
+  uint32_t mantissa = s & 0x007fffff;
+
+  if (exp == 0x7f800000) {
+    if (mantissa) {
+      return FP_NAN;
+    }
+    else {
+      return FP_INFINITE;
+    }
+  }
+  else if (!exp) {
+    if (mantissa) {
+      return FP_SUBNORMAL;
+    }
+    else {
+      return FP_ZERO;
+    }
+  }
+  return FP_NORMAL;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace cutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/functional.h

@@ -0,0 +1,1106 @@
+  /***************************************************************************************************
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Define basic numeric operators
+
+    This is inspired by the Standard Library's <functional> header.
+*/
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/platform/platform.h"
+#if defined(__CUDACC_RTC__)
+#include "cutlass/floating_point_nvrtc.h"
+#endif
+
+#include <cuda_runtime.h>
+
+#if defined(CUTLASS_ARCH_WMMA_ENABLED)
+#include <mma.h>
+#endif // defined(CUTLASS_ARCH_WMMA_ENABLED)
+
+#ifdef _MSC_VER
+// Provides support for alternate operators such as 'and', 'or', ...
+#include <ciso646>
+#include <intrin.h>
+#endif // _MSC_VER
+
+#if defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) ||\
+    defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM103F_ENABLED)
+#  define CUTLASS_ARCH_CREDUX_ENABLED
+#endif
+
+namespace cutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+  CUTLASS_HOST_DEVICE int32_t popcount(int32_t x) {
+    #if defined(__CUDA_ARCH__)
+    return __popc(x);
+    #elif defined(__GNUC__) || defined(__clang__)
+    return __builtin_popcount(x);
+    #elif (defined(_MSC_VER) && !defined(_M_ARM64))
+    return __popcnt(x);
+    #else
+    int32_t count = 0;
+    while (x) {
+      count += x & 1;
+      x >>= 1;
+    }
+    return count;
+    #endif
+  }
+
+  CUTLASS_HOST_DEVICE int64_t popcount(int64_t x) {
+    #if defined(__CUDA_ARCH__)
+    return __popcll(x);
+    #elif defined(__GNUC__) || defined(__clang__)
+    return __builtin_popcountll(x);
+    #elif (defined(_MSC_VER) && !defined(_M_ARM64))
+    return __popcnt64(x);
+    #else
+    int64_t count = 0;
+    while (x) {
+      count += x & 1;
+      x >>= 1;
+    }
+    return count;
+    #endif
+  }
+
+} // namespace detail
+  
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+struct absolute_value_op {
+  CUTLASS_HOST_DEVICE
+  T operator()(T lhs) const {
+    return abs(lhs);
+  }
+};
+
+template <>
+struct absolute_value_op<float> {
+  CUTLASS_HOST_DEVICE
+  float operator()(float lhs) const { return fabs(lhs); }
+};
+
+template <typename T>
+struct plus {
+  CUTLASS_HOST_DEVICE
+  T operator()(T lhs, T const &rhs) const {
+    lhs += rhs;
+    return lhs;
+  }
+};
+
+template <typename T>
+struct minus {
+  CUTLASS_HOST_DEVICE
+  T operator()(T lhs, T const &rhs) const {
+    lhs -= rhs;
+    return lhs;
+  }
+};
+
+template <typename T>
+struct multiplies {
+  CUTLASS_HOST_DEVICE
+  T operator()(T lhs, T const &rhs) const {
+    lhs *= rhs;
+    return lhs;
+  }
+};
+
+template <typename T>
+struct scale {
+  T const scaling_factor_;
+
+  CUTLASS_HOST_DEVICE
+  scale(float scaling_factor) : scaling_factor_(scaling_factor) {
+  }
+
+  T operator()(T const &rhs) const {
+    T result = rhs * scaling_factor_;
+    return result;
+  }
+};
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 530
+/// Partial specializations needed when __CUDA_NO_HALF2_OPERATORS__ is set
+template<>
+struct plus<__half2> {
+  CUTLASS_HOST_DEVICE
+  __half2 operator()(__half2 lhs, __half2 const &rhs) const {
+    return __hadd2(lhs, rhs);
+  }
+};
+
+template<>
+struct minus<__half2> {
+  CUTLASS_HOST_DEVICE
+  __half2 operator()(__half2 lhs, __half2 const &rhs) const {
+    return __hsub2(lhs, rhs);
+  }
+};
+
+template<>
+struct multiplies<__half2> {
+  CUTLASS_HOST_DEVICE
+  __half2 operator()(__half2 lhs, __half2 const &rhs) const {
+    return __hmul2(lhs, rhs);
+  }
+};
+
+/// Partial specializations needed when __CUDA_NO_HALF_OPERATORS__ is set
+template<>
+struct plus<__half> {
+  CUTLASS_HOST_DEVICE
+  __half operator()(__half lhs, __half const &rhs) const {
+    return __hadd(lhs, rhs);
+  }
+};
+
+template<>
+struct minus<__half> {
+  CUTLASS_HOST_DEVICE
+  __half operator()(__half lhs, __half const &rhs) const {
+    return __hsub(lhs, rhs);
+  }
+};
+
+template<>
+struct multiplies<__half> {
+  CUTLASS_HOST_DEVICE
+  __half operator()(__half lhs, __half const &rhs) const {
+    return __hmul(lhs, rhs);
+  }
+};
+#endif // defined(__CUDA_ARCH__)
+
+
+/// Squares with optional conversion
+template <typename T, typename Output = T>
+struct square {
+  CUTLASS_HOST_DEVICE
+  Output operator()(T lhs) const {
+    multiplies<Output> mul_op;
+
+    Output y = Output(lhs);
+    return mul_op(y, y);
+  }
+};
+
+/// Returns the magnitude squared of an element.
+template <typename T, typename Output = T>
+struct magnitude_squared {
+  CUTLASS_HOST_DEVICE
+  Output operator()(T lhs) const {
+    multiplies<Output> mul_op;
+
+    Output y = Output(lhs);
+    return mul_op(y, y);
+  }
+};
+
+/// Computes the square of a difference with optional conversion
+template <typename T, typename Output = T>
+struct square_difference {
+  CUTLASS_HOST_DEVICE
+  Output operator()(T lhs, T rhs) const {
+    multiplies<Output> mul_op;
+
+    Output y = Output(lhs) - Output(rhs);
+    return mul_op(y, y);
+  }
+};
+
+/// Computes the square of a difference with optional conversion
+template <typename T, typename Output = T>
+struct magnitude_squared_difference {
+  CUTLASS_HOST_DEVICE
+  Output operator()(T lhs, T rhs) const {
+    multiplies<Output> mul_op;
+
+    Output y = Output(lhs) - Output(rhs);
+    return mul_op(y, y);
+  }
+};
+
+// Computes the reciprocal square root
+template <typename T>
+struct inverse_square_root;
+
+template <>
+struct inverse_square_root<float> {
+  CUTLASS_HOST_DEVICE
+  float operator()(float const &lhs) const {
+#if defined(__CUDA_ARCH__)
+    return rsqrtf(lhs);
+#else
+    return 1.f / std::sqrt(lhs);
+#endif
+  }
+};
+
+template <>
+struct inverse_square_root<half_t> {
+  CUTLASS_HOST_DEVICE
+  half_t operator()(half_t const &lhs) const {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ > 520)
+    auto result = hrsqrt(reinterpret_cast<__half const &>(lhs));
+    return reinterpret_cast<half_t const &>(result);
+#else
+    return half_t(1.f / std::sqrt(half_t::convert(lhs)));
+#endif
+  }
+};
+
+/// Divides
+template <typename T>
+struct divides {
+  CUTLASS_HOST_DEVICE
+  T operator()(T lhs, T const &rhs) const {
+    lhs /= rhs;
+    return lhs;
+  }
+};
+
+/// reciprocal_approximate
+template <typename T>
+struct reciprocal_approximate {
+  CUTLASS_HOST_DEVICE
+  T operator()(T lhs) const {
+    return divides<T>{}(T(1), lhs);
+  }
+};
+
+template <>
+struct reciprocal_approximate <float> {
+  CUTLASS_HOST_DEVICE
+  float operator()(float lhs) const {
+    float ret;
+    #if defined(__CUDA_ARCH__)
+      asm volatile ("rcp.approx.f32 %0, %1;\n" : "=f"(ret) : "f"(lhs));
+    #else
+      ret = 1.0f / lhs;
+    #endif
+    return ret;
+  }
+};
+
+
+template <>
+struct reciprocal_approximate<cutlass::float_ue8m0_t> {
+  CUTLASS_HOST_DEVICE
+  cutlass::float_ue8m0_t operator()(cutlass::float_ue8m0_t lhs) const {
+    return cutlass::float_ue8m0_t::bitcast(static_cast<uint8_t>(static_cast<uint8_t>(254u) - lhs.storage));
+  }
+};
+
+
+/// reciprocal_approximate with ftz
+template<typename T>
+struct reciprocal_approximate_ftz :  reciprocal_approximate<T>
+{};
+
+template <>
+struct reciprocal_approximate_ftz <float> {
+  CUTLASS_HOST_DEVICE
+  float operator()(float lhs) const {
+    float ret;
+    #if defined(__CUDA_ARCH__)
+      asm volatile ("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(ret) : "f"(lhs));
+    #else
+      if (std::fpclassify(lhs) == FP_SUBNORMAL) {
+        lhs = 0.0f;
+      }
+      ret = 1.0f / lhs;
+      if (std::fpclassify(ret) == FP_SUBNORMAL) {
+        ret = 0.0f;
+      }
+    #endif
+    return ret;
+  }
+};
+
+/// Negate
+template <typename T>
+struct negate {
+  CUTLASS_HOST_DEVICE
+  T operator()(T lhs) const {
+    return -lhs;
+  }
+};
+
+/// Greater equal
+template <typename T>
+struct greater_equal {
+  CUTLASS_HOST_DEVICE
+  bool operator()(T const &lhs, T const &rhs) const {
+    return (lhs >= rhs);
+  }
+};
+
+/// Greater
+template <typename T>
+struct greater {
+  CUTLASS_HOST_DEVICE
+  bool operator()(T const &lhs, T const &rhs) const {
+    return (lhs > rhs);
+  }
+};
+
+/// Less equal
+template <typename T>
+struct less_equal {
+  CUTLASS_HOST_DEVICE
+  bool operator()(T const &lhs, T const &rhs) const {
+    return (lhs <= rhs);
+  }
+};
+
+/// Less
+template <typename T>
+struct less {
+  CUTLASS_HOST_DEVICE
+  bool operator()(T const &lhs, T const &rhs) const {
+    return (lhs < rhs);
+  }
+};
+
+template <typename T, bool PropagateNaN = false>
+struct maximum {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &lhs, T const &rhs) const {
+    if constexpr (PropagateNaN && cutlass::platform::is_floating_point<T>::value) {
+      using CUTLASS_CMATH_NAMESPACE :: isnan;
+
+      // Call isnan unqualified, so argument-dependent lookup (ADL)
+      // will find overloads such as cutlass::isnan(half_t).
+      // Calling ::isnan or std::isnan directly would force
+      // implicit conversions to float of custom number types
+      // in the cutlass namespace (e.g., cutlass::half_t).
+      return lhs > rhs || isnan(lhs) ? lhs : rhs;
+    }
+    else {
+      return (lhs < rhs ? rhs : lhs);
+    }
+
+    CUTE_GCC_UNREACHABLE;
+  }
+};
+
+// This is a subclass and not an alias
+// in order to work around a known Clang issue,
+// where a template template parameter with one template parameter
+// does not match classes that take multiple template parameters
+// but have defaults for all but the first.
+template<typename T>
+struct maximum_with_default_nan_propagation : public maximum<T>
+{};
+
+template <>
+struct maximum<float, false> {
+  CUTLASS_HOST_DEVICE
+  float operator()(float const &lhs, float const &rhs) const {
+    return fmaxf(lhs, rhs);
+  }
+};
+
+template <>
+struct maximum<float, true> {
+  CUTLASS_HOST_DEVICE
+  float operator()(float lhs, float rhs) const {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+    float res;
+    asm volatile("max.NaN.f32 %0, %1, %2;\n" : "=f"(res) : "f"(lhs), "f"(rhs));
+    return res;
+#else
+    using CUTLASS_CMATH_NAMESPACE :: isnan;
+
+    return lhs > rhs || isnan(lhs) ? lhs : rhs;
+#endif
+  }
+};
+
+// This is a subclass and not an alias
+// in order to work around a known Clang issue,
+// where a template template parameter with one template parameter
+// does not match classes that take multiple template parameters
+// but have defaults for all but the first.
+template <typename T>
+struct maximum_with_nan_propagation : maximum<T, true>
+{};
+
+// This alias exists for backwards compatibility only.
+// Please use the correctly spelled class template above.
+template <typename T>
+using maximum_with_nan_propogation = maximum_with_nan_propagation<T>;
+
+template <typename T, bool PropagateNaN = false>
+struct minimum {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &lhs, T const &rhs) const {
+    if constexpr (PropagateNaN && cutlass::platform::is_floating_point<T>::value) {
+      using CUTLASS_CMATH_NAMESPACE :: isnan;
+
+      return lhs < rhs || isnan(lhs) ? lhs : rhs;
+    }
+    else {
+      return (rhs < lhs ? rhs : lhs);
+    }
+  }
+};
+
+template <>
+struct minimum<float, false> {
+  CUTLASS_HOST_DEVICE
+  float operator()(float const &lhs, float const &rhs) const {
+    return fminf(lhs, rhs);
+  }
+};
+
+template <>
+struct minimum<float, true> {
+  CUTLASS_HOST_DEVICE
+  float operator()(float lhs, float rhs) const {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+    float res;
+    asm volatile("min.NaN.f32 %0, %1, %2;\n" : "=f"(res) : "f"(lhs), "f"(rhs));
+    return res;
+#else
+    // No need for ADL; call std::isnan(float) on host and ::isnan(float) on device.
+    return lhs < rhs || (CUTLASS_CMATH_NAMESPACE :: isnan(lhs)) ? lhs : rhs;
+#endif
+  }
+};
+
+template <typename T>
+struct minimum_with_nan_propagation : minimum<T, true> 
+{};
+
+template <typename T, bool PropagateNaN = false>
+struct maximum_absolute_value {
+  CUTLASS_HOST_DEVICE
+  float operator()(T const &lhs, T const &rhs) const {
+    absolute_value_op<T> abs_op;
+    maximum<T, PropagateNaN> max_op;
+
+    return max_op(abs_op(lhs), abs_op(rhs));
+  }
+};
+
+// assumes the left operand is already an absolute value
+template <typename T, bool PropagateNaN = false>
+struct maximum_absolute_value_reduction {
+  CUTLASS_HOST_DEVICE
+  float operator()(T const &lhs, T const &rhs) const {
+    absolute_value_op<T> abs_op;
+    maximum<T, PropagateNaN> max_op;
+
+    return max_op(lhs, abs_op(rhs));
+  }
+};
+
+/// Fused multiply-add
+template <typename A, typename B = A, typename C = A>
+struct multiply_add {
+  CUTLASS_HOST_DEVICE
+  C operator()(A const &a, B const &b, C const &c) const {
+    return C(a) * C(b) + c;
+  }
+};
+
+template <typename T>
+struct square_and_plus {
+  CUTLASS_HOST_DEVICE
+  T operator()(T lhs, T const &rhs) const {
+    multiply_add<T> multiply_add_op;
+    return multiply_add_op(rhs, rhs, lhs);
+  }
+};
+
+// Fused multiply-add that takes exactly one template parameter.
+// This is useful for working around a known Clang issue,
+// where a template template parameter with one template parameter
+// does not match classes that take multiple template parameters
+// but have defaults for all but the first.
+template <typename A>
+struct homogeneous_multiply_add : public multiply_add<A, A, A>
+{};
+
+/// Fused multiply-add
+template <typename A, typename B = A, typename C = A>
+struct multiply_add_relu0 {
+  CUTLASS_HOST_DEVICE
+  C operator()(A const &a, B const &b, C const &c) const {
+    maximum<C> mx;
+    return mx(C(a) * C(b) + c, C(0));
+  }
+};
+
+/// Guarded-multiply-add
+template <typename A, typename B = A, typename C = A>
+struct guarded_multiply_add {
+  CUTLASS_HOST_DEVICE
+  C operator()(A const &a, B const &b, C const &c) const {
+    using CUTLASS_CMATH_NAMESPACE :: isnan;
+
+    if (isnan(a) || isnan(b)) {
+      return C(0);
+    }
+    return C(a) * C(b) + c;
+  }
+};
+
+/// Guarded-multiply-add
+template <>
+struct guarded_multiply_add<half_t, half_t, half_t> {
+  CUTLASS_HOST_DEVICE
+  half_t operator()(half_t const &a, half_t const &b, half_t const &c) const {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+    half_t result;
+    asm ("fma.rn.oob.f16 %0, %1, %2, %3;\n"
+      : "=h"(*reinterpret_cast<uint16_t*>(&result))
+      : "h"(*reinterpret_cast<uint16_t const*>(&a)), "h"(*reinterpret_cast<uint16_t const*>(&b)), "h"(*reinterpret_cast<uint16_t const*>(&c)));
+    return result;
+#else
+    // Namespace-qualifying isnan as cutlass::isnan saves the compiler
+    // the trouble of argument-dependent lookup.  Calling std::isnan or
+    // ::isnan here would result in unwanted implicit conversion to float.
+    if (cutlass::isnan(a) || cutlass::isnan(b)) {
+      return half_t(0);
+    }
+    return a * b + c;
+#endif
+  }
+};
+
+/// Guarded-multiply-add-relu0
+template <typename A, typename B = A, typename C = A>
+struct guarded_multiply_add_relu0 {
+  CUTLASS_HOST_DEVICE
+  C operator()(A const &a, B const &b, C const &c) const {
+    using CUTLASS_CMATH_NAMESPACE :: isnan;
+
+    if (isnan(a) || isnan(b)) {
+      return C(0);
+    }
+    maximum<C> mx;
+    return mx(C(a) * C(b) + c, C(0));
+  }
+};
+
+template <>
+struct guarded_multiply_add_relu0<half_t, half_t, half_t> {
+  CUTLASS_HOST_DEVICE
+  half_t operator()(half_t const &a, half_t const &b, half_t const &c) const {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+    half_t result;
+    asm ("fma.rn.oob.relu.f16 %0, %1, %2, %3;\n"
+      : "=h"(*reinterpret_cast<uint16_t*>(&result))
+      : "h"(*reinterpret_cast<uint16_t const*>(&a)), "h"(*reinterpret_cast<uint16_t const*>(&b)), "h"(*reinterpret_cast<uint16_t const*>(&c)));
+    return result;
+#else
+    if (cutlass::isnan(a) || cutlass::isnan(b)) {
+      return half_t(0);
+    }
+    maximum<half_t> mx;
+    return mx(a * b + c, half_t(0));
+#endif
+  }
+};
+
+
+/// Fused and-popc-add
+template <typename A, typename B = A, typename C = A>
+struct and_popc_add {
+  CUTLASS_HOST_DEVICE
+  C operator()(A const &a, B const &b, C const &c) const {
+    A and_result = a & b;
+    int32_t popc_result = detail::popcount(and_result);
+    return C(popc_result) + c;
+  }
+};
+
+/// Fused and-add
+template <typename T>
+struct and_add {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b, T const &c) const {
+    return ((a & b) + c);
+  }
+};
+
+
+
+/// Fused xor-popc-add
+template <typename A, typename B = A, typename C = A>
+struct xor_popc_add {
+  CUTLASS_HOST_DEVICE
+  C operator()(A const &a, B const &b, C const &c) const {
+    A xor_result = a ^ b;
+    int32_t popc_result = detail::popcount(xor_result);
+    return C(popc_result) + c;
+  }
+};
+
+/// Fused xor-add
+template <typename T>
+struct xor_add {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b, T const &c) const {
+    return ((a ^ b) + c);
+  }
+};
+
+
+/// Fused or-popc-add
+template <typename A, typename B = A, typename C = A>
+struct or_popc_add {
+  CUTLASS_HOST_DEVICE
+  C operator()(A const &a, B const &b, C const &c) const {
+    A or_result = a | b;
+    int32_t popc_result = detail::popcount(or_result);
+    return C(popc_result) + c;
+  }
+};
+
+
+/// Fused or-add
+template <typename T>
+struct or_add {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b, T const &c) const {
+    return ((a | b) + c);
+  }
+};
+
+namespace detail {
+
+// Whether namespace-unqualified conj(t) for t of type T is
+// well-formed.  This says whether the compiler can find
+// namespace-unqualified conj(T) via argument-dependent lookup.
+// If so, then CUTLASS assumes that conj(t) returns
+// the complex conjugate of t.
+template <typename T, typename Enable = void>
+struct has_unqualified_conj : cutlass::platform::false_type
+{};
+
+template<typename T>
+struct has_unqualified_conj<
+    T,
+    decltype(static_cast<void>(conj(cutlass::platform::declval<T>())), void())
+  > : cutlass::platform::true_type
+{};
+
+template <typename T>
+constexpr bool has_unqualified_conj_v = has_unqualified_conj<T>::value;
+  
+} // namespace detail
+
+// forward declaration (needed for conjugate below)
+template<class T>
+CUTLASS_HOST_DEVICE T conj(T const& z);
+
+namespace detail {
+
+// Whether cutlass::conj(t) for t of type T is well-formed.
+// If so, then CUTLASS assumes that cutlass::conj(t)
+// returns the complex conjugate of t.
+template <typename T, typename Enable = void>
+struct has_cutlass_conj : cutlass::platform::false_type
+{};
+
+template<typename T>
+struct has_cutlass_conj<
+    T,
+    decltype(cutlass::conj(cutlass::platform::declval<T>()), void())
+  > : cutlass::platform::true_type
+{};
+
+template <typename T>
+constexpr bool has_cutlass_conj_v = has_cutlass_conj<T>::value;
+
+} // namespace detail
+  
+// Return the complex conjugate of the input.
+//
+// If the struct hasn't already been specialized for type T, then
+//
+// 1. for arithmetic types, return z;
+//
+// 2. for types where either (namespace-unqualified) conj(z) or
+//    cutlass::conj(z) is well formed, declare "using cutlass::conj;"
+//    and return conj(z); and
+//
+// 3. for everything else, return z.
+//
+// Regarding (1), the C++ Standard Library makes std::conj always
+// return std::complex, even for (noncomplex) arithmetic types.
+// cutlass::conj(T t) needs to return type T.  This follows the
+// convention of linear algebra software like the BLAS, where
+// "conjugate transpose" means the same thing as "transpose" for a
+// matrix of noncomplex numbers.
+//
+// Case (2) covers std::complex, cuda::std::complex, and non-Standard
+// (including user-defined) complex number types (for which "conj(z)"
+// is findable via argument-dependent lookup).  cutlass::conj has a
+// totally generic overload, but a more type-specific overload in any
+// namespace will take precedence.
+//
+// Case (3) covers non-Standard non-complex number types.
+//
+// Users should not generally need to specialize this struct for their
+// own custom complex or noncomplex types.  The idiomatic way to
+// identify a type T as "complex" is to make namespace-unqualified
+// calls to conj(T) findable via argument-dependent lookup.
+template <typename T>
+struct conjugate {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const& z) const {
+    if constexpr (cutlass::platform::is_arithmetic_v<T>) {
+      return z;
+    }
+    else if constexpr (detail::has_unqualified_conj_v<T> || detail::has_cutlass_conj_v<T>) {
+      using cutlass::conj;
+      return conj(z);
+    }
+    else {
+      return z;
+    }
+  }
+};
+
+template <typename T>
+struct first {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const & first, T const &...) const {
+    return first;
+  }
+  CUTLASS_HOST_DEVICE
+  T operator()(T const & first) const {
+    return first;
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+struct logical_and {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b) const {
+    return ((static_cast<bool>(a) && static_cast<bool>(b)) ? T(1) : T());
+  }
+};
+
+template <typename T>
+struct logical_or {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b) const {
+    return ((static_cast<bool>(a) || static_cast<bool>(b)) ? T(1) : T());
+  }
+};
+
+template <typename T>
+struct logical_not {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &a) const {
+    return T(!(a));
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+struct bit_and {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b) const {
+    return a & b;
+  }
+};
+
+template <typename T>
+struct bit_or {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b) const {
+    return a | b;
+  }
+};
+
+template <typename T>
+struct bit_not {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &a) const {
+    return ~a;
+  }
+};
+
+template <typename T>
+struct bit_xor {
+  CUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b) const {
+    return a ^ b;
+  }
+};
+
+//////////////////////////////////////////////////////////////////////////////////////////////////
+/// Atomic reductions
+
+template <typename T>
+struct atomic_add
+{
+  CUTLASS_DEVICE
+  void operator()(T *ptr, const T &data)
+  {
+#if defined(__CUDA_ARCH__)
+    atomicAdd(ptr, data);
+#else
+    CUTLASS_UNUSED(ptr);
+    CUTLASS_UNUSED(data);
+    CUTLASS_NOT_IMPLEMENTED();
+#endif
+  }
+};
+
+template<>
+struct atomic_add<double>
+{
+  CUTLASS_DEVICE
+  void operator()(double *ptr, const double &data)
+  {
+#if !defined(__CUDA_ARCH__)
+    CUTLASS_UNUSED(ptr);
+    CUTLASS_UNUSED(data);
+    CUTLASS_NOT_IMPLEMENTED();
+#elif (__CUDA_ARCH__ >= 600)
+    atomicAdd(ptr, data);
+#else
+    // Use CAS loop
+    unsigned long long int* ptr_int = reinterpret_cast<unsigned long long int*>(ptr);
+    unsigned long long int old_int = *ptr_int;
+    unsigned long long int assumed_int;
+
+    do {
+      double update = data + __longlong_as_double(old_int);
+      assumed_int = old_int;
+      old_int = atomicCAS(ptr_int, assumed_int, __double_as_longlong(update));
+    } while (assumed_int != old_int);
+#endif // (__CUDA_ARCH__ >= 600)
+  }
+};
+
+template<>
+struct atomic_add<half2>
+{
+  CUTLASS_DEVICE
+  void operator()(half2 *ptr, const half2 &data)
+  {
+#if !defined(__CUDA_ARCH__) || (defined(__CUDA_ARCH__)  && (__CUDA_ARCH__ < 600))
+      CUTLASS_UNUSED(ptr);
+      CUTLASS_UNUSED(data);
+      CUTLASS_NOT_IMPLEMENTED();
+#else
+    // Vector-2 atomic reduction requires .target sm_60 or higher
+    uint32_t word = reinterpret_cast<const uint32_t&>(data);
+    asm volatile ("red.gpu.global.add.noftz.f16x2 [%0], %1;\n" : : "l"(ptr), "r"(word));
+#endif // (__CUDA_ARCH__ >= 600)
+  }
+};
+
+template <typename T>
+using red [[deprecated("use atomic_add instead")]] = atomic_add<T>;
+
+template <typename T>
+struct atomic_maximum {
+  CUTLASS_DEVICE
+  T operator()(T *ptr, T value) const {
+#if defined(__CUDA_ARCH__)
+    return atomicMax(ptr, value);
+#else
+    CUTLASS_UNUSED(ptr);
+    CUTLASS_UNUSED(value);
+    CUTLASS_NOT_IMPLEMENTED();
+    return 0;
+#endif
+  }
+};
+
+template <>
+struct atomic_maximum<float> {
+  CUTLASS_DEVICE
+  float operator()(float *ptr, float value) const {
+#if defined(__CUDA_ARCH__)
+    // In device code, make sure that we do NOT try to use
+    // std::signbit, as that won't work if building with NVRTC.
+    // Instead, prefix "::" to call signbit from the global namespace,
+    // which CUDA guarantees to work in device code without including
+    // any headers.
+    //
+    return ! ::signbit(value) ?
+      __int_as_float(atomicMax((int*)ptr, __float_as_int(value))) :
+      __uint_as_float(atomicMin((unsigned int*)ptr, __float_as_uint(value)));
+#else
+    CUTLASS_UNUSED(ptr);
+    CUTLASS_UNUSED(value);
+    CUTLASS_NOT_IMPLEMENTED();
+    return 0;
+#endif
+  }
+};
+
+// is_atomic
+template <class Fn>
+struct is_atomic : platform::false_type {};
+template <class T>
+struct is_atomic<atomic_add<T>> : platform::true_type {};
+template <class T>
+struct is_atomic<atomic_maximum<T>> : platform::true_type {};
+
+
+//////////////////////////////////////////////////////////////////////////////////////////////////
+/// Parallel Synchronization and Communication Instructions
+template <typename T>
+struct redux_abs_max_nan_propagation_sync_warp;
+
+template <>
+struct redux_abs_max_nan_propagation_sync_warp <float>{
+  CUTLASS_DEVICE
+  float operator()(float const &lhs) const {
+#if defined(CUTLASS_ARCH_CREDUX_ENABLED)
+    float result;
+    asm volatile("redux.sync.max.abs.NaN.f32 %0, %1, 0xffffffff;\n" : "=f"(result) : "f"(lhs));
+    return result;
+#elif defined(__CUDA_ARCH__)
+    cutlass::maximum<float, /*PropagateNaN*/true> max_op;
+    int shuffle_width = 32;
+    float abs_max = cutlass::absolute_value_op<float>{}(lhs);
+    CUTLASS_PRAGMA_UNROLL
+    for(int offset = shuffle_width / 2; offset > 0; offset /= 2) {
+      float value = __shfl_down_sync(0xffffffff, abs_max, offset, shuffle_width);
+      abs_max = max_op(abs_max,value);
+    }
+    // Broadcast the maximum to all threads participating in the reduction.
+    abs_max = __shfl_sync(0xffffffff, abs_max, 0, shuffle_width);
+    return abs_max;
+#else
+    CUTLASS_UNUSED(lhs);
+    CUTLASS_NOT_IMPLEMENTED();
+    return 0;
+#endif
+  }
+};
+
+template <typename T>
+struct redux_abs_max_nan_propagation_sync_warp_t0t15_t16t31;
+
+template <>
+struct redux_abs_max_nan_propagation_sync_warp_t0t15_t16t31<float>{
+  CUTLASS_DEVICE
+  float operator()(float const &max) const {
+#if defined(CUTLASS_ARCH_CREDUX_ENABLED)
+    int half_warp_idx = threadIdx.x / (NumThreadsPerWarp / 2);
+    bool first_half_threads = (half_warp_idx % 2) == 0;
+    float value0 =  first_half_threads ? max : 0;
+    float v0 = cutlass::redux_abs_max_nan_propagation_sync_warp<float>{}(value0);
+
+    float value1 = !first_half_threads ? max : 0;
+    float v1 = cutlass::redux_abs_max_nan_propagation_sync_warp<float>{}(value1);
+    return first_half_threads ? v0: v1;
+    
+#elif defined(__CUDA_ARCH__)
+    float abs_max = cutlass::absolute_value_op<float>{}(max);
+    cutlass::maximum<float, /*PropagateNaN*/true> max_op;
+    constexpr int shuffle_width = 16;
+    CUTLASS_PRAGMA_UNROLL
+    for(int offset = shuffle_width/2; offset > 0; offset /= 2) {
+      float value = __shfl_down_sync(0xffffffff, abs_max, offset, shuffle_width);
+        abs_max  = max_op(abs_max,value);
+    }
+    // Broadcast the maximum to all threads participating in the reduction.
+    abs_max = __shfl_sync(0xffffffff, abs_max, 0, shuffle_width);
+    return abs_max;
+#else 
+    CUTLASS_UNUSED(max);
+    CUTLASS_NOT_IMPLEMENTED();
+    return 0;
+#endif
+  }
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Partial specializations for nvcuda::wmma::fragment<Use, m, n, k, T, Layout>
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+#if defined(CUTLASS_ARCH_WMMA_ENABLED)
+
+template<typename Use, int m, int n, int k, typename T, typename Layout>
+struct plus<nvcuda::wmma::fragment<Use, m, n, k, T, Layout>>
+{
+  using Fragment = nvcuda::wmma::fragment<Use, m, n, k, T, Layout>;
+  using ElementType = typename Fragment::element_type;
+
+  CUTLASS_HOST_DEVICE
+  Fragment operator()(Fragment const &lhs, Fragment const &rhs) const
+  {
+    Fragment result;
+    plus<ElementType> scalar_op;
+
+    ElementType *result_elts = reinterpret_cast<ElementType*>(&result);
+    const ElementType *lhs_elts = reinterpret_cast<const ElementType*>(&lhs);
+    const ElementType *rhs_elts = reinterpret_cast<const ElementType*>(&rhs);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Fragment::num_elements; i++) {
+      result_elts[i] = scalar_op(lhs_elts[i], rhs_elts[i]);
+    }
+
+    return result;
+  }
+};
+
+#endif // defined(CUTLASS_ARCH_WMMA_ENABLED)
+
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/gemm/collective/collective_builder.hpp

@@ -0,0 +1,62 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+#include "cutlass/gemm/collective/collective_mma_decl.hpp"
+#include "cutlass/gemm/collective/collective_mma.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include "cutlass/gemm/collective/collective_builder_decl.hpp"
+#include "cutlass/gemm/collective/builders/sm90_gmma_builder.inl"
+#include "cutlass/gemm/collective/builders/sm90_sparse_gmma_builder.inl"
+#if !defined(__CUDACC_RTC__) 
+#include "cutlass/gemm/collective/builders/sm100_umma_builder.inl"              
+#include "cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl"       
+#include "cutlass/gemm/collective/builders/sm100_sparse_umma_builder.inl"
+#include "cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl"  
+#include "cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl"
+#include "cutlass/gemm/collective/builders/sm100_blockscaled_sparse_umma_builder.inl"
+#include "cutlass/gemm/collective/builders/sm100_simt_builder.inl"
+#include "cutlass/gemm/collective/builders/sm100_mixed_input_umma_builder.inl"       
+#include "cutlass/gemm/collective/builders/sm100_cpasync_umma_builder.inl"
+#include "cutlass/gemm/collective/builders/sm100_mixed_tma_cpasync_umma_builder.inl"
+#include "cutlass/gemm/collective/builders/sm100_blockscaled_mixed_tma_cpasync_umma_builder.inl"
+#include "cutlass/gemm/collective/builders/sm103_blockscaled_umma_builder.inl"
+#include "cutlass/gemm/collective/builders/sm120_mma_builder.inl"
+#include "cutlass/gemm/collective/builders/sm120_blockscaled_mma_builder.inl"
+#include "cutlass/gemm/collective/builders/sm120_sparse_mma_builder.inl"
+#include "cutlass/gemm/collective/builders/sm120_blockscaled_sparse_mma_builder.inl"
+#include "cutlass/gemm/collective/builders/sm120_blockwise_mma_builder.inl"
+#endif
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/gemm/collective/collective_builder_decl.hpp

@@ -0,0 +1,100 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <cute/numeric/integral_constant.hpp>
+#include <cutlass/detail/dependent_false.hpp>
+
+namespace cutlass::gemm::collective {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Used to specify stage counts or dispatch to automatic computation of stage count
+template<int num_stages>
+struct StageCount {
+  static constexpr int value = num_stages;
+
+  StageCount() = default;
+  explicit StageCount(cute::Int<num_stages>) {}
+};
+
+template<int carveout_bytes>
+struct StageCountAutoCarveout {
+  static constexpr int bytes = carveout_bytes;
+
+  StageCountAutoCarveout() = default;
+  explicit StageCountAutoCarveout(cute::Int<carveout_bytes>) {}
+};
+
+namespace detail {
+
+// Forward Declaration
+template<class CollectiveEpilogue>
+constexpr int
+compute_carveout_from_epi();
+
+} // namespace detail
+
+template<class CollectiveEpilogue>
+struct StageCountAutoCarveoutEpi : StageCountAutoCarveout<detail::compute_carveout_from_epi<CollectiveEpilogue>()> {};
+
+using StageCountAuto = StageCountAutoCarveout<0>;
+
+// Used to automatically let the builder pick the kernel schedule.
+// Can be overridden with kernel schedule tags in cutlass/gemm/dispatch_policy.hpp
+struct KernelScheduleAuto final {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  class ArchTag,
+  class OpClass,
+  class ElementA,
+  class GmemLayoutA,
+  int AlignmentA,
+  class ElementB,
+  class GmemLayoutB,
+  int AlignmentB,
+  class ElementAccumulator,
+  class TileShape_MNK,
+  class ClusterShape_MNK,
+  class StageCountType,
+  class KernelScheduleType,
+  class Enable = void
+>
+struct CollectiveBuilder {
+  static_assert(sizeof(ElementA) == 0, "Could not build a collective for given parameters.");
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace cutlass::gemm::collective
+

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/gemm/collective/collective_mma.hpp

@@ -0,0 +1,84 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "cutlass/gemm/collective/collective_mma_decl.hpp"
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include "cutlass/gemm/collective/sm70_mma_twostage.hpp"
+#include "cutlass/gemm/collective/sm80_mma_multistage.hpp"
+#include "cutlass/gemm/collective/sm80_mma_array_multistage.hpp"
+#include "cutlass/gemm/collective/sm90_mma_multistage_gmma_ss_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm90_mma_multistage_gmma_rs_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp"
+#include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp" 
+#include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized_fp8.hpp"
+#include "cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input.hpp"
+#include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp"
+#include "cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8.hpp"
+#include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp"
+#include "cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp"
+#if !defined(__CUDACC_RTC__)
+#include "cutlass/gemm/collective/sm100_mma_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm100_mma_warpspecialized_emulated.hpp"
+#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp"
+#include "cutlass/gemm/collective/sm100_sparse_mma_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm100_blockscaled_sparse_mma_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp" 
+#include "cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp" 
+#include "cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp"
+#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_blockwise_scaling.hpp"
+#include "cutlass/gemm/collective/sm100_mma_warpspecialized_mixed_input.hpp"
+#include "cutlass/gemm/collective/sm100_mma_cpasync_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm100_mma_mixed_tma_cpasync_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm100_blockscaled_mma_mixed_tma_cpasync_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm103_blockscaled_mma_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm103_blockscaled_mma_array_warpspecialized.hpp"
+#include "cutlass/gemm/collective/sm120_mma_tma.hpp"
+#include "cutlass/gemm/collective/sm120_blockscaled_mma_tma.hpp"
+#include "cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp"
+#include "cutlass/gemm/collective/sm120_sparse_mma_tma.hpp"
+#include "cutlass/gemm/collective/sm120_blockscaled_sparse_mma_tma.hpp"
+#include "cutlass/gemm/collective/sm120_mma_tma_blockwise_scaling.hpp"
+#include "cutlass/gemm/collective/sm120_mma_array_tma_blockwise_scaling.hpp"
+#endif // !defined(__CUDACC_RTC__) 
+
+
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/gemm/collective/collective_mma_decl.hpp

@@ -0,0 +1,64 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <cute/numeric/integral_constant.hpp>
+#include <cutlass/detail/dependent_false.hpp>
+
+namespace cutlass::gemm::collective {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  class DispatchPolicy,
+  class TileShape,
+  class ElementA,
+  class StrideA,
+  class ElementB,
+  class StrideB,
+  class TiledMma,
+  class GmemTiledCopyA,
+  class SmemLayoutAtomA,
+  class SmemCopyAtomA,
+  class TransformA,
+  class GmemTiledCopyB,
+  class SmemLayoutAtomB,
+  class SmemCopyAtomB,
+  class TransformB
+>
+struct CollectiveMma {
+  static_assert(cutlass::detail::dependent_false<ElementA>, "Could not find a mainloop specialization.");
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace cutlass::gemm::collective
+

build/torch210-cxx11-cu128-aarch64-linux/include/third-party/cutlass/include/cutlass/gemm/collective/fp8_accumulation.hpp

@@ -0,0 +1,279 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include "cute/algorithm/clear.hpp"
+#include "cute/tensor.hpp"
+
+//////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////FP8 Accumulation///////////////////////////
+//////////////////////////////////////////////////////////////////////////////
+/// This class provides API to promote (add) or scale (multiply_add) the results 
+/// from the tensor core accumulators to the main accumulators when the number 
+/// of MMAs reaches the max number of MMA interval specified by user, after that
+/// the tensor core accumulators are zeroed.
+//////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass::gemm::collective {
+
+template <
+    class EngineAccum,
+    class LayoutAccum>
+struct GmmaFP8Accumulation {  
+  using TensorAccum = cute::Tensor<EngineAccum, LayoutAccum>;
+  using ElementAccumulator = typename EngineAccum::value_type;
+
+  static_assert(is_static<LayoutAccum>::value, "Accumulator Layout should be static");
+  static_assert(is_rmem<TensorAccum>::value , "Accumulator tensor must be rmem resident.");
+
+private:
+  TensorAccum& accum_;
+  TensorAccum accum_temp_;
+
+  uint32_t accum_promotion_interval_;         // defines the max num of executed MMAs after which accum should be promoted.
+  uint32_t mma_count_per_mainloop_iteration_; // num of MMAs per k_tile of mainloop
+  uint32_t mma_count_;                        // current executed MMAs
+  uint32_t reset_accum_flag_;                 // accum needs to be zeroed or not. 
+
+  // promote or `add` the partial accumulators to main accumulator (FADD).
+  CUTLASS_DEVICE
+  void promote_core() {
+    warpgroup_wait<0>();
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < size(accum_); ++i) {
+      accum_(i) += accum_temp_(i);
+    }
+  }
+
+  // `multiply` scale the partial accumulators and `add` to main accumulator (FFMA).
+  CUTLASS_DEVICE
+  void scale_core(ElementAccumulator const &scale) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < size(accum_); ++i) {
+      accum_(i) += accum_temp_(i) * scale;
+    }
+  }
+
+  template <
+    class EngineScale,
+    class LayoutScale>
+  CUTLASS_DEVICE
+  void scale_core(const cute::Tensor<EngineScale, LayoutScale> &scale) {
+    using TensorScale = cute::Tensor<EngineScale, LayoutScale>;
+
+    static_assert(is_static<LayoutScale>::value, "Scale Layout should be static");
+    static_assert(is_rmem<TensorScale>::value , "Scale tensor must be rmem resident.");
+
+    static_assert(LayoutAccum{}.shape() == LayoutScale{}.shape(), "Accumulator and scale must have same shape.");
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < size(accum_); ++i) {
+      accum_(i) += accum_temp_(i) * scale(i);
+    }
+  }
+
+  template <
+    class EngineScaleA,
+    class LayoutScaleA,
+    class EngineScaleB,
+    class LayoutScaleB>
+  CUTLASS_DEVICE
+  void scale_core(const cute::Tensor<EngineScaleA, LayoutScaleA> &scaleA, const cute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
+    using TensorScaleA = cute::Tensor<EngineScaleA, LayoutScaleA>;
+    using TensorScaleB = cute::Tensor<EngineScaleB, LayoutScaleB>;
+
+    static_assert(is_static<LayoutScaleA>::value, "ScaleA Layout should be static");
+    static_assert(is_static<LayoutScaleB>::value, "ScaleB Layout should be static");
+    static_assert(is_rmem<TensorScaleA>::value, "ScaleA tensor must be rmem resident.");
+    static_assert(is_rmem<TensorScaleB>::value, "ScaleB tensor must be rmem resident.");
+
+    static_assert(LayoutAccum{}.shape() == LayoutScaleA{}.shape(), "Accumulator and scaleA must have same shape.");
+    static_assert(LayoutAccum{}.shape() == LayoutScaleB{}.shape(), "Accumulator and scaleB must have same shape.");
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < size(accum_); ++i) {
+      accum_(i) += accum_temp_(i) * scaleA(i) * scaleB(i);
+    }
+  }
+
+public:
+  CUTLASS_DEVICE
+  GmmaFP8Accumulation(
+      TensorAccum &accum,
+      uint32_t accum_promotion_interval,
+      uint32_t mma_count_per_mainloop_iteration)
+      : accum_(accum), 
+        accum_promotion_interval_(accum_promotion_interval),
+        mma_count_per_mainloop_iteration_(mma_count_per_mainloop_iteration),
+        mma_count_(0), 
+        reset_accum_flag_(0) 
+  {
+    accum_temp_ = cute::make_fragment_like(accum);
+  }
+
+  //
+  // Methods (Common)
+  //
+
+  CUTLASS_DEVICE 
+  TensorAccum& operator()() {
+    return accum_temp_;
+  }
+
+  /// prepare the MMA accumulators when initialization or zeroing is required.
+  CUTLASS_DEVICE
+  bool prepare_if_needed() { 
+    return reset_accum_flag_;
+  }
+
+  //
+  // Methods (for FADD version)
+  //
+
+  /// promote (add) the results from the MMA accumulators to main accumulator if needed.
+  CUTLASS_DEVICE
+  void promote_if_needed() {
+    mma_count_ += mma_count_per_mainloop_iteration_;
+    reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
+    if (reset_accum_flag_) {
+      promote_core();
+      mma_count_ = 0;
+    }
+  }
+
+  /// promote (add) the residue results from the MMA accumulators to main accumulator if needed.
+  CUTLASS_DEVICE
+  void promote_residue_if_needed() {
+    if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) {
+      promote_core();
+    }
+  }
+
+  //
+  // Methods (for FFMA version)
+  //
+
+  /// scale (multiply_add) the results from the MMA accumulators to main accumulator if needed.
+  CUTLASS_DEVICE
+  void scale_if_needed(ElementAccumulator const &scale) {
+    mma_count_ += mma_count_per_mainloop_iteration_;
+    reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
+    if (reset_accum_flag_) {
+      scale_core(scale);
+      mma_count_ = 0;
+    }
+  }
+
+  template <
+    class EngineScale,
+    class LayoutScale>
+  CUTLASS_DEVICE
+  void scale_if_needed(const cute::Tensor<EngineScale, LayoutScale> &scale) {
+    mma_count_ += mma_count_per_mainloop_iteration_;
+    reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
+    if (reset_accum_flag_) {
+      scale_core(scale);
+      mma_count_ = 0;
+    }
+  }
+
+  template <
+    class EngineScaleA,
+    class LayoutScaleA,
+    class EngineScaleB,
+    class LayoutScaleB>
+  CUTLASS_DEVICE
+  void scale_if_needed(const cute::Tensor<EngineScaleA, LayoutScaleA> &scaleA, const cute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
+    mma_count_ += mma_count_per_mainloop_iteration_;
+    reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
+    if (reset_accum_flag_) {
+      scale_core(scaleA, scaleB);
+      mma_count_ = 0;
+    }
+  }
+  
+  /// scale (multiply_add) the results from the MMA accumulators to main accumulator without checking the counter.
+  CUTLASS_DEVICE
+  void scale(ElementAccumulator const &scale) {
+    scale_core(scale);
+  }
+
+  template <
+    class EngineScale,
+    class LayoutScale>
+  CUTLASS_DEVICE
+  void scale(const cute::Tensor<EngineScale, LayoutScale> &scale) {
+    scale_core(scale);
+  }
+
+  template <
+    class EngineScaleA,
+    class LayoutScaleA,
+    class EngineScaleB,
+    class LayoutScaleB>
+  CUTLASS_DEVICE
+  void scale(const cute::Tensor<EngineScaleA, LayoutScaleA> &scaleA, const cute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
+    scale_core(scaleA, scaleB);
+  }
+
+  /// scale (multiply_add) the residue results from the MMA accumulators to main accumulator if needed.
+  CUTLASS_DEVICE
+  void scale_residue_if_needed(ElementAccumulator const &scale) {
+    if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) {
+      scale_core(scale);
+    }
+  }
+
+  template <
+    class EngineScale,
+    class LayoutScale>
+  CUTLASS_DEVICE
+  void scale_residue_if_needed(const cute::Tensor<EngineScale, LayoutScale> &scale) {
+    if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) {
+      scale_core(scale);
+    }
+  }
+
+  template <
+    class EngineScaleA,
+    class LayoutScaleA,
+    class EngineScaleB,
+    class LayoutScaleB>
+  CUTLASS_DEVICE
+  void scale_residue_if_needed(const cute::Tensor<EngineScaleA, LayoutScaleA> &scaleA, const cute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
+    if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) {
+      scale_core(scaleA, scaleB);
+    }
+  }
+};
+
+} // namespace cutlass::gemm::collective