kInstrK Class — pytorch Architecture

Source Code

aten/src/ATen/native/transformers/cuda/mem_eff_attention/iterators/default_warp_iterator_from_smem.h lines 59–76

template <typename RegularWarpIterator, typename Policy, int kInstrK>
struct DefaultWarpIteratorAFromSharedMemory<
    cutlass::gemm::GemmShape<32, 32, 32>,
    cutlass::gemm::GemmShape<16, 8, kInstrK>,
    RegularWarpIterator,
    Policy,
    typename platform::enable_if<(
        sizeof_bits<typename RegularWarpIterator::Element>::value == 16 &&
        Policy::Operator::Policy::OpDelta::kRow == 1)>::type> {
  using OpDelta = typename Policy::Operator::Policy::OpDelta;
  using WarpShape = cutlass::MatrixShape<32, 32>;
  using InstructionShape = cutlass::gemm::GemmShape<16, 8, kInstrK>;

  using WarpIterator = cutlass::gemm::warp::WarpIteratorFromSmem<
      cutlass::gemm::Operand::kA,
      typename RegularWarpIterator::Element,
      cutlass::MatrixShape<InstructionShape::kM, InstructionShape::kK>>;
};