2025-10-27 nightly release (000f49551bfda0683d395e1fe67f53cfb785868f)

.lintrunner.toml

@@ -833,8 +833,7 @@ exclude_patterns = [
 command = [
     'python3',
     'tools/linter/adapters/grep_linter.py',
-    '--pattern=cudaSetDevice(',
-    '--pattern=cudaGetDevice(',
+    '--pattern=(cudaSetDevice|cudaGetDevice)\\(',
     '--linter-name=RAWCUDADEVICE',
     '--error-name=raw CUDA API usage',
     """--error-description=\

aten/src/ATen/core/jit_type_base.h

@@ -185,11 +185,11 @@ struct TORCH_API Type {
         : repr_(nullptr) {}
 
     /* implicit */ SingletonOrSharedTypePtr(SingletonTypePtr<T> p)
-        : repr_(p) {}
+        : repr_(makeSingletonSharedPtr(p.get())) {}
 
     template <typename U, std::enable_if_t<std::is_convertible_v<U*, T*>, bool> = true>
     /* implicit */ SingletonOrSharedTypePtr(SingletonTypePtr<U> p)
-        : repr_(SingletonTypePtr<T>(p.get())) {}
+        : repr_(makeSingletonSharedPtr(static_cast<T*>(p.get()))) {}
 
 
     // We need to support construction from T* for pybind. The problem
@@ -202,8 +202,8 @@ struct TORCH_API Type {
     // Case 2: if T is exactly Type, we need to do a dynamic_cast to
     // check if it's a SharedType and do the right thing.
     //
-    // Case 3: Otherwise, T is not a SharedType. (debug-check this
-    // assumption!) Use a singleton pointer.
+    // Case 3: Otherwise, T is not a SharedType. Use a singleton
+    // pointer.
 
     template <typename U = T, std::enable_if_t<std::is_base_of_v<SharedType, U>, bool> = true>
     /* implicit */ SingletonOrSharedTypePtr(T* p) : SingletonOrSharedTypePtr(static_cast<typename detail::as_shared_type<U>::type>(p)->shared_from_this()) {}
@@ -211,15 +211,15 @@ struct TORCH_API Type {
     template <typename U = T, std::enable_if_t<std::is_same_v<Type, U>, bool> = true>
     /* implicit */ SingletonOrSharedTypePtr(T* p) {
       if (auto* shared_p = dynamic_cast<typename detail::as_shared_type<U>::type>(p)) {
-        repr_ = Repr(shared_p->shared_from_this());
+        repr_ = shared_p->shared_from_this();
       } else {
-        repr_ = Repr(p);
+        repr_ = makeSingletonSharedPtr(p);
       }
     }
 
     template <typename U = T, std::enable_if_t<!std::is_same_v<Type, U> && !std::is_base_of_v<SharedType, U>, bool> = true>
     /* implicit */ SingletonOrSharedTypePtr(T* p)
-        : repr_(p) {
+        : repr_(makeSingletonSharedPtr(p)) {
       TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dynamic_cast<typename detail::as_shared_type<U>::type>(p) == nullptr);
     }
 
@@ -230,19 +230,19 @@ struct TORCH_API Type {
     ~SingletonOrSharedTypePtr() = default;
 
     T* get() const {
-      return repr_.isSharedAndNonNull() ? repr_.shared_.repr_.get() : static_cast<T*>(repr_.rawRepr().first);
+      return repr_.get();
     }
 
     operator bool() const {
-      return repr_.isNonNull();
+      return repr_ != nullptr;
     }
 
     bool operator==(std::nullptr_t) const {
-      return !repr_.isNonNull();
+      return repr_ == nullptr;
     }
 
     bool operator!=(std::nullptr_t) const {
-      return repr_.isNonNull();
+      return repr_ != nullptr;
     }
 
     template <typename U = T, std::enable_if_t<!std::is_same_v<std::remove_const_t<U>, void>, bool> = true>
@@ -255,138 +255,14 @@ struct TORCH_API Type {
     }
 
   private:
-    // NOTE: SharedPtrWrapper exists to work around a baffling bug in
-    // nvcc; see comment in destroy() below.
-    struct SharedPtrWrapper {
-      SharedPtrWrapper(std::shared_ptr<T> &&x)
-          : repr_(std::move(x)) {}
-      std::shared_ptr<T> repr_;
-    };
-    union Repr {
-      Repr() : Repr(nullptr) {}
+    // Use shared_ptr's aliasing constructor to create a non-owning pointer
+    // to a singleton. The lifetime is tied to the null shared_ptr, so there's
+    // no reference counting overhead for the singleton itself.
+    static std::shared_ptr<T> makeSingletonSharedPtr(T* ptr) {
+      return std::shared_ptr<T>(std::shared_ptr<T>(), ptr);
+    }
 
-      explicit Repr(std::shared_ptr<T> x)
-          : shared_(std::move(x)) {}
-
-      explicit Repr(std::nullptr_t)
-          : singletonRepr_(nullptr) {}
-
-      explicit Repr(SingletonTypePtr<T> p)
-          : singletonRepr_(p.get()) {}
-
-      ~Repr() {
-        destroy();
-      }
-
-      // NOTE: the only non-UB way to access our null state is through
-      // rawRepr(), because our copy operation doesn't preserve which
-      // union member is active for null pointers.
-      Repr(const Repr& rhs) {
-        if (rhs.isSharedAndNonNull()) {
-          new (&shared_) SharedPtrWrapper(rhs.shared_);
-        } else {
-          singletonRepr_.singleton_ = static_cast<T*>(rhs.rawRepr().first);
-          TORCH_INTERNAL_ASSERT_DEBUG_ONLY(rhs.singletonRepr_.unused_ == nullptr);
-          singletonRepr_.unused_ = nullptr;
-        }
-      }
-
-      Repr(Repr&& rhs) noexcept {
-        if (rhs.isSharedAndNonNull()) {
-          new (&shared_) SharedPtrWrapper(std::move(rhs.shared_));
-        } else {
-          singletonRepr_.singleton_ = static_cast<T*>(rhs.rawRepr().first);
-          TORCH_INTERNAL_ASSERT_DEBUG_ONLY(rhs.singletonRepr_.unused_ == nullptr);
-          singletonRepr_.unused_ = nullptr;
-        }
-      }
-
-      Repr& operator=(const Repr& rhs) {
-        if (&rhs == this) {
-          return *this;
-        }
-        if (rhs.isSharedAndNonNull()) {
-          if (isSharedAndNonNull()) {
-            shared_ = rhs.shared_;
-          } else {
-            new (&shared_) SharedPtrWrapper(rhs.shared_);
-          }
-        } else {
-          if (isSharedAndNonNull()) {
-            destroy();
-          }
-          singletonRepr_.singleton_ = static_cast<T*>(rhs.rawRepr().first);
-          TORCH_INTERNAL_ASSERT_DEBUG_ONLY(rhs.rawRepr().nullIfSingleton_ == nullptr);
-          singletonRepr_.unused_ = nullptr;
-        }
-        return *this;
-      }
-
-      Repr& operator=(Repr&& rhs) noexcept {
-        if (&rhs == this) {
-          return *this;
-        }
-        if (rhs.isSharedAndNonNull()) {
-          if (isSharedAndNonNull()) {
-            shared_ = std::move(rhs.shared_);
-          } else {
-            new (&shared_) SharedPtrWrapper(std::move(rhs.shared_));
-          }
-        } else {
-          if (isSharedAndNonNull()) {
-            destroy();
-          }
-          singletonRepr_.singleton_ = static_cast<T*>(rhs.rawRepr().first);
-          TORCH_INTERNAL_ASSERT_DEBUG_ONLY(rhs.rawRepr().nullIfSingleton_ == nullptr);
-          singletonRepr_.unused_ = nullptr;
-        }
-        return *this;
-      }
-
-      SharedPtrWrapper shared_;
-
-      struct SingletonRepr {
-        explicit SingletonRepr(T* s) : singleton_(s) {}
-        T* singleton_;
-        void* unused_ = nullptr;
-      } singletonRepr_;
-      struct RawRepr {
-        void* first;
-        void* nullIfSingleton_;
-      };
-
-      // It is UB to read the singleton part of Repr if it was
-      // constructed as a shared_ptr and vice versa, but memcpying out
-      // the representation is always OK, so here's an accessor to obey
-      // the letter of the law.
-      RawRepr rawRepr() const {
-        RawRepr repr{};
-        memcpy(&repr, reinterpret_cast<const char *>(this), sizeof(RawRepr));
-        return repr;
-      }
-
-      bool isNonNull() const {
-        auto repr = rawRepr();
-        TORCH_INTERNAL_ASSERT_DEBUG_ONLY(repr.nullIfSingleton_ == nullptr || repr.first != nullptr);
-        return repr.first != nullptr;
-      }
-
-      bool isSharedAndNonNull() const {
-        return rawRepr().nullIfSingleton_ != nullptr;
-      }
-
-     private:
-      void destroy() {
-        if (isSharedAndNonNull()) {
-          // Without SharedPtrWrapper, this line would read
-          // `shared_.~shared_ptr()` and nvcc would complain with
-          // "error: expected primary-expression before '>' token"
-          // referring to the "t" in "shared_ptr". SharedPtrWrapper
-          // exists to work around this compiler bug.
-          shared_.~SharedPtrWrapper();
-        }
-      }
-    } repr_;
+    std::shared_ptr<T> repr_;
   };
 
   using TypePtr = SingletonOrSharedTypePtr<Type>;

aten/src/ATen/cpu/vec/sve/vec_float.h

@@ -104,71 +104,6 @@ class Vectorized<float> {
     }
     return b;
   }
-  // Implementation is picked from
-  // https://github.com/ARM-software/ComputeLibrary/blob/v25.01/src/core/NEON/SVEMath.inl#L105
-  inline svfloat32_t svexp_f32_z(svbool_t pg, svfloat32_t x) const {
-    const auto c1 =
-        svreinterpret_f32_u32(svdup_n_u32(0x3f7ffff6)); // x^1: 0x1.ffffecp-1f
-    const auto c2 =
-        svreinterpret_f32_u32(svdup_n_u32(0x3efffedb)); // x^2: 0x1.fffdb6p-2f
-    const auto c3 =
-        svreinterpret_f32_u32(svdup_n_u32(0x3e2aaf33)); // x^3: 0x1.555e66p-3f
-    const auto c4 =
-        svreinterpret_f32_u32(svdup_n_u32(0x3d2b9f17)); // x^4: 0x1.573e2ep-5f
-    const auto c5 =
-        svreinterpret_f32_u32(svdup_n_u32(0x3c072010)); // x^5: 0x1.0e4020p-7f
-    const auto shift = svreinterpret_f32_u32(
-        svdup_n_u32(0x4b00007f)); // 2^23 + 127 = 0x1.0000fep23f
-    const auto inv_ln2 = svreinterpret_f32_u32(
-        svdup_n_u32(0x3fb8aa3b)); // 1 / ln(2) = 0x1.715476p+0f
-    const auto neg_ln2_hi = svreinterpret_f32_u32(svdup_n_u32(
-        0xbf317200)); // -ln(2) from bits  -1 to -19: -0x1.62e400p-1f
-    const auto neg_ln2_lo = svreinterpret_f32_u32(svdup_n_u32(
-        0xb5bfbe8e)); // -ln(2) from bits -20 to -42: -0x1.7f7d1cp-20f
-    const auto inf = svdup_n_f32(std::numeric_limits<float>::infinity());
-    const auto max_input = svdup_n_f32(88.37f); // Approximately ln(2^127.5)
-    const auto zero = svdup_n_f32(0.f);
-    const auto min_input = svdup_n_f32(-86.64f); // Approximately ln(2^-125)
-    // Range reduction:
-    //   e^x = 2^n * e^r
-    // where:
-    //   n = floor(x / ln(2))
-    //   r = x - n * ln(2)
-    //
-    // By adding x / ln(2) with 2^23 + 127 (shift):
-    //   * As FP32 fraction part only has 23-bits, the addition of 2^23 + 127
-    //   forces decimal part
-    //     of x / ln(2) out of the result. The integer part of x / ln(2) (i.e.
-    //     n) + 127 will occupy the whole fraction part of z in FP32 format.
-    //     Subtracting 2^23 + 127 (shift) from z will result in the integer part
-    //     of x / ln(2) (i.e. n) because the decimal part has been pushed out
-    //     and lost.
-    //   * The addition of 127 makes the FP32 fraction part of z ready to be
-    //   used as the exponent
-    //     in FP32 format. Left shifting z by 23 bits will result in 2^n.
-    const auto z = svmla_f32_z(pg, shift, x, inv_ln2);
-    const auto n = svsub_f32_z(pg, z, shift);
-    const auto scale = svreinterpret_f32_u32(
-        svlsl_n_u32_z(pg, svreinterpret_u32_f32(z), 23)); // 2^n
-    // The calculation of n * ln(2) is done using 2 steps to achieve accuracy
-    // beyond FP32. This outperforms longer Taylor series (3-4 tabs) both in
-    // term of accuracy and performance.
-    const auto r_hi = svmla_f32_z(pg, x, n, neg_ln2_hi);
-    const auto r = svmla_f32_z(pg, r_hi, n, neg_ln2_lo);
-    // Compute the truncated Taylor series of e^r.
-    //   poly = scale * (1 + c1 * r + c2 * r^2 + c3 * r^3 + c4 * r^4 + c5 * r^5)
-    const auto r2 = svmul_f32_z(pg, r, r);
-    const auto p1 = svmul_f32_z(pg, c1, r);
-    const auto p23 = svmla_f32_z(pg, c2, c3, r);
-    const auto p45 = svmla_f32_z(pg, c4, c5, r);
-    const auto p2345 = svmla_f32_z(pg, p23, p45, r2);
-    const auto p12345 = svmla_f32_z(pg, p1, p2345, r2);
-    auto poly = svmla_f32_z(pg, scale, p12345, scale);
-    // Handle underflow and overflow.
-    poly = svsel_f32(svcmplt_f32(pg, x, min_input), zero, poly);
-    poly = svsel_f32(svcmpgt_f32(pg, x, max_input), inf, poly);
-    return poly;
-  }
   static Vectorized<float> loadu(const void* ptr, int64_t count = size()) {
     if (count == size())
       return svld1_f32(ptrue, reinterpret_cast<const float*>(ptr));
@@ -313,11 +248,41 @@ class Vectorized<float> {
     return USE_SLEEF(
         Vectorized<float>(Sleef_expm1fx_u10sve(values)), map(std::expm1));
   }
+  // Implementation copied from Arm Optimized Routines:
+  // https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/sve/expf.c
   Vectorized<float> exp_u20() const {
-    return exp();
+    // special case to handle special inputs that are too large or too small
+    // i.e. where there's at least one element x, s.t. |x| >= 87.3...
+    svbool_t is_special_case = svacgt(svptrue_b32(), values, 0x1.5d5e2ap+6f);
+    if (svptest_any(svptrue_b32(), is_special_case)) {
+      return exp();
+    }
+    const svfloat32_t ln2_hi = svdup_n_f32(0x1.62e4p-1f);
+    const svfloat32_t ln2_lo = svdup_n_f32(0x1.7f7d1cp-20f);
+    const svfloat32_t c1 = svdup_n_f32(0.5f);
+    const svfloat32_t inv_ln2 = svdup_n_f32(0x1.715476p+0f);
+
+    const float shift = 0x1.803f8p17f;
+
+    /* n = round(x/(ln2/N)).  */
+    svfloat32_t z = svmad_x(svptrue_b32(), inv_ln2, values, shift);
+    svfloat32_t n = svsub_x(svptrue_b32(), z, shift);
+
+    /* r = x - n*ln2/N.  */
+    svfloat32_t r = values;
+    r = svmls_x(svptrue_b32(), r, n, ln2_hi);
+    r = svmls_x(svptrue_b32(), r, n, ln2_lo);
+
+    /* scale = 2^(n/N).  */
+    svfloat32_t scale = svexpa(svreinterpret_u32(z));
+
+    /* poly(r) = exp(r) - 1 ~= r + 0.5 r^2.  */
+    svfloat32_t r2 = svmul_x(svptrue_b32(), r, r);
+    svfloat32_t poly = svmla_x(svptrue_b32(), r, r2, c1);
+    return svmla_x(svptrue_b32(), scale, scale, poly);
   }
   Vectorized<float> fexp_u20() const {
-    return exp();
+    return exp_u20();
   }
   Vectorized<float> fmod(const Vectorized<float>& q) const {USE_SLEEF(
       { return Vectorized<float>(Sleef_fmodfx_sve(values, q)); },
@@ -453,9 +418,11 @@ class Vectorized<float> {
         ptrue, svmax_f32_z(ptrue, values, CONST_MIN_TANH), CONST_MAX_TANH);
 
     // Step 2: Calculate exp(2 * x), where x is the clamped value.
-    // svmul_f32_z computes 2 * x, and svexp_f32_z computes the exponential of
-    // the result.
-    svfloat32_t exp2x = svexp_f32_z(ptrue, svmul_f32_z(ptrue, CONST_2, x));
+    // svmul_f32_z computes 2 * x, and exp_u20() computes the exponential of
+    // the result (via Vectorized<float>, then auto-converts back to
+    // svfloat32_t).
+    svfloat32_t exp2x =
+        Vectorized<float>(svmul_f32_z(ptrue, CONST_2, x)).exp_u20();
 
     // Step 3: Calculate the numerator of the tanh function, which is exp(2x)
     // - 1.

aten/src/ATen/cpu/vec/vec128/vec128_float_neon.h

@@ -307,11 +307,49 @@ class Vectorized<float> {
   DEFINE_SLEEF_COMPATIBLE_UNARY_ELEMENTWISE_FUNC(exp)
   DEFINE_SLEEF_COMPATIBLE_UNARY_ELEMENTWISE_FUNC(exp2)
   DEFINE_SLEEF_COMPATIBLE_UNARY_ELEMENTWISE_FUNC(expm1)
+  // Implementation copied from Arm Optimized Routine
+  // https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/expf.c
   Vectorized<float> exp_u20() const {
-    return exp();
+    // bail out to sleef if it's a special case:
+    // i.e. there's an input s.t. |input| > 87.3....
+    const float32x4_t special_bound = vdupq_n_f32(0x1.5d5e2ap+6f);
+    uint32x4_t cmp = vcagtq_f32(values, special_bound);
+    if (vpaddd_u64(vreinterpretq_u64_u32(cmp)) != 0) {
+      return exp();
+    }
+
+    const float32x4_t inv_ln2 = vdupq_n_f32(0x1.715476p+0f);
+    const float ln2_hi = 0x1.62e4p-1f;
+    const float ln2_lo = 0x1.7f7d1cp-20f;
+    const float c0 = 0x1.0e4020p-7f;
+    const float c2 = 0x1.555e66p-3f;
+    const float32x4_t ln2_c02 = {ln2_hi, ln2_lo, c0, c2};
+
+    const uint32x4_t exponent_bias = vdupq_n_u32(0x3f800000);
+    const float32x4_t c1 = vdupq_n_f32(0x1.573e2ep-5f);
+    const float32x4_t c3 = vdupq_n_f32(0x1.fffdb6p-2f);
+    const float32x4_t c4 = vdupq_n_f32(0x1.ffffecp-1f);
+
+    /* exp(x) = 2^n (1 + poly(r)), with 1 + poly(r) in [1/sqrt(2),sqrt(2)]
+      x = ln2*n + r, with r in [-ln2/2, ln2/2].  */
+
+    float32x4_t n = vrndaq_f32(vmulq_f32(values, inv_ln2));
+    float32x4_t r = vfmsq_laneq_f32(values, n, ln2_c02, 0);
+    r = vfmsq_laneq_f32(r, n, ln2_c02, 1);
+    uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_s32(vcvtq_s32_f32(n)), 23);
+    float32x4_t scale = vreinterpretq_f32_u32(vaddq_u32(e, exponent_bias));
+
+    float32x4_t r2 = vmulq_f32(r, r);
+    float32x4_t p = vfmaq_laneq_f32(c1, r, ln2_c02, 2);
+    float32x4_t q = vfmaq_laneq_f32(c3, r, ln2_c02, 3);
+    q = vfmaq_f32(q, p, r2);
+    p = vmulq_f32(c4, r);
+    float32x4_t poly = vfmaq_f32(p, q, r2);
+
+    return vfmaq_f32(scale, poly, scale);
   }
   Vectorized<float> fexp_u20() const {
-    return exp();
+    return exp_u20();
   }
   DEFINE_SLEEF_COMPATIBLE_BINARY_ELEMENTWISE_FUNC_WITH_SLEEF_NAME(
       fmod,

aten/src/ATen/cuda/CUDAScaledBlas.cpp

@@ -0,0 +1,270 @@
+#include <cstdint>
+#include <c10/util/typeid.h>
+#include <c10/util/Exception.h>
+#include <c10/util/SmallVector.h>
+#include <c10/core/Scalar.h>
+#include <c10/core/ScalarType.h>
+#include <c10/util/Exception.h>
+#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
+#include <ATen/core/Tensor.h>
+#include <ATen/core/NamedTensor.h>
+#include <ATen/Dispatch.h>
+#include <ATen/ExpandUtils.h>
+#include <ATen/OpMathType.h>
+#include <ATen/TensorUtils.h>
+#include <ATen/cuda/CUDABlas.h>
+#include <ATen/cuda/tunable/Tunable.h>
+#include <ATen/cuda/tunable/TunableGemm.h>
+#include <ATen/native/Resize.h>
+#include <c10/util/MaybeOwned.h>
+#include <ATen/native/GroupedMMUtils.h>
+#include <ATen/native/cuda/RowwiseScaledMM.h>
+#include <ATen/native/cuda/ScaledGroupMM.h>
+#include <ATen/native/cuda/GroupMM.h>
+#include <ATen/ceil_div.h>
+
+#ifdef USE_FBGEMM_GENAI
+#include <fbgemm_gpu/torch_ops.h>
+#endif
+
+#ifndef AT_PER_OPERATOR_HEADERS
+#include <ATen/Functions.h>
+#include <ATen/NativeFunctions.h>
+#else
+#include <ATen/ops/_addmm_activation_native.h>
+#include <ATen/ops/_efficientzerotensor.h>
+#include <ATen/ops/_scaled_mm_native.h>
+#include <ATen/ops/_unsafe_view_native.h>
+#include <ATen/ops/abs.h>
+#include <ATen/ops/addmm_native.h>
+#include <ATen/ops/addmv_native.h>
+#include <ATen/ops/baddbmm_native.h>
+#include <ATen/ops/bmm_native.h>
+#include <ATen/ops/copy_native.h>
+#include <ATen/ops/dot_native.h>
+#include <ATen/ops/empty.h>
+#include <ATen/ops/empty_strided.h>
+#include <ATen/ops/gelu.h>
+#include <ATen/ops/max.h>
+#include <ATen/ops/mm_native.h>
+#include <ATen/ops/mul.h>
+#include <ATen/ops/relu.h>
+#include <ATen/ops/ones.h>
+#include <ATen/ops/scalar_tensor_native.h>
+#include <ATen/ops/vdot_native.h>
+#endif
+
+using at::blas::ScalingType;
+using at::blas::SwizzleType;
+
+namespace at::cuda::scaled {
+
+/**
+ * Both inputs must be fp8,
+ * Each needs a single scale, {Tensorwise (float)}
+ */
+bool check_tensorwise_recipe(c10::ScalarType type_a,
+                             std::vector<ScalingType>& recipe_a,
+                             ArrayRef<Tensor>& scales_a,
+                             c10::ScalarType type_b,
+                             std::vector<ScalingType>& recipe_b,
+                             ArrayRef<Tensor>& scales_b) {
+  // both types must be fp8
+  if (!isFloat8Type(type_a) || !isFloat8Type(type_b)) {
+    return false;
+  }
+
+  // 1 scale each, {Tensorwise, float}
+  if (scales_a.size() != 1 || recipe_a.size() != 1 || scales_b.size() != 1 || recipe_b.size() != 1) {
+    return false;
+  }
+  // Need {Blockwise_1x32, e8m0} for A & B
+  if (recipe_a[0] != ScalingType::TensorWise) return false;
+  if (scales_a[0].scalar_type() != ScalarType::Float) return false;
+  if (recipe_b[0] != ScalingType::TensorWise) return false;
+  if (scales_b[0].scalar_type() != ScalarType::Float) return false;
+
+  return true;
+}
+
+/**
+ * Both inputs must be fp8,
+ * Each needs scales, {Rowwise (float)}
+ */
+bool check_rowwise_recipe(c10::ScalarType type_a,
+                             std::vector<ScalingType>& recipe_a,
+                             ArrayRef<Tensor>& scales_a,
+                             c10::ScalarType type_b,
+                             std::vector<ScalingType>& recipe_b,
+                             ArrayRef<Tensor>& scales_b) {
+  // both types must be fp8
+  if (!isFloat8Type(type_a) || !isFloat8Type(type_b)) {
+    return false;
+  }
+
+  // 1 scale each, {Tensorwise, float}
+  if (scales_a.size() != 1 || recipe_a.size() != 1 || scales_b.size() != 1 || recipe_b.size() != 1) {
+    return false;
+  }
+
+  // Need {RowWise, dp32} for A & B
+  if (recipe_a[0] != ScalingType::RowWise) return false;
+  if (scales_a[0].scalar_type() != ScalarType::Float) return false;
+  if (recipe_b[0] != ScalingType::RowWise) return false;
+  if (scales_b[0].scalar_type() != ScalarType::Float) return false;
+
+  return true;
+}
+
+
+/**
+ * Two-level scaling, canonical NVFP4
+ * Both inputs must be fp4
+ * A, B need 2 scales, {Blockwise_1x16 (e4m3), Tensorwise (fp32)}
+ */
+bool check_nvfp4_recipe(c10::ScalarType type_a,
+                        std::vector<ScalingType>& recipe_a,
+                        ArrayRef<Tensor>& scales_a,
+                        c10::ScalarType type_b,
+                        std::vector<ScalingType>& recipe_b,
+                        ArrayRef<Tensor>& scales_b) {
+  // both types must be fp4
+  if (type_a != ScalarType::Float4_e2m1fn_x2 || type_b != ScalarType::Float4_e2m1fn_x2) {
+    return false;
+  }
+
+  // 2 scales, 2 recipes for each input
+  if (scales_a.size() != 2 || recipe_a.size() != 2 || scales_b.size() != 2 || recipe_b.size() != 2) {
+    return false;
+  }
+
+  // Need {Blockwise_1x16, e4m3 for scale[0], Tensorwise, fp32 for scale[1]}
+  if (recipe_a[0] != ScalingType::BlockWise1x16 || recipe_a[1] != ScalingType::TensorWise) return false;
+  if (scales_a[0].scalar_type() != ScalarType::Float8_e4m3fn || scales_a[1].scalar_type() != ScalarType::Float) return false;
+  if (recipe_b[0] != ScalingType::BlockWise1x16 || recipe_b[1] != ScalingType::TensorWise) return false;
+  if (scales_b[0].scalar_type() != ScalarType::Float8_e4m3fn || scales_b[1].scalar_type() != ScalarType::Float) return false;
+
+  return true;
+}
+
+/**
+ * Single-level scaling, what PyT currently understands
+ * Both inputs must be fp4
+ * A, B need 1 scale, {Blockwise_1x16 (e4m3)}
+ */
+bool check_nvfp4_recipe_single_scale
+                       (c10::ScalarType type_a,
+                        std::vector<ScalingType>& recipe_a,
+                        ArrayRef<Tensor>& scales_a,
+                        c10::ScalarType type_b,
+                        std::vector<ScalingType>& recipe_b,
+                        ArrayRef<Tensor>& scales_b) {
+  // both types must be fp4
+  if (type_a != ScalarType::Float4_e2m1fn_x2 || type_b != ScalarType::Float4_e2m1fn_x2) {
+    return false;
+  }
+
+  // 2 scales, 2 recipes for each input
+  if (scales_a.size() != 1 || recipe_a.size() != 1 || scales_b.size() != 1 || recipe_b.size() != 1) {
+    return false;
+  }
+
+  // Need {Blockwise_1x16, e4m3 for scale[0], Tensorwise, fp32 for scale[1]}
+  if (recipe_a[0] != ScalingType::BlockWise1x16) return false;
+  if (scales_a[0].scalar_type() != ScalarType::Float8_e4m3fn) return false;
+  if (recipe_b[0] != ScalingType::BlockWise1x16) return false;
+  if (scales_b[0].scalar_type() != ScalarType::Float8_e4m3fn) return false;
+
+  return true;
+}
+
+/**
+ * Both inputs must be fp8
+ * A, B must only have 1 scale each, A: {Blockwise_1x128 (float), B: {Blockwise_128x128 (float)
+ */
+bool check_deepseek_recipe(ScalingType expected_recipe_a,
+                           ScalingType expected_recipe_b,
+                           c10::ScalarType type_a,
+                           std::vector<ScalingType>& recipe_a,
+                           ArrayRef<Tensor>& scales_a,
+                           c10::ScalarType type_b,
+                           std::vector<ScalingType>& recipe_b,
+                           ArrayRef<Tensor>& scales_b) {
+  // both types must be fp8
+  if (type_a != ScalarType::Float8_e4m3fn || type_b != ScalarType::Float8_e4m3fn) {
+    return false;
+  }
+
+  // 1 scales, 1 recipes for each input
+  if (scales_a.size() != 1 || recipe_a.size() != 1 || scales_b.size() != 1 || recipe_b.size() != 1) {
+    return false;
+  }
+
+  // Need {Blockwise_1x128, float} for A, {Blockwise_128x128, float} for B
+  if (recipe_a[0] != expected_recipe_a) return false;
+  if (scales_a[0].scalar_type() != ScalarType::Float) return false;
+  if (recipe_b[0] != expected_recipe_b) return false;
+  if (scales_b[0].scalar_type() != ScalarType::Float) return false;
+
+  return true;
+}
+
+/**
+ * Both inputs must be fp8
+ * A, B must have 1 scale each, {Blockwise_1x32, e8m0}
+ */
+bool check_mxfp8_recipe(c10::ScalarType type_a,
+                        std::vector<ScalingType>& recipe_a,
+                        ArrayRef<Tensor>& scales_a,
+                        c10::ScalarType type_b,
+                        std::vector<ScalingType>& recipe_b,
+                        ArrayRef<Tensor>& scales_b) {
+  // both types must be fp8
+  if (type_a != ScalarType::Float8_e4m3fn || type_b != ScalarType::Float8_e4m3fn) {
+    return false;
+  }
+
+  // 1 scales, 1 recipes for each input
+  if (scales_a.size() != 1 || recipe_a.size() != 1 || scales_b.size() != 1 || recipe_b.size() != 1) {
+    return false;
+  }
+
+  // Need {Blockwise_1x32, e8m0} for A & B
+  if (recipe_a[0] != ScalingType::BlockWise1x32) return false;
+  if (scales_a[0].scalar_type() != ScalarType::Float8_e8m0fnu) return false;
+  if (recipe_b[0] != ScalingType::BlockWise1x32) return false;
+  if (scales_b[0].scalar_type() != ScalarType::Float8_e8m0fnu) return false;
+
+  return true;
+}
+
+/**
+ * Both inputs must be fp4
+ * A, B must have 1 scale each, {Blockwise_1x32, e8m0}
+ */
+bool check_mxfp4_recipe(c10::ScalarType type_a,
+                        std::vector<ScalingType>& recipe_a,
+                        ArrayRef<Tensor>& scales_a,
+                        c10::ScalarType type_b,
+                        std::vector<ScalingType>& recipe_b,
+                        ArrayRef<Tensor>& scales_b) {
+  // both types must be fp4
+  if (type_a != ScalarType::Float4_e2m1fn_x2 || type_b != ScalarType::Float4_e2m1fn_x2) {
+    return false;
+  }
+
+  // 1 scales, 1 recipes for each input
+  if (scales_a.size() != 1 || recipe_a.size() != 1 || scales_b.size() != 1 || recipe_b.size() != 1) {
+    return false;
+  }
+
+  // Need {Blockwise_1x32, e8m0} for A & B
+  if (recipe_a[0] != ScalingType::BlockWise1x32) return false;
+  if (scales_a[0].scalar_type() != ScalarType::Float8_e8m0fnu) return false;
+  if (recipe_b[0] != ScalingType::BlockWise1x32) return false;
+  if (scales_b[0].scalar_type() != ScalarType::Float8_e8m0fnu) return false;
+
+  return true;
+}
+
+} // namespace at::native::cuda::blas::scaled

aten/src/ATen/cuda/CUDAScaledBlas.h

@@ -0,0 +1,174 @@
+#include <cstdint>
+#include <c10/util/typeid.h>
+#include <c10/util/Exception.h>
+#include <c10/util/SmallVector.h>
+#include <c10/core/Scalar.h>
+#include <c10/core/ScalarType.h>
+#include <c10/util/Exception.h>
+#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
+#include <ATen/core/Tensor.h>
+#include <ATen/core/NamedTensor.h>
+#include <ATen/Dispatch.h>
+#include <ATen/ExpandUtils.h>
+#include <ATen/OpMathType.h>
+#include <ATen/TensorUtils.h>
+#include <ATen/cuda/CUDABlas.h>
+#include <ATen/cuda/tunable/Tunable.h>
+#include <ATen/cuda/tunable/TunableGemm.h>
+#include <ATen/native/Resize.h>
+#include <c10/util/MaybeOwned.h>
+#include <ATen/native/GroupedMMUtils.h>
+#include <ATen/native/cuda/RowwiseScaledMM.h>
+#include <ATen/native/cuda/ScaledGroupMM.h>
+#include <ATen/native/cuda/GroupMM.h>
+#include <ATen/ceil_div.h>
+
+#ifdef USE_FBGEMM_GENAI
+#include <fbgemm_gpu/torch_ops.h>
+#endif
+
+#ifndef AT_PER_OPERATOR_HEADERS
+#include <ATen/Functions.h>
+#include <ATen/NativeFunctions.h>
+#else
+#include <ATen/ops/_addmm_activation_native.h>
+#include <ATen/ops/_efficientzerotensor.h>
+#include <ATen/ops/_scaled_mm_native.h>
+#include <ATen/ops/_unsafe_view_native.h>
+#include <ATen/ops/abs.h>
+#include <ATen/ops/addmm_native.h>
+#include <ATen/ops/addmv_native.h>
+#include <ATen/ops/baddbmm_native.h>
+#include <ATen/ops/bmm_native.h>
+#include <ATen/ops/copy_native.h>
+#include <ATen/ops/dot_native.h>
+#include <ATen/ops/empty.h>
+#include <ATen/ops/empty_strided.h>
+#include <ATen/ops/gelu.h>
+#include <ATen/ops/max.h>
+#include <ATen/ops/mm_native.h>
+#include <ATen/ops/mul.h>
+#include <ATen/ops/relu.h>
+#include <ATen/ops/ones.h>
+#include <ATen/ops/scalar_tensor_native.h>
+#include <ATen/ops/vdot_native.h>
+#endif
+
+using at::blas::ScalingType;
+using at::blas::SwizzleType;
+
+namespace at::cuda::scaled {
+
+static bool _scaled_mm_allowed_device(bool sm90_only=false, bool sm100_only=false) {
+#ifdef USE_ROCM
+    static const std::vector<std::string> archs = {
+        "gfx942",
+#if ROCM_VERSION >= 60300
+        "gfx1200", "gfx1201",
+#endif
+#if ROCM_VERSION >= 60500
+        "gfx950"
+#endif
+    };
+    return at::detail::getCUDAHooks().isGPUArch(archs);
+#else
+    auto dprops = at::cuda::getCurrentDeviceProperties();
+
+    if (sm90_only || sm100_only) {
+      return (sm90_only && dprops->major == 9) || (sm100_only && dprops->major == 10);
+    } else {
+      return dprops->major >= 9 || (dprops->major == 8 && dprops->minor == 9);
+    }
+#endif
+}
+
+#ifdef USE_ROCM
+static bool _scaled_mm_is_fnuz() {
+    return at::detail::getCUDAHooks().isGPUArch({"gfx942"});
+}
+#endif
+/**
+ * Track concrete implementations available
+ */
+enum class ScaledGemmImplementation {
+  NONE = 0,
+  TENSORWISE_TENSORWISE = 1,
+  ROWWISE_ROWWISE = 2,
+  BLOCK_128x128_1x128 = 3,
+  BLOCK_1x128_128x128 = 4,
+  BLOCK_1x128_1x128 = 5,
+  MXFP8_MXFP8 = 6,
+  NVFP4_NVFP4 = 7,
+  NVFP4_NVFP4_SINGLE_SCALE = 8,
+  MXFP4_MXFP4 = 9,
+};
+
+/**
+ * Convert passed int (enum) from python back into a
+ * strictly-typed enum
+ */
+template <class EnumType, class ArrayType>
+std::vector<EnumType> convert_int_to_enum(ArrayType& v) {
+  std::vector<EnumType> converted;
+  converted.reserve(v.size());
+
+  for (auto vi : v) {
+    converted.push_back(static_cast<EnumType>(vi));
+  }
+  return converted;
+}
+
+bool check_tensorwise_recipe(c10::ScalarType,
+                             std::vector<ScalingType>&,
+                             ArrayRef<Tensor>&,
+                             c10::ScalarType,
+                             std::vector<ScalingType>&,
+                             ArrayRef<Tensor>&);
+
+
+bool check_rowwise_recipe(c10::ScalarType,
+                             std::vector<ScalingType>&,
+                             ArrayRef<Tensor>&,
+                             c10::ScalarType,
+                             std::vector<ScalingType>&,
+                             ArrayRef<Tensor>&);
+
+bool check_nvfp4_recipe(c10::ScalarType,
+                        std::vector<ScalingType>&,
+                        ArrayRef<Tensor>&,
+                        c10::ScalarType,
+                        std::vector<ScalingType>&,
+                        ArrayRef<Tensor>&);
+
+bool check_nvfp4_recipe_single_scale
+                       (c10::ScalarType,
+                        std::vector<ScalingType>&,
+                        ArrayRef<Tensor>&,
+                        c10::ScalarType,
+                        std::vector<ScalingType>&,
+                        ArrayRef<Tensor>&);
+
+bool check_deepseek_recipe(ScalingType,
+                           ScalingType,
+                           c10::ScalarType,
+                           std::vector<ScalingType>&,
+                           ArrayRef<Tensor>&,
+                           c10::ScalarType,
+                           std::vector<ScalingType>&,
+                           ArrayRef<Tensor>&);
+
+bool check_mxfp8_recipe(c10::ScalarType,
+                        std::vector<ScalingType>&,
+                        ArrayRef<Tensor>&,
+                        c10::ScalarType,
+                        std::vector<ScalingType>&,
+                        ArrayRef<Tensor>&);
+
+bool check_mxfp4_recipe(c10::ScalarType,
+                        std::vector<ScalingType>&,
+                        ArrayRef<Tensor>&,
+                        c10::ScalarType,
+                        std::vector<ScalingType>&,
+                        ArrayRef<Tensor>&);
+
+} // namespace at::native::cuda::blas::scaled

aten/src/ATen/native/cuda/Blas.cpp

@@ -13,6 +13,7 @@
 #include <ATen/OpMathType.h>
 #include <ATen/TensorUtils.h>
 #include <ATen/cuda/CUDABlas.h>
+#include <ATen/cuda/CUDAScaledBlas.h>
 #include <ATen/cuda/tunable/Tunable.h>
 #include <ATen/cuda/tunable/TunableGemm.h>
 #include <ATen/native/Resize.h>
@@ -1628,104 +1629,6 @@ _scaled_mm_out_cuda(const Tensor& mat1, const Tensor& mat2,
   return _scaled_gemm(mat1, mat2, scale_a, scale_b, scaling_choice_a, scaling_choice_b, bias, use_fast_accum, out);
 }
 
-namespace {
-  void _check_scales_fp8_rowwise(const Tensor& mat, const Tensor& scale, const int dim, const int arg_idx, const int scale_multiplier=1) {
-    // Checks scales for 2d or 3d target tensors (`mat`).
-    if (mat.dim() == 2) {
-      TORCH_CHECK(
-          scale.dim() == 1,
-          "scale must be a 1D tensor, but got ",
-          scale.dim(),
-          "D, arg ",
-          arg_idx);
-      TORCH_CHECK(
-          scale.is_contiguous(), "scale must be contiguous for arg ", arg_idx);
-      TORCH_CHECK(
-          scale.size(0) == mat.size(dim) * scale_multiplier,
-          "scale must have the same length as mat for arg ",
-          arg_idx);
-    } else {
-      TORCH_CHECK(
-          scale.dim() == 2,
-          "scale must be a 2D tensor, but got ",
-          scale.dim(),
-          "D for arg ",
-          arg_idx);
-      TORCH_CHECK(
-          scale.stride(1) == 1,
-          "scale must be contiguous in the last dimension for arg ",
-          arg_idx);
-      TORCH_CHECK(
-          scale.size(0) == mat.size(0),
-          "scale must have the same batch dimension as mat for arg ",
-          arg_idx);
-      TORCH_CHECK(
-          scale.size(1) == mat.size(1 + dim),
-          "scale must have the same first dimension as mat for arg ",
-          arg_idx);
-    }
-  }
-
-  void _check_scales_mxfp8(const Tensor& mat, const Tensor& scale, const int dim, const int arg_idx) {
-    // Checks scales for 2d or 3d target tensors (`mat`).
-    if (mat.dim() == 2) {
-      // For MXFP8, 2d tensors have variable size groups represented as subtensors,
-      // that are converted to blocked padded format individually,
-      // so we can't check the scale sizes without doing a d2h sync to get the group sizes here.
-      TORCH_CHECK(
-        scale.dim() == mat.dim(),
-        "for mxfp8, scale must have same number of dimensions as parent tensor, but got mat.dim() = ", mat.dim(), " and scale.dim() = ", scale.dim(), " for arg ", arg_idx);
-
-      // LHS mat shape (M, total_K) -> scale shape (rounded_up(M, 128), rounded_up_per_group(K/32, 4))
-      // RHS mat shape (total_K, N) -> scale shape (rounded_up(N, 128), rounded_up_per_group(K/32, 4))
-      //   * weight is transposed prior to the call, scale stays non-transposed.
-      bool LHS = arg_idx == 0;
-      int scale_dim_to_check = 0;
-      int mat_dim_to_check = LHS ? 0 : 1;
-      TORCH_CHECK(
-          scale.size(scale_dim_to_check) >= mat.size(mat_dim_to_check),
-          "for mxfp8, arg ", arg_idx, " tensor shape (", mat.size(0), ", ", mat.size(1), ") ",
-          "must have scale.shape[", scale_dim_to_check, "] >= ", mat.size(mat_dim_to_check), " but got scale.shape=(", scale.size(0), ", ", scale.size(1), ")");
-    } else {
-      // For MXFP8, 3d tensors have static group sizes (stack of 2d tensors),
-      // so we can check the exact expected scale sizes here without a d2h sync.
-      auto round_up = [](auto x, auto y) {
-          return ((x + y - 1) / y) * y;
-      };
-
-      // TODO: this is for 3d tensor in 2d-3d case specifically.
-      // We'll need to support 3d-3d and 3d-2d cases once mxfp8 grouped gemm supports them.
-      int64_t G = mat.size(0);
-      int64_t K = mat.size(1);
-      int64_t N = mat.size(2);
-      int64_t blocked_scale_K = round_up(K/32, 4);
-      int64_t blocked_scale_N = round_up(N, 128);
-
-      // fbgemm expects stack of flattened blocked scales for 3d tensor, shape (G, blocked_scale_K * blocked_scale_N).
-      TORCH_CHECK(
-        scale.dim() == mat.dim() - 1,
-        "for mxfp8 2d-3d grouped GEMM, the 3d tensor of shape (G,K,N) must have a 2d scale of shape (G, blocked_scale_K * blocked_scale_N), but scale is ", scale.dim(), "D for arg ", arg_idx
-      );
-      TORCH_CHECK(
-        scale.size(0) == G && scale.size(1) == blocked_scale_K * blocked_scale_N,
-        "for mxfp8, the tensor shape (", G, ", ", K, ", ", N, ") must have scale shape (", G, ",", blocked_scale_K, ",", blocked_scale_N, ") for arg ", arg_idx
-      );
-    }
-  }
-
-  void check_scale(const Tensor& mat, const Tensor& scale, const int dim, const int arg_idx, const int scale_multiplier=1) {
-    bool using_fp8_rowwise = scale.scalar_type() == kFloat;
-    bool using_mxfp8 = scale.scalar_type() == at::kFloat8_e8m0fnu;
-    if (using_fp8_rowwise) {
-      _check_scales_fp8_rowwise(mat, scale, dim, arg_idx, scale_multiplier);
-    } else if (using_mxfp8) {
-      _check_scales_mxfp8(mat, scale, dim, arg_idx);
-    } else {
-      TORCH_CHECK(false, "scale must be float32 or float8_e8m0fnu, but got ", scale.dtype());
-    }
-  }
-}
-
 Tensor
 _scaled_mm_cuda(const Tensor& mat_a, const Tensor& mat_b,
           const Tensor& scale_a,
@@ -1740,261 +1643,26 @@ _scaled_mm_cuda(const Tensor& mat_a, const Tensor& mat_b,
   return _scaled_mm_out_cuda(mat_a, mat_b, scale_a, scale_b, bias, scale_result, out_dtype, use_fast_accum, out);
 }
 
-/**
- * Track concrete implementations available
- */
-enum class ScaledGemmImplementation {
-  NONE = 0,
-  TENSORWISE_TENSORWISE = 1,
-  ROWWISE_ROWWISE = 2,
-  BLOCK_128x128_1x128 = 3,
-  BLOCK_1x128_128x128 = 4,
-  BLOCK_1x128_1x128 = 5,
-  MXFP8_MXFP8 = 6,
-  NVFP4_NVFP4 = 7,
-  NVFP4_NVFP4_SINGLE_SCALE = 8,
-  MXFP4_MXFP4 = 9,
-};
-
-/**
- * Convert passed int (enum) from python back into a
- * strictly-typed enum
- */
-template <class EnumType, class ArrayType>
-std::vector<EnumType> convert_int_to_enum(ArrayType& v) {
-  std::vector<EnumType> converted;
-  converted.reserve(v.size());
-
-  for (auto vi : v) {
-    converted.push_back(static_cast<EnumType>(vi));
-  }
-  return converted;
-}
-
-/**
- * Both inputs must be fp8,
- * Each needs a single scale, {Tensorwise (float)}
- */
-bool check_tensorwise_recipe(c10::ScalarType type_a,
-                             std::vector<ScalingType>& recipe_a,
-                             ArrayRef<Tensor>& scales_a,
-                             c10::ScalarType type_b,
-                             std::vector<ScalingType>& recipe_b,
-                             ArrayRef<Tensor>& scales_b) {
-  // both types must be fp8
-  if (!isFloat8Type(type_a) || !isFloat8Type(type_b)) {
-    return false;
-  }
-
-  // 1 scale each, {Tensorwise, float}
-  if (scales_a.size() != 1 || recipe_a.size() != 1 || scales_b.size() != 1 || recipe_b.size() != 1) {
-    return false;
-  }
-  // Need {Blockwise_1x32, e8m0} for A & B
-  if (recipe_a[0] != ScalingType::TensorWise) return false;
-  if (scales_a[0].scalar_type() != ScalarType::Float) return false;
-  if (recipe_b[0] != ScalingType::TensorWise) return false;
-  if (scales_b[0].scalar_type() != ScalarType::Float) return false;
-
-  return true;
-}
-
-/**
- * Both inputs must be fp8,
- * Each needs scales, {Rowwise (float)}
- */
-bool check_rowwise_recipe(c10::ScalarType type_a,
-                             std::vector<ScalingType>& recipe_a,
-                             ArrayRef<Tensor>& scales_a,
-                             c10::ScalarType type_b,
-                             std::vector<ScalingType>& recipe_b,
-                             ArrayRef<Tensor>& scales_b) {
-  // both types must be fp8
-  if (!isFloat8Type(type_a) || !isFloat8Type(type_b)) {
-    return false;
-  }
-
-  // 1 scale each, {Tensorwise, float}
-  if (scales_a.size() != 1 || recipe_a.size() != 1 || scales_b.size() != 1 || recipe_b.size() != 1) {
-    return false;
-  }
-
-  // Need {RowWise, dp32} for A & B
-  if (recipe_a[0] != ScalingType::RowWise) return false;
-  if (scales_a[0].scalar_type() != ScalarType::Float) return false;
-  if (recipe_b[0] != ScalingType::RowWise) return false;
-  if (scales_b[0].scalar_type() != ScalarType::Float) return false;
-
-  return true;
-}
-
-
-/**
- * Two-level scaling, canonical NVFP4
- * Both inputs must be fp4
- * A, B need 2 scales, {Blockwise_1x16 (e4m3), Tensorwise (fp32)}
- */
-bool check_nvfp4_recipe(c10::ScalarType type_a,
-                        std::vector<ScalingType>& recipe_a,
-                        ArrayRef<Tensor>& scales_a,
-                        c10::ScalarType type_b,
-                        std::vector<ScalingType>& recipe_b,
-                        ArrayRef<Tensor>& scales_b) {
-  // both types must be fp4
-  if (type_a != ScalarType::Float4_e2m1fn_x2 || type_b != ScalarType::Float4_e2m1fn_x2) {
-    return false;
-  }
-
-  // 2 scales, 2 recipes for each input
-  if (scales_a.size() != 2 || recipe_a.size() != 2 || scales_b.size() != 2 || recipe_b.size() != 2) {
-    return false;
-  }
-
-  // Need {Blockwise_1x16, e4m3 for scale[0], Tensorwise, fp32 for scale[1]}
-  if (recipe_a[0] != ScalingType::BlockWise1x16 || recipe_a[1] != ScalingType::TensorWise) return false;
-  if (scales_a[0].scalar_type() != ScalarType::Float8_e4m3fn || scales_a[1].scalar_type() != ScalarType::Float) return false;
-  if (recipe_b[0] != ScalingType::BlockWise1x16 || recipe_b[1] != ScalingType::TensorWise) return false;
-  if (scales_b[0].scalar_type() != ScalarType::Float8_e4m3fn || scales_b[1].scalar_type() != ScalarType::Float) return false;
-
-  return true;
-}
-
-/**
- * Single-level scaling, what PyT currently understands
- * Both inputs must be fp4
- * A, B need 1 scale, {Blockwise_1x16 (e4m3)}
- */
-bool check_nvfp4_recipe_single_scale
-                       (c10::ScalarType type_a,
-                        std::vector<ScalingType>& recipe_a,
-                        ArrayRef<Tensor>& scales_a,
-                        c10::ScalarType type_b,
-                        std::vector<ScalingType>& recipe_b,
-                        ArrayRef<Tensor>& scales_b) {
-  // both types must be fp4
-  if (type_a != ScalarType::Float4_e2m1fn_x2 || type_b != ScalarType::Float4_e2m1fn_x2) {
-    return false;
-  }
-
-  // 2 scales, 2 recipes for each input
-  if (scales_a.size() != 1 || recipe_a.size() != 1 || scales_b.size() != 1 || recipe_b.size() != 1) {
-    return false;
-  }
-
-  // Need {Blockwise_1x16, e4m3 for scale[0], Tensorwise, fp32 for scale[1]}
-  if (recipe_a[0] != ScalingType::BlockWise1x16) return false;
-  if (scales_a[0].scalar_type() != ScalarType::Float8_e4m3fn) return false;
-  if (recipe_b[0] != ScalingType::BlockWise1x16) return false;
-  if (scales_b[0].scalar_type() != ScalarType::Float8_e4m3fn) return false;
-
-  return true;
-}
-
-/**
- * Both inputs must be fp8
- * A, B must only have 1 scale each, A: {Blockwise_1x128 (float), B: {Blockwise_128x128 (float)
- */
-bool check_deepseek_recipe(ScalingType expected_recipe_a,
-                           ScalingType expected_recipe_b,
-                           c10::ScalarType type_a,
-                           std::vector<ScalingType>& recipe_a,
-                           ArrayRef<Tensor>& scales_a,
-                           c10::ScalarType type_b,
-                           std::vector<ScalingType>& recipe_b,
-                           ArrayRef<Tensor>& scales_b) {
-  // both types must be fp8
-  if (type_a != ScalarType::Float8_e4m3fn || type_b != ScalarType::Float8_e4m3fn) {
-    return false;
-  }
-
-  // 1 scales, 1 recipes for each input
-  if (scales_a.size() != 1 || recipe_a.size() != 1 || scales_b.size() != 1 || recipe_b.size() != 1) {
-    return false;
-  }
-
-  // Need {Blockwise_1x128, float} for A, {Blockwise_128x128, float} for B
-  if (recipe_a[0] != expected_recipe_a) return false;
-  if (scales_a[0].scalar_type() != ScalarType::Float) return false;
-  if (recipe_b[0] != expected_recipe_b) return false;
-  if (scales_b[0].scalar_type() != ScalarType::Float) return false;
-
-  return true;
-}
-
-/**
- * Both inputs must be fp8
- * A, B must have 1 scale each, {Blockwise_1x32, e8m0}
- */
-bool check_mxfp8_recipe(c10::ScalarType type_a,
-                        std::vector<ScalingType>& recipe_a,
-                        ArrayRef<Tensor>& scales_a,
-                        c10::ScalarType type_b,
-                        std::vector<ScalingType>& recipe_b,
-                        ArrayRef<Tensor>& scales_b) {
-  // both types must be fp8
-  if (type_a != ScalarType::Float8_e4m3fn || type_b != ScalarType::Float8_e4m3fn) {
-    return false;
-  }
-
-  // 1 scales, 1 recipes for each input
-  if (scales_a.size() != 1 || recipe_a.size() != 1 || scales_b.size() != 1 || recipe_b.size() != 1) {
-    return false;
-  }
-
-  // Need {Blockwise_1x32, e8m0} for A & B
-  if (recipe_a[0] != ScalingType::BlockWise1x32) return false;
-  if (scales_a[0].scalar_type() != ScalarType::Float8_e8m0fnu) return false;
-  if (recipe_b[0] != ScalingType::BlockWise1x32) return false;
-  if (scales_b[0].scalar_type() != ScalarType::Float8_e8m0fnu) return false;
-
-  return true;
-}
-
-/**
- * Both inputs must be fp4
- * A, B must have 1 scale each, {Blockwise_1x32, e8m0}
- */
-bool check_mxfp4_recipe(c10::ScalarType type_a,
-                        std::vector<ScalingType>& recipe_a,
-                        ArrayRef<Tensor>& scales_a,
-                        c10::ScalarType type_b,
-                        std::vector<ScalingType>& recipe_b,
-                        ArrayRef<Tensor>& scales_b) {
-  // both types must be fp4
-  if (type_a != ScalarType::Float4_e2m1fn_x2 || type_b != ScalarType::Float4_e2m1fn_x2) {
-    return false;
-  }
-
-  // 1 scales, 1 recipes for each input
-  if (scales_a.size() != 1 || recipe_a.size() != 1 || scales_b.size() != 1 || recipe_b.size() != 1) {
-    return false;
-  }
-
-  // Need {Blockwise_1x32, e8m0} for A & B
-  if (recipe_a[0] != ScalingType::BlockWise1x32) return false;
-  if (scales_a[0].scalar_type() != ScalarType::Float8_e8m0fnu) return false;
-  if (recipe_b[0] != ScalingType::BlockWise1x32) return false;
-  if (scales_b[0].scalar_type() != ScalarType::Float8_e8m0fnu) return false;
-
-  return true;
-}
-
 using acceptance_fn = std::function<bool(c10::ScalarType, std::vector<ScalingType>&, ArrayRef<Tensor>&, c10::ScalarType, std::vector<ScalingType>&, ArrayRef<Tensor>&)>;
 using namespace std::placeholders;
 
+namespace scaled_blas = at::cuda::scaled;
+using scaled_blas::ScaledGemmImplementation;
+using scaled_blas::convert_int_to_enum;
+
 std::array<std::tuple<std::string, acceptance_fn, ScaledGemmImplementation>, 9> scale_kernel_dispatch = {{
-  { "tensorwise_tensorwise", check_tensorwise_recipe, ScaledGemmImplementation::TENSORWISE_TENSORWISE },
-  { "rowwise_rowwise", check_rowwise_recipe, ScaledGemmImplementation::ROWWISE_ROWWISE},
-  { "block_1x128_128x128", std::bind(check_deepseek_recipe, ScalingType::BlockWise1x128, ScalingType::BlockWise128x128, _1, _2, _3, _4, _5, _6),
+  { "tensorwise_tensorwise", scaled_blas::check_tensorwise_recipe, ScaledGemmImplementation::TENSORWISE_TENSORWISE },
+  { "rowwise_rowwise", scaled_blas::check_rowwise_recipe, ScaledGemmImplementation::ROWWISE_ROWWISE},
+  { "block_1x128_128x128", std::bind(scaled_blas::check_deepseek_recipe, ScalingType::BlockWise1x128, ScalingType::BlockWise128x128, _1, _2, _3, _4, _5, _6),
     ScaledGemmImplementation::BLOCK_1x128_128x128},
-  { "block_128x128_1x128", std::bind(check_deepseek_recipe, ScalingType::BlockWise128x128, ScalingType::BlockWise1x128, _1, _2, _3, _4, _5, _6),
+  { "block_128x128_1x128", std::bind(scaled_blas::check_deepseek_recipe, ScalingType::BlockWise128x128, ScalingType::BlockWise1x128, _1, _2, _3, _4, _5, _6),
     ScaledGemmImplementation::BLOCK_128x128_1x128},
-  { "block_1x128_1x128", std::bind(check_deepseek_recipe, ScalingType::BlockWise1x128, ScalingType::BlockWise1x128, _1, _2, _3, _4, _5, _6),
+  { "block_1x128_1x128", std::bind(scaled_blas::check_deepseek_recipe, ScalingType::BlockWise1x128, ScalingType::BlockWise1x128, _1, _2, _3, _4, _5, _6),
     ScaledGemmImplementation::BLOCK_1x128_1x128},
-  { "nvfp4_nvfp4", check_nvfp4_recipe, ScaledGemmImplementation::NVFP4_NVFP4},
-  { "nvfp4_nvfp4_single_scale", check_nvfp4_recipe_single_scale, ScaledGemmImplementation::NVFP4_NVFP4_SINGLE_SCALE },
-  { "mxfp8_mxfp8", check_mxfp8_recipe, ScaledGemmImplementation::MXFP8_MXFP8},
-  { "mxfp4_mxfp4", check_mxfp4_recipe, ScaledGemmImplementation::MXFP4_MXFP4}}};
+  { "nvfp4_nvfp4", scaled_blas::check_nvfp4_recipe, ScaledGemmImplementation::NVFP4_NVFP4},
+  { "nvfp4_nvfp4_single_scale", scaled_blas::check_nvfp4_recipe_single_scale, ScaledGemmImplementation::NVFP4_NVFP4_SINGLE_SCALE },
+  { "mxfp8_mxfp8", scaled_blas::check_mxfp8_recipe, ScaledGemmImplementation::MXFP8_MXFP8},
+  { "mxfp4_mxfp4", scaled_blas::check_mxfp4_recipe, ScaledGemmImplementation::MXFP4_MXFP4}}};
 
 Tensor&
 _scaled_tensorwise_tensorwise(
@@ -2596,410 +2264,6 @@ _scaled_mm_cuda_v2(
                       out);
 }
 
-// 2d-2d and 2d-3d
-// scaling=MXFP8
-// CUDA-only
-Tensor&
-_mx8_mx8_bf16_grouped_mm_fbgemm(
-        const Tensor& mat_a,
-        const Tensor& mat_b,
-        const Tensor& scale_a,
-        const SwizzleType& swizzle_a,
-        const Tensor& scale_b,
-        const SwizzleType& swizzle_b,
-        const std::optional<at::Tensor>& offs,
-        Tensor& out) {
-    const bool a_is_2d = mat_a.dim() == 2;
-    const bool b_is_2d = mat_b.dim() == 2;
-    bool b_is_3d = mat_b.dim() == 3;
-    bool is_2d_2d = a_is_2d && b_is_2d;
-    bool is_2d_3d = a_is_2d && b_is_3d;
-    TORCH_CHECK_VALUE(is_2d_2d || is_2d_3d, "MXFP8 grouped GEMM currently only supports 2d-2d and 2d-3d cases");
-    TORCH_CHECK_VALUE(offs.has_value(), "MXFP8 2d-2d and 2d-3d grouped GEMMs requires offsets");
-    TORCH_CHECK_VALUE(out.scalar_type() == at::kBFloat16, "Only bf16 out_dtype is supported for MXFP8 grouped gemm");
-    // MXFP8 expects float8_e8m0fnu scales.
-    TORCH_CHECK_VALUE(scale_a.scalar_type() == at::kFloat8_e8m0fnu && scale_b.scalar_type() == at::kFloat8_e8m0fnu,
-        "For MXFP8 grouped gemm, both scales must be float8_e8m0fnu tensors.");
-#ifdef USE_ROCM
-    TORCH_CHECK_VALUE(swizzle_a == SwizzleType::NO_SWIZZLE && swizzle_b == SwizzleType::NO_SWIZZLE,
-        "For ROCM MXFP8 grouped gemm, both scale swizzle types must be SWIZZLE_NONE");
-#else
-    TORCH_CHECK_VALUE(swizzle_a == SwizzleType::SWIZZLE_32_4_4 && swizzle_b == SwizzleType::SWIZZLE_32_4_4,
-        "For CUDA MXFP8 grouped gemm, both scale swizzle types must be SWIZZLE_32_4_4");
-#endif
-
-#if defined(USE_FBGEMM_GENAI) and !defined(USE_ROCM)
-    fbgemm_gpu::mx8mx8bf16_grouped_mm(
-        mat_a,
-        mat_b,
-        scale_a,
-        scale_b,
-        offs.value(),
-        out);
-#else
-    TORCH_CHECK_NOT_IMPLEMENTED(false, "mxfp8_mxfp8 grouped gemm requires compile with USE_FBGEMM_GENAI");
-#endif
-    return out;
-}
-
-// 2d-2d and 2d-3d cases
-// scaling=rowwise
-// CUDA-only
-Tensor&
-_f8_f8_bf16_rowwise_grouped_mm_cuda(
-          const Tensor& mat_a,
-          const Tensor& mat_b,
-          const Tensor& scale_a,
-          const Tensor& scale_b,
-          const std::optional<Tensor>& offs,
-          const std::optional<Tensor>& bias,
-          const bool use_fast_accum,
-          Tensor& out) {
-  TORCH_CHECK_VALUE(mat_a.dtype() == at::kFloat8_e4m3fn, "Expected mat_a to be Float8_e4m3 matrix got ", mat_a.scalar_type());
-  TORCH_CHECK_VALUE(mat_b.dtype() == at::kFloat8_e4m3fn, "Expected mat_a to be Float8_e4m3 matrix got ", mat_b.scalar_type());
-
-  at::cuda::detail::f8f8bf16_grouped_mm(
-      mat_a,
-      mat_b,
-      scale_a,
-      scale_b,
-      offs,
-      bias,
-      use_fast_accum,
-      out);
-    return out;
-}
-
-// 2d-2d and 2d-3d cases
-// scaling=rowwise
-// only being called for rocm
-Tensor&
-_f8_f8_bf16_rowwise_grouped_mm_rocm(
-      const Tensor& mat_a,
-      const Tensor& mat_b,
-      const Tensor& scale_a,
-      const Tensor& scale_b,
-      const std::optional<Tensor>& offs,
-      Tensor& out) {
-  TORCH_CHECK_VALUE(mat_a.dtype() == at::kFloat8_e4m3fnuz, "Expected mat_a to be Float8_e4m3fnuz matrix got ", mat_a.scalar_type());
-  TORCH_CHECK_VALUE(mat_b.dtype() == at::kFloat8_e4m3fnuz, "Expected mat_a to be Float8_e4m3fnuz matrix got ", mat_b.scalar_type());
-
-#if defined(USE_FBGEMM_GENAI) && defined(USE_ROCM)
-  fbgemm_gpu::f8f8bf16_rowwise_grouped_mm(
-      mat_a,
-      // FBGEMM expects B matrix shape to be (.., N, K)
-      mat_b.transpose(-2, -1),
-      scale_a,
-      scale_b,
-      offs,
-      out);
-#else
-  TORCH_CHECK_NOT_IMPLEMENTED(false, "grouped gemm is not supported without USE_FBGEMM_GENAI on ROCM")
-#endif
-  return out;
-
-}
-
-// Dispatch f8 x f8 -> bf16 row-wise scaled to rocm/cuda
-Tensor&
-_f8_f8_bf16_rowwise_grouped_mm(
-      const Tensor& mat_a,
-      const Tensor& mat_b,
-      const Tensor& scale_a,
-      const Tensor& scale_b,
-      const std::optional<Tensor>& offs,
-      const std::optional<Tensor>& bias,
-      bool use_fast_accum,
-      Tensor& out) {
-  // FP8 per-tensor and per-row scaling expect fp32 scales.
-  TORCH_CHECK_VALUE(scale_a.scalar_type() == kFloat && scale_b.scalar_type() == kFloat,
-      "For grouped FP8 rowwise, both scales must be float32 tensors");
-#ifndef USE_ROCM
-  return _f8_f8_bf16_rowwise_grouped_mm_cuda(
-      mat_a,
-      mat_b,
-      scale_a,
-      scale_b,
-      offs,
-      bias,
-      use_fast_accum,
-      out);
-#else
-  // NOTE: ignore use_fast_accum
-  TORCH_CHECK_VALUE(!bias.has_value(), "ROCM grouped gemm does not support bias")
-  return _f8_f8_bf16_rowwise_grouped_mm_rocm(
-      mat_a,
-      mat_b,
-      scale_a,
-      scale_b,
-      offs,
-      out);
-#endif
-}
-
-Tensor
-_scaled_grouped_mm_cuda(
-        const Tensor& mat_a,
-        const Tensor& mat_b,
-        const Tensor& scale_a,
-        const Tensor& scale_b,
-        const std::optional<at::Tensor>& offs,
-        const std::optional<at::Tensor>& bias,
-        const std::optional<at::Tensor>& scale_result,
-        std::optional<c10::ScalarType> out_dtype,
-        bool use_fast_accum) {
-  bool allowed_device = _scaled_mm_allowed_device(/*sm90_only*/true, /*sm100_only*/true);
-  TORCH_CHECK_VALUE(allowed_device, "torch._scaled_grouped_mm is only supported on CUDA devices with compute capability = [9.0, 10.0], or ROCm MI300+");
-
-  TORCH_CHECK_VALUE(!check_valid_strides_and_return_transposed(mat_a), "Expected mat1 to not be transposed");
-  TORCH_CHECK_VALUE(check_valid_strides_and_return_transposed(mat_b), "Expected mat2 to be transposed");
-  TORCH_CHECK_VALUE(mat_a.dim() == 2 || mat_a.dim() == 3, "mat_a has to be 2 or 3d");
-  TORCH_CHECK_VALUE(mat_b.dim() == 2 || mat_b.dim() == 3, "mat_b has to be 2 or 3d");
-  const bool a_is_2d = mat_a.dim() == 2;
-  const bool b_is_2d = mat_b.dim() == 2;
-
-  // NOTE(slayton): For sub-1B formats want contraction_dim argument?
-  if (!a_is_2d || !b_is_2d) {
-    TORCH_CHECK_VALUE(mat_a.size(-1) == mat_b.size(-2), "contraction dimension of mat_a and mat_b must match");
-  }
-  TORCH_CHECK_VALUE(
-    mat_a.size(-1) % 16 == 0,
-    "Expected trailing dimension of mat_a to be divisible by 16 ",
-    "but got mat1 shape: (",
-    mat_a.sizes(),
-    ").");
-  TORCH_CHECK_VALUE(mat_b.size(-2) % 16 == 0 && mat_b.size(-1) % 16 == 0,
-    "Expected mat_b shape to be divisible by 16 ",
-    "but got mat_b shape: (",
-    mat_b.sizes(),
-    ").");
-
-
-  TORCH_CHECK_VALUE(!bias.has_value(), "Bias not supported yet");
-  TORCH_CHECK_VALUE(!scale_result.has_value(), "Scale result not supported yet");
-  TORCH_CHECK_VALUE(offs.has_value() ==  (a_is_2d || b_is_2d), "Have to provide offsets if there is a 2d matrix");
-
-  // NOTE: mxfp8 x mxfp8 requires (and asserts later) that offsets is present.
-  //       for rowwise, no offsets implies 3d-3d and is handled by lower-level
-  //       routines
-  if (offs.has_value()) {
-    TORCH_CHECK_VALUE(offs->dim() == 1, "offs has to be 1D");
-    TORCH_CHECK_VALUE(offs->dtype() == at::kInt, "Offsets have to be int32");
-  }
-  // FP8 per-tensor and per-row scaling expect fp32 scales.
-  // MXFP8 expects float8_e8m0fnu scales.
-  TORCH_CHECK_VALUE(
-      (scale_a.scalar_type() == kFloat && scale_b.scalar_type() == kFloat) ||
-      (scale_a.scalar_type() == at::kFloat8_e8m0fnu && scale_b.scalar_type() == at::kFloat8_e8m0fnu),
-      "For FP8 tensorwise and rowwise, both scales must both be float32 tensors. For MXFP8, scales must both be float8_e8m0fnu tensors.");
-
-  const int scale_multiplier = (mat_a.dim() == 2 && mat_b.dim() == 2) ? offs->size(0) : 1;
-  check_scale(mat_a, scale_a, 0 ,0, scale_multiplier);
-  check_scale(mat_b, scale_b, 1, 1, scale_multiplier);
-
-  const auto out_dtype_ = out_dtype.value_or(kBFloat16);
-  TORCH_CHECK_VALUE(out_dtype_ == kBFloat16, "Only bf16 high precision output types are supported for grouped gemm");
-
-  Tensor out = create_grouped_gemm_output_tensor(mat_a, mat_b, offs, out_dtype_);
-
-#if defined(USE_FBGEMM_GENAI) && defined(USE_CUDA) && !defined(USE_ROCM)
-  // MXFP8 grouped GEMM dispatching
-  bool is_mx8mx8bf16 = (
-    mat_a.scalar_type() == at::kFloat8_e4m3fn && mat_b.scalar_type() == at::kFloat8_e4m3fn &&
-    scale_a.scalar_type() == at::kFloat8_e8m0fnu && scale_b.scalar_type() == at::kFloat8_e8m0fnu
-  );
-#else
-  bool is_mx8mx8bf16 = false;
-#endif
-
-  if (is_mx8mx8bf16) {
-    // Note: Passing implied SwizzleType here, correctness of scale previously checked
-    //       in `check_scale` call
-    return _mx8_mx8_bf16_grouped_mm_fbgemm(
-        mat_a,
-        mat_b,
-        scale_a,
-        SwizzleType::SWIZZLE_32_4_4,
-        scale_b,
-        SwizzleType::SWIZZLE_32_4_4,
-        offs.value(),
-        out);
-  }
-
-  // If we're not MXFP8, then we're row-wise scaling.
-  return _f8_f8_bf16_rowwise_grouped_mm(
-      mat_a,
-      mat_b,
-      scale_a,
-      scale_b,
-      offs,
-      bias,
-      use_fast_accum,
-      out);
-}
-
-namespace {
-
-std::array<std::tuple<std::string, acceptance_fn, ScaledGemmImplementation>, 2> scale_grouped_kernel_dispatch = {{
-  { "rowwise_rowwise", check_rowwise_recipe, ScaledGemmImplementation::ROWWISE_ROWWISE},
-  { "mxfp8_mxfp8", check_mxfp8_recipe, ScaledGemmImplementation::MXFP8_MXFP8}}};
-
-} // anonymous namespace
-
-Tensor
-_scaled_grouped_mm_cuda_v2(
-          const Tensor& mat_a, const Tensor& mat_b,
-          ArrayRef<Tensor> scale_a,
-          IntArrayRef scale_recipe_a,
-          IntArrayRef swizzle_a,
-          ArrayRef<Tensor> scale_b,
-          IntArrayRef scale_recipe_b,
-          IntArrayRef swizzle_b,
-          const std::optional<Tensor>& offs,
-          const std::optional<Tensor>& bias,
-          const std::optional<c10::ScalarType> out_dtype,
-          IntArrayRef contraction_dim,
-          bool use_fast_accum) {
-  bool allowed_device = _scaled_mm_allowed_device(/*sm90_only*/true, /*sm100_only*/true);
-  TORCH_CHECK_VALUE(allowed_device, "torch._scaled_grouped_mm is only supported on CUDA devices with compute capability = [9.0, 10.0], or ROCm MI300+");
-
-  TORCH_CHECK_VALUE(!check_valid_strides_and_return_transposed(mat_a), "Expected mat1 to not be transposed");
-  TORCH_CHECK_VALUE(check_valid_strides_and_return_transposed(mat_b), "Expected mat2 to be transposed");
-  TORCH_CHECK_VALUE(mat_a.dim() == 2 || mat_a.dim() == 3, "mat_a has to be 2 or 3d");
-  TORCH_CHECK_VALUE(mat_b.dim() == 2 || mat_b.dim() == 3, "mat_b has to be 2 or 3d");
-  const bool a_is_2d = mat_a.dim() == 2;
-  const bool b_is_2d = mat_b.dim() == 2;
-
-  // NOTE(slayton): For sub-1B formats want contraction_dim argument?
-  if (!a_is_2d || !b_is_2d) {
-    if (contraction_dim.size() > 0) {
-      const int dim_a = contraction_dim[0], dim_b = mat_b.size(contraction_dim[1]);
-      TORCH_CHECK_VALUE(mat_a.size(dim_a) == mat_b.size(dim_b),
-          "Contraction dimensions (", dim_a, ",", dim_b, ") of mat_a and mat_b must match, got: ", mat_a.size(dim_a), " and ",
-          mat_b.size(dim_b));
-      // Note: only (-1, -2) is currently supported
-      TORCH_CHECK_VALUE(dim_a == -1 && dim_b == -2, "Curently contraction dims must be (-1, -2) only");
-    } else {
-      TORCH_CHECK_VALUE(mat_a.size(-1) == mat_b.size(-2), "contraction dimension of mat_a and mat_b must match");
-    }
-  }
-  TORCH_CHECK_VALUE(
-    mat_a.size(-1) % 16 == 0,
-    "Expected trailing dimension of mat_a to be divisible by 16 ",
-    "but got mat1 shape: (",
-    mat_a.sizes(),
-    ").");
-  TORCH_CHECK_VALUE(mat_b.size(-2) % 16 == 0 && mat_b.size(-1) % 16 == 0,
-    "Expected mat_b shape to be divisible by 16 ",
-    "but got mat_b shape: (",
-    mat_b.sizes(),
-    ").");
-
-  TORCH_CHECK_VALUE(!bias.has_value(), "Bias not supported yet");
-  TORCH_CHECK_VALUE(offs.has_value() ==  (a_is_2d || b_is_2d), "Have to provide offsets if there is a 2d matrix");
-
-  // NOTE: mxfp8 x mxfp8 requires (and asserts later) that offsets is present.
-  //       for rowwise, no offsets implies 3d-3d and is handled by lower-level
-  //       routines
-  if (offs.has_value()) {
-    TORCH_CHECK_VALUE(offs->dim() == 1, "offs has to be 1D");
-    TORCH_CHECK_VALUE(offs->dtype() == at::kInt, "Offsets have to be int32");
-  }
-
-  const auto out_dtype_ = out_dtype.value_or(kBFloat16);
-  TORCH_CHECK_VALUE(out_dtype_ == kBFloat16, "Only bf16 high precision output types are supported for grouped gemm");
-
-  Tensor out = create_grouped_gemm_output_tensor(mat_a, mat_b, offs, out_dtype_);
-
-  // Conversion of implicitly-defined enums to explicit
-  auto scale_recipe_a_enum = convert_int_to_enum<ScalingType>(scale_recipe_a);
-  auto swizzle_a_enum = convert_int_to_enum<SwizzleType>(swizzle_a);
-  auto scale_recipe_b_enum = convert_int_to_enum<ScalingType>(scale_recipe_b);
-  auto swizzle_b_enum = convert_int_to_enum<SwizzleType>(swizzle_b);
-
-  // at this point we can start working out what we want to be doing
-  // Try to do as few steps as possible.
-  // NOTE: support is deliberately sparse, can explicitly enumerate all combinations allowed.
-  // Do this via a list of defined (name, acceptance, concrete_impl) tuples.
-  ScaledGemmImplementation gemm_impl = ScaledGemmImplementation::NONE;
-  for (const auto& fn_entry : scale_grouped_kernel_dispatch) {
-    const auto [name, accept_fn, scaled_gemm_impl] = fn_entry;
-    bool ok = accept_fn(mat_a.scalar_type(),
-                        scale_recipe_a_enum,
-                        scale_a,
-                        mat_b.scalar_type(),
-                        scale_recipe_b_enum,
-                        scale_b);
-    if (ok) {
-      gemm_impl = scaled_gemm_impl;
-      break;
-    }
-  }
-  TORCH_CHECK_VALUE(gemm_impl != ScaledGemmImplementation::NONE,
-      "No gemm implementation was found");
-
-  switch (gemm_impl) {
-    case ScaledGemmImplementation::ROWWISE_ROWWISE: {
-      const int scale_multiplier = (mat_a.dim() == 2 && mat_b.dim() == 2) ? offs->size(0) : 1;
-      _check_scales_fp8_rowwise(mat_a, scale_a[0], 0 /* dim */ , 0 /* arg_idx */, scale_multiplier);
-      _check_scales_fp8_rowwise(mat_b, scale_b[0], 1 /* dim */ , 1 /* arg_idx */, scale_multiplier);
-      return _f8_f8_bf16_rowwise_grouped_mm(
-          mat_a,
-          mat_b,
-          scale_a[0],
-          scale_b[0],
-          offs,
-          bias,
-          use_fast_accum,
-          out);
-    }
-    case ScaledGemmImplementation::MXFP8_MXFP8: {
-      _check_scales_mxfp8(mat_a, scale_a[0], 0 /* dim */, 0 /* arg_idx */);
-      _check_scales_mxfp8(mat_b, scale_b[0], 1 /* dim */, 1 /* arg_idx */);
-      return _mx8_mx8_bf16_grouped_mm_fbgemm(
-          mat_a,
-          mat_b,
-          scale_a[0],
-          swizzle_a_enum[0],
-          scale_b[0],
-          swizzle_b_enum[0],
-          offs.value(),
-          out);
-    }
-    default:
-      TORCH_CHECK_NOT_IMPLEMENTED(false,
-          "_scaled_grouped_mm_cuda_v2 is in an inconsistent state - should never reach here");
-  }
-}
-
-Tensor _grouped_mm_cuda(const Tensor& mat_a, const Tensor& mat_b,
-const std::optional<at::Tensor>& offs,
-const std::optional<at::Tensor>& bias,
-std::optional<c10::ScalarType> out_dtype) {
-  _grouped_mm_validate_inputs(mat_a, mat_b, offs, bias, out_dtype);
-  bool a_b_and_out_are_bf16 = (
-    mat_a.dtype() == at::kBFloat16 &&
-    mat_b.dtype() == at::kBFloat16 &&
-    out_dtype.value_or(at::kBFloat16) == at::kBFloat16
-  );
-#ifndef USE_ROCM
-  bool use_fast_path = _scaled_mm_allowed_device(/*sm90_only*/true, /*sm100_only*/true) && a_b_and_out_are_bf16;
-#else
-  // _scaled_mm_allowed_device is used here within _grouped_mm_cuda which seems incorrect since scale is not used.
-  // the _grouped_mm_fallback should be safe for any ROCm GPU since it's just calling typical mm/bmm
-  bool use_fast_path = false;
-#endif
-  const auto out_dtype_ = _resolve_grouped_mm_out_dtype(mat_a, mat_b, out_dtype);
-  Tensor out = create_grouped_gemm_output_tensor(mat_a, mat_b, offs, out_dtype_);
-  if (use_fast_path) {
-    // fast path, no d2h sync needed
-    at::cuda::detail::bf16bf16_grouped_mm(mat_a, mat_b, offs, bias, out);
-  } else {
-    _grouped_mm_fallback(mat_a, mat_b, offs, bias, out_dtype, out);
-  }
-  return out;
-}
-
 static void baddbmm_bmm_out_dtype_checks(const Tensor& batch1, const Tensor& batch2, const Scalar& beta, const Scalar& alpha, const at::ScalarType out_dtype, bool is_bmm, const std::optional<Tensor>& self_baddbmm = std::nullopt) {
   // ref ATen/native/LinearAlgebra.cpp common_checks_baddbmm_bmm
   TORCH_CHECK(batch1.dim() == 3, "batch1 must be a 3D tensor");

aten/src/ATen/native/cuda/GroupedBlas.cpp

@@ -0,0 +1,574 @@
+#include <cstdint>
+#include <c10/util/typeid.h>
+#include <c10/util/Exception.h>
+#include <c10/util/SmallVector.h>
+#include <c10/core/Scalar.h>
+#include <c10/core/ScalarType.h>
+#include <c10/util/Exception.h>
+#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
+#include <ATen/core/Tensor.h>
+#include <ATen/core/NamedTensor.h>
+#include <ATen/Dispatch.h>
+#include <ATen/ExpandUtils.h>
+#include <ATen/OpMathType.h>
+#include <ATen/TensorUtils.h>
+#include <ATen/cuda/CUDABlas.h>
+#include <ATen/cuda/CUDAScaledBlas.h>
+#include <ATen/cuda/tunable/Tunable.h>
+#include <ATen/cuda/tunable/TunableGemm.h>
+#include <ATen/native/Resize.h>
+#include <c10/util/MaybeOwned.h>
+#include <ATen/native/GroupedMMUtils.h>
+#include <ATen/native/cuda/RowwiseScaledMM.h>
+#include <ATen/native/cuda/ScaledGroupMM.h>
+#include <ATen/native/cuda/GroupMM.h>
+#include <ATen/ceil_div.h>
+
+#ifdef USE_FBGEMM_GENAI
+#include <fbgemm_gpu/torch_ops.h>
+#endif
+
+#ifndef AT_PER_OPERATOR_HEADERS
+#include <ATen/Functions.h>
+#include <ATen/NativeFunctions.h>
+#else
+#include <ATen/ops/_addmm_activation_native.h>
+#include <ATen/ops/_efficientzerotensor.h>
+#include <ATen/ops/_scaled_mm_native.h>
+#include <ATen/ops/_unsafe_view_native.h>
+#include <ATen/ops/abs.h>
+#include <ATen/ops/addmm_native.h>
+#include <ATen/ops/addmv_native.h>
+#include <ATen/ops/baddbmm_native.h>
+#include <ATen/ops/bmm_native.h>
+#include <ATen/ops/copy_native.h>
+#include <ATen/ops/dot_native.h>
+#include <ATen/ops/empty.h>
+#include <ATen/ops/empty_strided.h>
+#include <ATen/ops/gelu.h>
+#include <ATen/ops/max.h>
+#include <ATen/ops/mm_native.h>
+#include <ATen/ops/mul.h>
+#include <ATen/ops/relu.h>
+#include <ATen/ops/ones.h>
+#include <ATen/ops/scalar_tensor_native.h>
+#include <ATen/ops/vdot_native.h>
+#endif
+
+using at::blas::ScalingType;
+using at::blas::SwizzleType;
+
+namespace scaled_blas = at::cuda::scaled;
+using scaled_blas::ScaledGemmImplementation;
+using scaled_blas::convert_int_to_enum;
+using scaled_blas::_scaled_mm_allowed_device;
+
+namespace at::native {
+
+namespace {
+
+// 2d-2d and 2d-3d
+// scaling=MXFP8
+// CUDA-only
+Tensor&
+_mx8_mx8_bf16_grouped_mm_fbgemm(
+        const Tensor& mat_a,
+        const Tensor& mat_b,
+        const Tensor& scale_a,
+        const SwizzleType& swizzle_a,
+        const Tensor& scale_b,
+        const SwizzleType& swizzle_b,
+        const std::optional<at::Tensor>& offs,
+        Tensor& out) {
+    const bool a_is_2d = mat_a.dim() == 2;
+    const bool b_is_2d = mat_b.dim() == 2;
+    bool b_is_3d = mat_b.dim() == 3;
+    bool is_2d_2d = a_is_2d && b_is_2d;
+    bool is_2d_3d = a_is_2d && b_is_3d;
+    TORCH_CHECK_VALUE(is_2d_2d || is_2d_3d, "MXFP8 grouped GEMM currently only supports 2d-2d and 2d-3d cases");
+    TORCH_CHECK_VALUE(offs.has_value(), "MXFP8 2d-2d and 2d-3d grouped GEMMs requires offsets");
+    TORCH_CHECK_VALUE(out.scalar_type() == at::kBFloat16, "Only bf16 out_dtype is supported for MXFP8 grouped gemm");
+    // MXFP8 expects float8_e8m0fnu scales.
+    TORCH_CHECK_VALUE(scale_a.scalar_type() == at::kFloat8_e8m0fnu && scale_b.scalar_type() == at::kFloat8_e8m0fnu,
+        "For MXFP8 grouped gemm, both scales must be float8_e8m0fnu tensors.");
+#ifdef USE_ROCM
+    TORCH_CHECK_VALUE(swizzle_a == SwizzleType::NO_SWIZZLE && swizzle_b == SwizzleType::NO_SWIZZLE,
+        "For ROCM MXFP8 grouped gemm, both scale swizzle types must be SWIZZLE_NONE");
+#else
+    TORCH_CHECK_VALUE(swizzle_a == SwizzleType::SWIZZLE_32_4_4 && swizzle_b == SwizzleType::SWIZZLE_32_4_4,
+        "For CUDA MXFP8 grouped gemm, both scale swizzle types must be SWIZZLE_32_4_4");
+#endif
+
+#if defined(USE_FBGEMM_GENAI) and !defined(USE_ROCM)
+    fbgemm_gpu::mx8mx8bf16_grouped_mm(
+        mat_a,
+        mat_b,
+        scale_a,
+        scale_b,
+        offs.value(),
+        out);
+#else
+    TORCH_CHECK_NOT_IMPLEMENTED(false, "mxfp8_mxfp8 grouped gemm requires compile with USE_FBGEMM_GENAI");
+#endif
+    return out;
+}
+
+// 2d-2d and 2d-3d cases
+// scaling=rowwise
+// CUDA-only
+Tensor&
+_f8_f8_bf16_rowwise_grouped_mm_cuda(
+          const Tensor& mat_a,
+          const Tensor& mat_b,
+          const Tensor& scale_a,
+          const Tensor& scale_b,
+          const std::optional<Tensor>& offs,
+          const std::optional<Tensor>& bias,
+          const bool use_fast_accum,
+          Tensor& out) {
+  TORCH_CHECK_VALUE(mat_a.dtype() == at::kFloat8_e4m3fn, "Expected mat_a to be Float8_e4m3 matrix got ", mat_a.scalar_type());
+  TORCH_CHECK_VALUE(mat_b.dtype() == at::kFloat8_e4m3fn, "Expected mat_a to be Float8_e4m3 matrix got ", mat_b.scalar_type());
+
+  at::cuda::detail::f8f8bf16_grouped_mm(
+      mat_a,
+      mat_b,
+      scale_a,
+      scale_b,
+      offs,
+      bias,
+      use_fast_accum,
+      out);
+    return out;
+}
+
+// 2d-2d and 2d-3d cases
+// scaling=rowwise
+// only being called for rocm
+Tensor&
+_f8_f8_bf16_rowwise_grouped_mm_rocm(
+      const Tensor& mat_a,
+      const Tensor& mat_b,
+      const Tensor& scale_a,
+      const Tensor& scale_b,
+      const std::optional<Tensor>& offs,
+      Tensor& out) {
+  TORCH_CHECK_VALUE(mat_a.dtype() == at::kFloat8_e4m3fnuz, "Expected mat_a to be Float8_e4m3fnuz matrix got ", mat_a.scalar_type());
+  TORCH_CHECK_VALUE(mat_b.dtype() == at::kFloat8_e4m3fnuz, "Expected mat_a to be Float8_e4m3fnuz matrix got ", mat_b.scalar_type());
+
+#if defined(USE_FBGEMM_GENAI) && defined(USE_ROCM)
+  fbgemm_gpu::f8f8bf16_rowwise_grouped_mm(
+      mat_a,
+      // FBGEMM expects B matrix shape to be (.., N, K)
+      mat_b.transpose(-2, -1),
+      scale_a,
+      scale_b,
+      offs,
+      out);
+#else
+  TORCH_CHECK_NOT_IMPLEMENTED(false, "grouped gemm is not supported without USE_FBGEMM_GENAI on ROCM")
+#endif
+  return out;
+
+}
+
+// Dispatch f8 x f8 -> bf16 row-wise scaled to rocm/cuda
+Tensor&
+_f8_f8_bf16_rowwise_grouped_mm(
+      const Tensor& mat_a,
+      const Tensor& mat_b,
+      const Tensor& scale_a,
+      const Tensor& scale_b,
+      const std::optional<Tensor>& offs,
+      const std::optional<Tensor>& bias,
+      bool use_fast_accum,
+      Tensor& out) {
+  // FP8 per-tensor and per-row scaling expect fp32 scales.
+  TORCH_CHECK_VALUE(scale_a.scalar_type() == kFloat && scale_b.scalar_type() == kFloat,
+      "For grouped FP8 rowwise, both scales must be float32 tensors");
+#ifndef USE_ROCM
+  return _f8_f8_bf16_rowwise_grouped_mm_cuda(
+      mat_a,
+      mat_b,
+      scale_a,
+      scale_b,
+      offs,
+      bias,
+      use_fast_accum,
+      out);
+#else
+  // NOTE: ignore use_fast_accum
+  TORCH_CHECK_VALUE(!bias.has_value(), "ROCM grouped gemm does not support bias")
+  return _f8_f8_bf16_rowwise_grouped_mm_rocm(
+      mat_a,
+      mat_b,
+      scale_a,
+      scale_b,
+      offs,
+      out);
+#endif
+}
+
+void _check_scales_fp8_rowwise(const Tensor& mat, const Tensor& scale, const int dim, const int arg_idx, const int scale_multiplier=1) {
+  // Checks scales for 2d or 3d target tensors (`mat`).
+  if (mat.dim() == 2) {
+    TORCH_CHECK(
+        scale.dim() == 1,
+        "scale must be a 1D tensor, but got ",
+        scale.dim(),
+        "D, arg ",
+        arg_idx);
+    TORCH_CHECK(
+        scale.is_contiguous(), "scale must be contiguous for arg ", arg_idx);
+    TORCH_CHECK(
+        scale.size(0) == mat.size(dim) * scale_multiplier,
+        "scale must have the same length as mat for arg ",
+        arg_idx);
+  } else {
+    TORCH_CHECK(
+        scale.dim() == 2,
+        "scale must be a 2D tensor, but got ",
+        scale.dim(),
+        "D for arg ",
+        arg_idx);
+    TORCH_CHECK(
+        scale.stride(1) == 1,
+        "scale must be contiguous in the last dimension for arg ",
+        arg_idx);
+    TORCH_CHECK(
+        scale.size(0) == mat.size(0),
+        "scale must have the same batch dimension as mat for arg ",
+        arg_idx);
+    TORCH_CHECK(
+        scale.size(1) == mat.size(1 + dim),
+        "scale must have the same first dimension as mat for arg ",
+        arg_idx);
+  }
+}
+
+void _check_scales_mxfp8(const Tensor& mat, const Tensor& scale, const int dim, const int arg_idx) {
+  // Checks scales for 2d or 3d target tensors (`mat`).
+  if (mat.dim() == 2) {
+    // For MXFP8, 2d tensors have variable size groups represented as subtensors,
+    // that are converted to blocked padded format individually,
+    // so we can't check the scale sizes without doing a d2h sync to get the group sizes here.
+    TORCH_CHECK(
+      scale.dim() == mat.dim(),
+      "for mxfp8, scale must have same number of dimensions as parent tensor, but got mat.dim() = ", mat.dim(), " and scale.dim() = ", scale.dim(), " for arg ", arg_idx);
+
+    // LHS mat shape (M, total_K) -> scale shape (rounded_up(M, 128), rounded_up_per_group(K/32, 4))
+    // RHS mat shape (total_K, N) -> scale shape (rounded_up(N, 128), rounded_up_per_group(K/32, 4))
+    //   * weight is transposed prior to the call, scale stays non-transposed.
+    bool LHS = arg_idx == 0;
+    int scale_dim_to_check = 0;
+    int mat_dim_to_check = LHS ? 0 : 1;
+    TORCH_CHECK(
+        scale.size(scale_dim_to_check) >= mat.size(mat_dim_to_check),
+        "for mxfp8, arg ", arg_idx, " tensor shape (", mat.size(0), ", ", mat.size(1), ") ",
+        "must have scale.shape[", scale_dim_to_check, "] >= ", mat.size(mat_dim_to_check), " but got scale.shape=(", scale.size(0), ", ", scale.size(1), ")");
+  } else {
+    // For MXFP8, 3d tensors have static group sizes (stack of 2d tensors),
+    // so we can check the exact expected scale sizes here without a d2h sync.
+    auto round_up = [](auto x, auto y) {
+        return ((x + y - 1) / y) * y;
+    };
+
+    // TODO: this is for 3d tensor in 2d-3d case specifically.
+    // We'll need to support 3d-3d and 3d-2d cases once mxfp8 grouped gemm supports them.
+    int64_t G = mat.size(0);
+    int64_t K = mat.size(1);
+    int64_t N = mat.size(2);
+    int64_t blocked_scale_K = round_up(K/32, 4);
+    int64_t blocked_scale_N = round_up(N, 128);
+
+    // fbgemm expects stack of flattened blocked scales for 3d tensor, shape (G, blocked_scale_K * blocked_scale_N).
+    TORCH_CHECK(
+      scale.dim() == mat.dim() - 1,
+      "for mxfp8 2d-3d grouped GEMM, the 3d tensor of shape (G,K,N) must have a 2d scale of shape (G, blocked_scale_K * blocked_scale_N), but scale is ", scale.dim(), "D for arg ", arg_idx
+    );
+    TORCH_CHECK(
+      scale.size(0) == G && scale.size(1) == blocked_scale_K * blocked_scale_N,
+      "for mxfp8, the tensor shape (", G, ", ", K, ", ", N, ") must have scale shape (", G, ",", blocked_scale_K, ",", blocked_scale_N, ") for arg ", arg_idx
+    );
+  }
+}
+
+void check_scale(const Tensor& mat, const Tensor& scale, const int dim, const int arg_idx, const int scale_multiplier=1) {
+  bool using_fp8_rowwise = scale.scalar_type() == kFloat;
+  bool using_mxfp8 = scale.scalar_type() == at::kFloat8_e8m0fnu;
+  if (using_fp8_rowwise) {
+    _check_scales_fp8_rowwise(mat, scale, dim, arg_idx, scale_multiplier);
+  } else if (using_mxfp8) {
+    _check_scales_mxfp8(mat, scale, dim, arg_idx);
+  } else {
+    TORCH_CHECK(false, "scale must be float32 or float8_e8m0fnu, but got ", scale.dtype());
+  }
+}
+
+} // namespace
+
+Tensor
+_scaled_grouped_mm_cuda(
+        const Tensor& mat_a,
+        const Tensor& mat_b,
+        const Tensor& scale_a,
+        const Tensor& scale_b,
+        const std::optional<at::Tensor>& offs,
+        const std::optional<at::Tensor>& bias,
+        const std::optional<at::Tensor>& scale_result,
+        std::optional<c10::ScalarType> out_dtype,
+        bool use_fast_accum) {
+  bool allowed_device = _scaled_mm_allowed_device(/*sm90_only*/true, /*sm100_only*/true);
+  TORCH_CHECK_VALUE(allowed_device, "torch._scaled_grouped_mm is only supported on CUDA devices with compute capability = [9.0, 10.0], or ROCm MI300+");
+
+  TORCH_CHECK_VALUE(!check_valid_strides_and_return_transposed(mat_a), "Expected mat1 to not be transposed");
+  TORCH_CHECK_VALUE(check_valid_strides_and_return_transposed(mat_b), "Expected mat2 to be transposed");
+  TORCH_CHECK_VALUE(mat_a.dim() == 2 || mat_a.dim() == 3, "mat_a has to be 2 or 3d");
+  TORCH_CHECK_VALUE(mat_b.dim() == 2 || mat_b.dim() == 3, "mat_b has to be 2 or 3d");
+  const bool a_is_2d = mat_a.dim() == 2;
+  const bool b_is_2d = mat_b.dim() == 2;
+
+  // NOTE(slayton): For sub-1B formats want contraction_dim argument?
+  if (!a_is_2d || !b_is_2d) {
+    TORCH_CHECK_VALUE(mat_a.size(-1) == mat_b.size(-2), "contraction dimension of mat_a and mat_b must match");
+  }
+  TORCH_CHECK_VALUE(
+    mat_a.size(-1) % 16 == 0,
+    "Expected trailing dimension of mat_a to be divisible by 16 ",
+    "but got mat1 shape: (",
+    mat_a.sizes(),
+    ").");
+  TORCH_CHECK_VALUE(mat_b.size(-2) % 16 == 0 && mat_b.size(-1) % 16 == 0,
+    "Expected mat_b shape to be divisible by 16 ",
+    "but got mat_b shape: (",
+    mat_b.sizes(),
+    ").");
+
+
+  TORCH_CHECK_VALUE(!bias.has_value(), "Bias not supported yet");
+  TORCH_CHECK_VALUE(!scale_result.has_value(), "Scale result not supported yet");
+  TORCH_CHECK_VALUE(offs.has_value() ==  (a_is_2d || b_is_2d), "Have to provide offsets if there is a 2d matrix");
+
+  // NOTE: mxfp8 x mxfp8 requires (and asserts later) that offsets is present.
+  //       for rowwise, no offsets implies 3d-3d and is handled by lower-level
+  //       routines
+  if (offs.has_value()) {
+    TORCH_CHECK_VALUE(offs->dim() == 1, "offs has to be 1D");
+    TORCH_CHECK_VALUE(offs->dtype() == at::kInt, "Offsets have to be int32");
+  }
+  // FP8 per-tensor and per-row scaling expect fp32 scales.
+  // MXFP8 expects float8_e8m0fnu scales.
+  TORCH_CHECK_VALUE(
+      (scale_a.scalar_type() == kFloat && scale_b.scalar_type() == kFloat) ||
+      (scale_a.scalar_type() == at::kFloat8_e8m0fnu && scale_b.scalar_type() == at::kFloat8_e8m0fnu),
+      "For FP8 tensorwise and rowwise, both scales must both be float32 tensors. For MXFP8, scales must both be float8_e8m0fnu tensors.");
+
+  const int scale_multiplier = (mat_a.dim() == 2 && mat_b.dim() == 2) ? offs->size(0) : 1;
+  check_scale(mat_a, scale_a, 0 ,0, scale_multiplier);
+  check_scale(mat_b, scale_b, 1, 1, scale_multiplier);
+
+  const auto out_dtype_ = out_dtype.value_or(kBFloat16);
+  TORCH_CHECK_VALUE(out_dtype_ == kBFloat16, "Only bf16 high precision output types are supported for grouped gemm");
+
+  Tensor out = create_grouped_gemm_output_tensor(mat_a, mat_b, offs, out_dtype_);
+
+#if defined(USE_FBGEMM_GENAI) && defined(USE_CUDA) && !defined(USE_ROCM)
+  // MXFP8 grouped GEMM dispatching
+  bool is_mx8mx8bf16 = (
+    mat_a.scalar_type() == at::kFloat8_e4m3fn && mat_b.scalar_type() == at::kFloat8_e4m3fn &&
+    scale_a.scalar_type() == at::kFloat8_e8m0fnu && scale_b.scalar_type() == at::kFloat8_e8m0fnu
+  );
+#else
+  bool is_mx8mx8bf16 = false;
+#endif
+
+  if (is_mx8mx8bf16) {
+    // Note: Passing implied SwizzleType here, correctness of scale previously checked
+    //       in `check_scale` call
+    return _mx8_mx8_bf16_grouped_mm_fbgemm(
+        mat_a,
+        mat_b,
+        scale_a,
+        SwizzleType::SWIZZLE_32_4_4,
+        scale_b,
+        SwizzleType::SWIZZLE_32_4_4,
+        offs.value(),
+        out);
+  }
+
+  // If we're not MXFP8, then we're row-wise scaling.
+  return _f8_f8_bf16_rowwise_grouped_mm(
+      mat_a,
+      mat_b,
+      scale_a,
+      scale_b,
+      offs,
+      bias,
+      use_fast_accum,
+      out);
+}
+
+namespace {
+
+using acceptance_fn = std::function<bool(c10::ScalarType, std::vector<ScalingType>&, ArrayRef<Tensor>&, c10::ScalarType, std::vector<ScalingType>&, ArrayRef<Tensor>&)>;
+
+std::array<std::tuple<std::string, acceptance_fn, ScaledGemmImplementation>, 2> scale_grouped_kernel_dispatch = {{
+  { "rowwise_rowwise", scaled_blas::check_rowwise_recipe, ScaledGemmImplementation::ROWWISE_ROWWISE},
+  { "mxfp8_mxfp8", scaled_blas::check_mxfp8_recipe, ScaledGemmImplementation::MXFP8_MXFP8}}};
+
+} // anonymous namespace
+
+Tensor
+_scaled_grouped_mm_cuda_v2(
+          const Tensor& mat_a, const Tensor& mat_b,
+          ArrayRef<Tensor> scale_a,
+          IntArrayRef scale_recipe_a,
+          IntArrayRef swizzle_a,
+          ArrayRef<Tensor> scale_b,
+          IntArrayRef scale_recipe_b,
+          IntArrayRef swizzle_b,
+          const std::optional<Tensor>& offs,
+          const std::optional<Tensor>& bias,
+          const std::optional<c10::ScalarType> out_dtype,
+          IntArrayRef contraction_dim,
+          bool use_fast_accum) {
+  bool allowed_device = _scaled_mm_allowed_device(/*sm90_only*/true, /*sm100_only*/true);
+  TORCH_CHECK_VALUE(allowed_device, "torch._scaled_grouped_mm is only supported on CUDA devices with compute capability = [9.0, 10.0], or ROCm MI300+");
+
+  TORCH_CHECK_VALUE(!check_valid_strides_and_return_transposed(mat_a), "Expected mat1 to not be transposed");
+  TORCH_CHECK_VALUE(check_valid_strides_and_return_transposed(mat_b), "Expected mat2 to be transposed");
+  TORCH_CHECK_VALUE(mat_a.dim() == 2 || mat_a.dim() == 3, "mat_a has to be 2 or 3d");
+  TORCH_CHECK_VALUE(mat_b.dim() == 2 || mat_b.dim() == 3, "mat_b has to be 2 or 3d");
+  const bool a_is_2d = mat_a.dim() == 2;
+  const bool b_is_2d = mat_b.dim() == 2;
+
+  // NOTE(slayton): For sub-1B formats want contraction_dim argument?
+  if (!a_is_2d || !b_is_2d) {
+    if (contraction_dim.size() > 0) {
+      const int dim_a = contraction_dim[0], dim_b = mat_b.size(contraction_dim[1]);
+      TORCH_CHECK_VALUE(mat_a.size(dim_a) == mat_b.size(dim_b),
+          "Contraction dimensions (", dim_a, ",", dim_b, ") of mat_a and mat_b must match, got: ", mat_a.size(dim_a), " and ",
+          mat_b.size(dim_b));
+      // Note: only (-1, -2) is currently supported
+      TORCH_CHECK_VALUE(dim_a == -1 && dim_b == -2, "Curently contraction dims must be (-1, -2) only");
+    } else {
+      TORCH_CHECK_VALUE(mat_a.size(-1) == mat_b.size(-2), "contraction dimension of mat_a and mat_b must match");
+    }
+  }
+  TORCH_CHECK_VALUE(
+    mat_a.size(-1) % 16 == 0,
+    "Expected trailing dimension of mat_a to be divisible by 16 ",
+    "but got mat1 shape: (",
+    mat_a.sizes(),
+    ").");
+  TORCH_CHECK_VALUE(mat_b.size(-2) % 16 == 0 && mat_b.size(-1) % 16 == 0,
+    "Expected mat_b shape to be divisible by 16 ",
+    "but got mat_b shape: (",
+    mat_b.sizes(),
+    ").");
+
+  TORCH_CHECK_VALUE(!bias.has_value(), "Bias not supported yet");
+  TORCH_CHECK_VALUE(offs.has_value() ==  (a_is_2d || b_is_2d), "Have to provide offsets if there is a 2d matrix");
+
+  // NOTE: mxfp8 x mxfp8 requires (and asserts later) that offsets is present.
+  //       for rowwise, no offsets implies 3d-3d and is handled by lower-level
+  //       routines
+  if (offs.has_value()) {
+    TORCH_CHECK_VALUE(offs->dim() == 1, "offs has to be 1D");
+    TORCH_CHECK_VALUE(offs->dtype() == at::kInt, "Offsets have to be int32");
+  }
+
+  const auto out_dtype_ = out_dtype.value_or(kBFloat16);
+  TORCH_CHECK_VALUE(out_dtype_ == kBFloat16, "Only bf16 high precision output types are supported for grouped gemm");
+
+  Tensor out = create_grouped_gemm_output_tensor(mat_a, mat_b, offs, out_dtype_);
+
+  // Conversion of implicitly-defined enums to explicit
+  auto scale_recipe_a_enum = convert_int_to_enum<ScalingType>(scale_recipe_a);
+  auto swizzle_a_enum = convert_int_to_enum<SwizzleType>(swizzle_a);
+  auto scale_recipe_b_enum = convert_int_to_enum<ScalingType>(scale_recipe_b);
+  auto swizzle_b_enum = convert_int_to_enum<SwizzleType>(swizzle_b);
+
+  // at this point we can start working out what we want to be doing
+  // Try to do as few steps as possible.
+  // NOTE: support is deliberately sparse, can explicitly enumerate all combinations allowed.
+  // Do this via a list of defined (name, acceptance, concrete_impl) tuples.
+  ScaledGemmImplementation gemm_impl = ScaledGemmImplementation::NONE;
+  for (const auto& fn_entry : scale_grouped_kernel_dispatch) {
+    const auto [name, accept_fn, scaled_gemm_impl] = fn_entry;
+    bool ok = accept_fn(mat_a.scalar_type(),
+                        scale_recipe_a_enum,
+                        scale_a,
+                        mat_b.scalar_type(),
+                        scale_recipe_b_enum,
+                        scale_b);
+    if (ok) {
+      gemm_impl = scaled_gemm_impl;
+      break;
+    }
+  }
+  TORCH_CHECK_VALUE(gemm_impl != ScaledGemmImplementation::NONE,
+      "No gemm implementation was found");
+
+  switch (gemm_impl) {
+    case ScaledGemmImplementation::ROWWISE_ROWWISE: {
+      const int scale_multiplier = (mat_a.dim() == 2 && mat_b.dim() == 2) ? offs->size(0) : 1;
+      _check_scales_fp8_rowwise(mat_a, scale_a[0], 0 /* dim */ , 0 /* arg_idx */, scale_multiplier);
+      _check_scales_fp8_rowwise(mat_b, scale_b[0], 1 /* dim */ , 1 /* arg_idx */, scale_multiplier);
+      return _f8_f8_bf16_rowwise_grouped_mm(
+          mat_a,
+          mat_b,
+          scale_a[0],
+          scale_b[0],
+          offs,
+          bias,
+          use_fast_accum,
+          out);
+    }
+    case ScaledGemmImplementation::MXFP8_MXFP8: {
+      _check_scales_mxfp8(mat_a, scale_a[0], 0 /* dim */, 0 /* arg_idx */);
+      _check_scales_mxfp8(mat_b, scale_b[0], 1 /* dim */, 1 /* arg_idx */);
+      return _mx8_mx8_bf16_grouped_mm_fbgemm(
+          mat_a,
+          mat_b,
+          scale_a[0],
+          swizzle_a_enum[0],
+          scale_b[0],
+          swizzle_b_enum[0],
+          offs.value(),
+          out);
+    }
+    default:
+      TORCH_CHECK_NOT_IMPLEMENTED(false,
+          "_scaled_grouped_mm_cuda_v2 is in an inconsistent state - should never reach here");
+  }
+}
+
+Tensor _grouped_mm_cuda(const Tensor& mat_a, const Tensor& mat_b,
+const std::optional<at::Tensor>& offs,
+const std::optional<at::Tensor>& bias,
+std::optional<c10::ScalarType> out_dtype) {
+  _grouped_mm_validate_inputs(mat_a, mat_b, offs, bias, out_dtype);
+  bool a_b_and_out_are_bf16 = (
+    mat_a.dtype() == at::kBFloat16 &&
+    mat_b.dtype() == at::kBFloat16 &&
+    out_dtype.value_or(at::kBFloat16) == at::kBFloat16
+  );
+#ifndef USE_ROCM
+  bool use_fast_path = _scaled_mm_allowed_device(/*sm90_only*/true, /*sm100_only*/true) && a_b_and_out_are_bf16;
+#else
+  // _scaled_mm_allowed_device is used here within _grouped_mm_cuda which seems incorrect since scale is not used.
+  // the _grouped_mm_fallback should be safe for any ROCm GPU since it's just calling typical mm/bmm
+  bool use_fast_path = false;
+#endif
+  const auto out_dtype_ = _resolve_grouped_mm_out_dtype(mat_a, mat_b, out_dtype);
+  Tensor out = create_grouped_gemm_output_tensor(mat_a, mat_b, offs, out_dtype_);
+  if (use_fast_path) {
+    // fast path, no d2h sync needed
+    at::cuda::detail::bf16bf16_grouped_mm(mat_a, mat_b, offs, bias, out);
+  } else {
+    _grouped_mm_fallback(mat_a, mat_b, offs, bias, out_dtype, out);
+  }
+  return out;
+}
+
+} // namespace at::native

aten/src/ATen/native/mkldnn/xpu/Blas.cpp

@@ -2,6 +2,7 @@
 #include <ATen/WrapDimUtilsMulti.h>
 #include <ATen/native/Resize.h>
 #include <ATen/native/mkldnn/xpu/detail/oneDNN.h>
+#include <ATen/native/xpu/Blas.h>
 #include <torch/library.h>
 #ifndef AT_PER_OPERATOR_HEADERS
 
@@ -50,9 +51,13 @@ Tensor& addmm_out(
       mat1.dtype(),
       " != ",
       mat2.dtype())
+
   // complex case
-  TORCH_CHECK(
-      !mat1.is_complex(), "Complex datatype matmul is not supported in oneDNN");
+  if (self.is_complex()) {
+    at::native::addmm_complex_out_xpu(self, mat1, mat2, beta, alpha, result);
+
+    return result;
+  }
 
   std::vector<int64_t> result_shape = {mat1.size(0), mat2.size(1)};
   result.resize_(result_shape);
@@ -167,8 +172,11 @@ Tensor& mm_out(const Tensor& self, const Tensor& mat2, Tensor& result) {
     return result;
   }
 
-  TORCH_CHECK(
-      !self.is_complex(), "Complex datatype matmul is not supported in oneDNN");
+  if (self.is_complex()) {
+    at::native::mm_complex_out_xpu(self, mat2, result);
+
+    return result;
+  }
 
   onednn::matmul(result, self, mat2, Tensor(), true, onednn::Attr());
   return result;
@@ -208,9 +216,12 @@ Tensor& baddbmm_out(
       input.sizes());
 
   // complex case
-  TORCH_CHECK(
-      !batch1.is_complex(),
-      "Complex datatype matmul is not supported in oneDNN");
+  if (input.is_complex()) {
+    at::native::baddbmm_complex_out_xpu(
+        input, batch1, batch2, beta, alpha, result);
+
+    return result;
+  }
 
   // general case
   onednn::Attr attr;
@@ -257,8 +268,13 @@ Tensor& bmm_out(const Tensor& self, const Tensor& batch2, Tensor& result) {
     return result;
   }
 
-  TORCH_CHECK(
-      !self.is_complex(), "Complex datatype matmul is not supported in oneDNN");
+  // complex case
+  if (self.is_complex()) {
+    at::native::bmm_complex_out_xpu(self, batch2, result);
+
+    return result;
+  }
+
   onednn::matmul(result, self, batch2, at::Tensor(), true, onednn::Attr());
   return result;
 }

aten/src/ATen/native/mps/kernels/BinaryKernel.metal

@@ -222,6 +222,13 @@ struct nextafter_functor {
   }
 };
 
+struct hypot_functor {
+  template <typename T>
+  inline T operator()(const T a, const T b) {
+    return static_cast<T>(precise::sqrt(float(a) * a + float(b) * b));
+  }
+};
+
 // Complex binary functors
 struct polar_functor {
   template <typename U>
@@ -362,6 +369,7 @@ struct igammac_functor {
   REGISTER_OPMATH_BINARY_OP(NAME, half, half);   \
   REGISTER_OPMATH_BINARY_OP(NAME, bfloat, bfloat)
 
+REGISTER_FLOAT_BINARY_OP(hypot);
 REGISTER_FLOAT_BINARY_OP(copysign);
 REGISTER_INT2FLOAT_BINARY_OP(copysign);
 REGISTER_FLOAT_BINARY_OP(fmax);

aten/src/ATen/native/mps/operations/BinaryKernel.mm

@@ -202,6 +202,10 @@ static void igammac_mps_kernel(TensorIteratorBase& iter) {
   lib.exec_binary_kernel(iter, "igammac");
 }
 
+static void hypot_mps_kernel(TensorIteratorBase& iter) {
+  lib.exec_binary_kernel(iter, "hypot");
+}
+
 REGISTER_DISPATCH(fmax_stub, &fmax_mps_kernel)
 REGISTER_DISPATCH(fmin_stub, &fmin_mps_kernel)
 REGISTER_DISPATCH(copysign_stub, &copysign_mps_kernel)
@@ -229,4 +233,5 @@ REGISTER_DISPATCH(fmod_stub, &fmod_mps_kernel)
 REGISTER_DISPATCH(remainder_stub, &remainder_mps_kernel)
 REGISTER_DISPATCH(igamma_stub, &igamma_mps_kernel)
 REGISTER_DISPATCH(igammac_stub, &igammac_mps_kernel)
+REGISTER_DISPATCH(hypot_stub, &hypot_mps_kernel)
 } // namespace at::native

aten/src/ATen/native/mps/operations/BinaryOps.mm

@@ -16,7 +16,6 @@
 #include <ATen/ops/eq_native.h>
 #include <ATen/ops/ge_native.h>
 #include <ATen/ops/gt_native.h>
-#include <ATen/ops/hypot_native.h>
 #include <ATen/ops/le_native.h>
 #include <ATen/ops/logaddexp2_native.h>
 #include <ATen/ops/logaddexp_native.h>
@@ -278,22 +277,6 @@ TORCH_IMPL_FUNC(pow_Scalar_out_mps)(const Scalar& base, const Tensor& exp, const
   }
 }
 
-TORCH_IMPL_FUNC(hypot_out_mps)(const Tensor& self, const Tensor& other, const Tensor& output) {
-  mps::BinaryOpBlock hypot_op_block = ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) {
-    MPSGraph* mpsGraph = cachedGraph->graph();
-    MPSGraphTensor* twoTensor = [mpsGraph constantWithScalar:2.0 shape:@[ @1 ] dataType:primaryCastTensor.dataType];
-    MPSGraphTensor* sumTensor = [mpsGraph additionWithPrimaryTensor:[mpsGraph powerWithPrimaryTensor:primaryCastTensor
-                                                                                     secondaryTensor:twoTensor
-                                                                                                name:nil]
-                                                    secondaryTensor:[mpsGraph powerWithPrimaryTensor:secondaryCastTensor
-                                                                                     secondaryTensor:twoTensor
-                                                                                                name:nil]
-                                                               name:nil];
-    return [mpsGraph squareRootWithTensor:sumTensor name:nil];
-  };
-  mps::binaryOpTensor(self, other, output, "hypot_out_mps", hypot_op_block);
-}
-
 TORCH_IMPL_FUNC(logaddexp_out_mps)(const Tensor& self, const Tensor& other, const Tensor& output) {
   mps::BinaryOpBlock logaddexp_op_block = ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) {
     MPSGraph* mpsGraph = cachedGraph->graph();

aten/src/ATen/native/native_functions.yaml

@@ -10040,8 +10040,7 @@
   structured: True
   structured_inherits: TensorIteratorBase
   dispatch:
-    CPU, CUDA: hypot_out
-    MPS: hypot_out_mps
+    CPU, CUDA, MPS: hypot_out
   tags: pointwise
 
 - func: hypot(Tensor self, Tensor other) -> Tensor

aten/src/ATen/native/transformers/cuda/sdp_utils.cpp

@@ -76,14 +76,21 @@ bool priority_order_init_ = false;
 // TODO(eqy): more benchmarking to determine whether this should include sm86/89
 // Needs to be kept in-sync with test_fused_chocie in test_transformers.py
 bool check_prefer_cudnn_attention() {
-  static const bool prefer_cudnn = c10::utils::check_env("TORCH_CUDNN_SDPA_PREFERRED") != false;
+  static const bool prefer_cudnn = c10::utils::check_env("TORCH_CUDNN_SDPA_DEPRIORITIZED") != true;
   if (!prefer_cudnn) {
     return false;
   }
 #if (defined(CUDNN_VERSION) && (CUDNN_VERSION >= 90900))
-  auto dprops = at::cuda::getCurrentDeviceProperties();
-  auto major = dprops->major;
-  return (major == 9 || major == 10) && !dprops->minor;
+  try {
+    auto dprops = at::cuda::getCurrentDeviceProperties();
+    auto major = dprops->major;
+    return (major == 9 || major == 10) && !dprops->minor;
+  } catch (c10::Error const& e) {
+#ifdef DEBUG
+    TORCH_WARN("check_prefer_cudnn_attention() caught exception ", e.what());
+#endif
+    return false;
+  }
 #else
   return false;
 #endif

aten/src/ATen/test/type_ptr_test.cpp

@@ -37,6 +37,10 @@ TEST(SingletonOrSharedTypePtr, Comparison) {
 
   EXPECT_NE(empty, p);
   EXPECT_NE(p, p2);
+
+  EXPECT_EQ(empty, empty);
+  EXPECT_EQ(p, p);
+  EXPECT_EQ(p2, p2);
 }
 
 TEST(SingletonOrSharedTypePtr, SingletonComparison) {
@@ -47,6 +51,8 @@ TEST(SingletonOrSharedTypePtr, SingletonComparison) {
   c10::TypePtr type = c10::NoneType::get();
   EXPECT_NE(type, c10::StringType::get());
   EXPECT_NE(type, c10::DeviceObjType::get());
+  EXPECT_EQ(type, type);
+  EXPECT_EQ(type, c10::NoneType::get());
 }
 
 

aten/src/ATen/test/vec_test_all_types.cpp

@@ -526,6 +526,41 @@ namespace {
             [](const vec& v) { return v.expm1(); },
             createDefaultUnaryTestCase<vec>(TestSeed(), false, true));
     }
+    TYPED_TEST(Exponents, ExpU20) {
+        using vec = TypeParam;
+        using VT = ValueType<TypeParam>;
+        using UVT = UvalueType<TypeParam>;
+
+        // Explicit edge values
+        VT v_too_small = VT(-100.0); // much less than -87.3
+        VT exp_too_small = std::exp(v_too_small);
+        VT v_neg_edge = VT(-0x1.5d5e2ap+6f);   // just at the edge
+        VT exp_neg_edge = std::exp(v_neg_edge);
+        VT v_zero = VT(0.0);         // middle, normal case
+        VT exp_zero = std::exp(v_zero);
+        VT v_pos_edge = VT(0x1.5d5e2ap+6f);    // just at the edge
+        VT exp_pos_edge = std::exp(v_pos_edge);
+        VT v_too_large = VT(100.0);  // much more than 87.3
+        VT exp_too_large = std::exp(v_too_large);
+
+        auto test_case = TestingCase<vec>::getBuilder()
+            // Randoms in normal range, but the .addCustom() below guarantees we hit the special/fallback cases
+            .addDomain(CheckWithinDomains<UVT>{{{-100, 100}}, false, getDefaultTolerance<UVT>()})
+            .addCustom({ {v_too_small}, exp_too_small })
+            .addCustom({ {v_neg_edge}, exp_neg_edge })
+            .addCustom({ {v_zero}, exp_zero })
+            .addCustom({ {v_pos_edge}, exp_pos_edge })
+            .addCustom({ {v_too_large}, exp_too_large })
+            .setTrialCount(65536)
+            .setTestSeed(TestSeed());
+
+        test_unary<vec>(
+            NAME_INFO(exp_u20_edge_cases),
+            RESOLVE_OVERLOAD(std::exp),
+            [](const vec& v) { return v.exp_u20(); },
+            test_case
+        );
+    }
     TYPED_TEST(ErrorFunctions, Erf) {
         using vec = TypeParam;
         test_unary<vec>(

docs/source/pytorch-api.md

@@ -41,6 +41,7 @@ torch.distributed.fsdp.fully_shard <distributed.fsdp.fully_shard>
 torch.distributed.tensor.parallel <distributed.tensor.parallel>
 torch.distributed.optim <distributed.optim>
 torch.distributed.pipelining <distributed.pipelining>
+torch.distributed._symmetric_memory <symmetric_memory>
 torch.distributed.checkpoint <distributed.checkpoint>
 torch.distributions <distributions>
 torch.compiler <torch.compiler>

docs/source/symmetric_memory.md

@@ -0,0 +1,380 @@
+```{eval-rst}
+.. role:: hidden
+    :class: hidden-section
+```
+
+# PyTorch Symmetric Memory
+
+:::{note}
+`torch.distributed._symmetric_memory` is currently in alpha state and under
+development. API changes may be possible.
+:::
+
+## Why Symmetric Memory?
+
+With rapidly evolving parallelization techniques, existing frameworks and
+libraries often struggle to keep up, and developers increasingly rely on custom
+implementations directly scheduling communications and computations. In recent
+years we’ve witnessed a shift from primarily relying on one-dimensional
+data-parallelism techniques to multi-dimensional parallelism ones. The latter
+have different latency requirements for different types of communications and
+thus require fine-grained overlapping of compute and communications.
+
+To minimize compute interference, they also require the use of copy engines and
+network interface cards (NICs) to drive communication. Network transport
+protocols such as remote direct memory access (RDMA) enhance the performance by
+enabling direct, high-speed, and low-latency communication between processors
+and memory. This increase in variety indicates the need for finer-grained
+communication primitives than are offered today by high-level collective APIs,
+ones that would enable developers to implement specific algorithms tailored for
+their use cases, such as low-latency collectives, fine-grained
+compute-communications overlap, or custom fusions.
+
+Furthermore, today’s advanced AI systems connect GPUs with high-bandwidth links
+(such as NVLinks, InfiniBand or RoCE), making GPU global memory directly
+accessible to peers. Such connections present a great opportunity for
+programmers to program the system as a single, gigantic GPU with vast accessible
+memory, instead of programming singular “GPU islands.”
+
+In this document, we will show how you can use PyTorch Symmetric Memory to
+program modern GPU systems as a “single GPU” and achieve fine-grained remote
+access.
+
+## What PyTorch Symmetric Memory unlocks?
+
+PyTorch Symmetric Memory unlocks three new capabilities:
+
+- **Customized communication patterns**: Increased flexibility in kernel writing
+allows developers to write custom kernels that implement their custom
+computations and communications, directly tailored to the need of the
+application. It will also be straightforward to add support for new data types
+along with the special compute that those data types might require, even if it’s
+not present yet in the standard libraries.
+
+- **In-kernel compute-comm fusion**: Device-initiated communication capability
+allows developers to write kernels with both computation and communication
+instructions, allowing for the fusion of computation and data movement in the
+smallest possible granularity.
+
+- **Low-latency remote access**: Network transport protocols like RDMA enhance the
+performance of symmetric memory in networked environments by enabling direct,
+high-speed, and low-latency communication between processors and memory. RDMA
+eliminates the overhead associated with the traditional network stack and CPU
+involvement. It also offloads data transfer from the compute to the NICs,
+freeing up compute resources for computational tasks.
+
+Next, we will show you how PyTorch Symmetric Memory (SymmMem) enables new
+applications with the above capabilities.
+
+## A “Hello World” example
+
+The PyTorch SymmMem programming model involves two key elements:
+
+- creating symmetric tensors
+- creating SymmMem kernels
+
+To create symmetric tensors, one can use the
+`torch.distributed._symmetric_memory` package:
+
+```python
+import torch.distributed._symmetric_memory as symm_mem
+
+t = symm_mem.empty(128, device=torch.device("cuda", rank))
+hdl = symm_mem.rendezvous(t, group)
+```
+
+The `symm_mem.empty` function creates a tensor that is backed by a symmetric
+memory allocation. The `rendezvous` function establishes a rendezvous with peers
+in the group, and returns a handle to the symmetric memory allocation. The
+handle provides method to access information related to the symmetric memory
+allocation, such as pointers to symmetric buffer on peer ranks, multicast
+pointer (if supported), and signal pads.
+
+The `empty` and `rendezvous` functions must be called in the same order on all
+ranks in the group.
+
+Then, collectives can be called on these tensors. For example, to perform a
+one-shot all-reduce:
+
+```python
+# Most SymmMem ops are under the torch.ops.symm_mem namespace
+torch.ops.symm_mem.one_shot_all_reduce(t, "sum", group)
+```
+
+Please note that `torch.ops.symm_mem` is an "op namespace" instead of a python
+module. Therefore, you can't import it by `import torch.ops.symm_mem`, neither
+can you import an op by `from torch.ops.symm_mem import one_shot_all_reduce`.
+You can call the op directly as in the example above.
+
+## Write your own kernel
+
+To write your own kernel doing communications with symmetric memory, you’ll need
+access to the addresses of mapped peer buffers and access to signal pads that
+are required for synchronization. In the kernel you’ll also need to perform
+correct synchronizations to make sure that peers are ready for communication,
+and signal to them that this GPU is ready.
+
+PyTorch Symmetric Memory provides CUDA Graph-compatible synchronization
+primitives that operate on the signal pad accompanying each symmetric memory
+allocation. Kernels using symmetric memory can be written both in CUDA and in
+Triton. Here’s an example allocating symmetric tensor and exchanging handles:
+
+```python
+import torch.distributed._symmetric_memory as symm_mem
+
+dist.init_process_group()
+rank = dist.get_rank()
+
+# Allocate a tensor
+t = symm_mem.empty(4096, device=f"cuda:{rank}")
+# Establish symmetric memory and obtain the handle
+hdl = symm_mem.rendezvous(t, dist.group.WORLD)
+```
+
+Access to buffer pointers, multimem pointer, and signal pads is provided via:
+
+```python
+hdl.buffer_ptrs
+hdl.multicast_ptr
+hdl.signal_pad_ptrs
+```
+
+Data pointed to by `buffer_ptrs` can be accessed just like regular local data,
+and any necessary compute can also be performed in the usual ways. As with local
+data, you can and should use vectorized accesses to improve efficiency.
+
+Symmetric memory is especially convenient for writing kernels in Triton. While
+previously Triton removed the barriers to writing efficient CUDA code, now
+communications can be added easily to Triton kernels. The kernel below
+demonstrates a low-latency, all-reduce kernel written in Triton.
+
+```python
+@triton.jit
+def one_shot_all_reduce_kernel(
+    buf_tuple,
+    signal_pad_ptrs,
+    output_ptr,
+    numel: tl.constexpr,
+    rank: tl.constexpr,
+    world_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    ptx_utils.symm_mem_sync(
+        signal_pad_ptrs, None, rank, world_size, hasSubsequenceMemAccess=True
+    )
+
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+
+    while block_start < numel:
+        offsets = block_start + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < numel
+        acc = tl.zeros((BLOCK_SIZE,), dtype=tl.bfloat16)
+
+        for i in tl.static_range(world_size):
+            buffer_rank = buf_tuple[i]
+            x = tl.load(buffer_rank + offsets, mask=mask)
+            acc += x
+
+        tl.store(output_ptr + offsets, acc, mask=mask)
+        block_start += tl.num_programs(axis=0) * BLOCK_SIZE
+
+    ptx_utils.symm_mem_sync(
+        signal_pad_ptrs, None, rank, world_size, hasPreviousMemAccess=True
+    )
+```
+
+Synchronizations at the beginning and the end of the kernel above guarantee that
+all the processes see consistent data. The bulk of the kernel is recognizable
+Triton code, and Triton will optimize it behind the scene, making sure memory
+accesses are performed in an efficient way with vectorization and unrolling. As
+with all Triton kernels, it is easily modifiable to add extra computations or
+change the communication algorithm. Visit
+https://github.com/meta-pytorch/kraken/blob/main/kraken to see additional
+utilities and examples of using symmetric memory to implement common patterns in
+Triton.
+
+## Scale out
+
+Large language models distribute experts onto more than 8 GPUs, hence requiring
+multi-node access capability. NICs capable of RDMA come to help. In addition,
+software libraries such as NVSHMEM or rocSHMEM abstract away the programming
+difference between intra-node access and inter-node access with primitives that
+are slightly higher level than pointer access, such as put and get.
+
+PyTorch provides NVSHMEM plugins to augment Triton kernels’ cross-node
+capabilities. As shown in the code snippet below, one can initiate a cross-node
+put command within the kernel.
+
+```python
+import torch.distributed._symmetric_memory._nvshmem_triton as nvshmem
+from torch.distributed._symmetric_memory._nvshmem_triton import requires_nvshmem
+
+@requires_nvshmem
+@triton.jit
+def my_put_kernel(
+    dest,
+    src,
+    nelems,
+    pe,
+):
+    nvshmem.put(dest, src, nelems, pe)
+```
+
+The `requires_nvshmem` decorator is used to indicate that the kernel requires
+the NVSHMEM device library as an external dependency. When Triton compiles the
+kernel, the decorator will search your system paths for the NVSHMEM device
+library. If it is available, Triton will include the necessary device assembly
+to use the NVSHMEM functions.
+
+## API Reference
+
+```{eval-rst}
+.. currentmodule:: torch.distributed._symmetric_memory
+```
+
+```{eval-rst}
+.. autofunction:: empty
+```
+
+```{eval-rst}
+.. autofunction:: rendezvous
+```
+
+```{eval-rst}
+.. autofunction:: is_nvshmem_available
+```
+
+```{eval-rst}
+.. autofunction:: set_backend
+```
+
+```{eval-rst}
+.. autofunction:: get_backend
+```
+
+## Op Reference
+:::{note}
+The following ops are hosted in the `torch.ops.symm_mem` namespace. You can call
+them directly via `torch.ops.symm_mem.<op_name>`.
+:::
+
+```{eval-rst}
+.. currentmodule:: torch.ops.symm_mem
+```
+
+```{eval-rst}
+.. py:function:: multimem_all_reduce_(input: Tensor, reduce_op: str, group_name: str) -> Tensor
+
+    Performs a multimem all-reduce operation on the input tensor. This operation
+    requires hardware support for multimem operations. On NVIDIA GPUs, NVLink
+    SHARP is required.
+
+    :param Tensor input: Input tensor to perform all-reduce on. Must be symmetric.
+    :param str reduce_op: Reduction operation to perform. Currently only "sum" is supported.
+    :param str group_name: Name of the group to perform all-reduce on.
+
+
+.. py:function:: multimem_all_gather_out(input: Tensor, group_name: str, out: Tensor) -> Tensor
+
+    Performs a multimem all-gather operation on the input tensor. This operation requires hardware support for multimem operations. On NVIDIA GPUs, NVLink SHARP is required.
+
+    :param Tensor input: Input tensor to perform all-gather on.
+    :param str group_name: Name of the group to perform all-gather on.
+    :param Tensor out: Output tensor to store the result of the all-gather operation. Must be symmetric.
+
+
+.. py:function:: one_shot_all_reduce(input: Tensor, reduce_op: str, group_name: str) -> Tensor
+
+    Performs a one-shot all-reduce operation on the input tensor.
+
+    :param Tensor input: Input tensor to perform all-reduce on. Must be symmetric.
+    :param str reduce_op: Reduction operation to perform. Currently only "sum" is supported.
+    :param str group_name: Name of the group to perform all-reduce on.
+
+
+.. py:function:: one_shot_all_reduce_out(input: Tensor, reduce_op: str, group_name: str, out: Tensor) -> Tensor
+
+    Performs a one-shot all-reduce operation based on the input tensor and writes the result to the output tensor.
+
+    :param Tensor input: Input tensor to perform all-reduce on. Must be symmetric.
+    :param str reduce_op: Reduction operation to perform. Currently only "sum" is supported.
+    :param str group_name: Name of the group to perform all-reduce on.
+    :param Tensor out: Output tensor to store the result of the all-reduce operation. Can be a regular tensor.
+
+
+.. py:function:: two_shot_all_reduce_(input: Tensor, reduce_op: str, group_name: str) -> Tensor
+
+    Performs a two-shot all-reduce operation on the input tensor.
+
+    :param Tensor input: Input tensor to perform all-reduce on. Must be symmetric.
+    :param str reduce_op: Reduction operation to perform. Currently only "sum" is supported.
+    :param str group_name: Name of the group to perform all-reduce on.
+
+
+.. py:function:: all_to_all_vdev(input: Tensor, out: Tensor, in_splits: Tensor, out_splits_offsets: Tensor, group_name: str) -> None
+
+    Performs an all-to-all-v operation using NVSHMEM, with split information provided on device.
+
+    :param Tensor input: Input tensor to perform all-to-all on. Must be symmetric.
+    :param Tensor out: Output tensor to store the result of the all-to-all operation. Must be symmetric.
+    :param Tensor in_splits: Tensor containing splits of data to send to each peer. Must be symmetric. Must be of size (group_size,). The splits are in the unit of elements in the 1st dimension.
+    :param Tensor out_splits_offsets: Tensor containing the splits and offsets of data received from each peer. Must be symmetric. Must be of size (2, group_size). The rows are (in order): output splits and output offsets.
+    :param str group_name: Name of the group to perform all-to-all on.
+
+
+.. py:function:: all_to_all_vdev_2d(input: Tensor, out: Tensor, in_splits: Tensor, out_splits_offsets: Tensor, group_name: str, [major_align: int = None]) -> None
+
+    Perform a 2D all-to-all-v operation using NVSHMEM, with split information provided on device. In Mixture of Experts models, this operation can be used to dispatch tokens.
+
+    :param Tensor input: Input tensor to perform all-to-all on. Must be symmetric.
+    :param Tensor out: Output tensor to store the result of the all-to-all operation. Must be symmetric.
+    :param Tensor in_splits: Tensor containing the splits of data to send to each expert. Must be symmetric. Must be of size (group_size * ne,), where ne is the number of experts per rank. The splits are in the unit of elements in the 1st dimension.
+    :param Tensor out_splits_offsets: Tensor containing the splits and offsets of data received from each peer. Must be symmetric. Must be of size (2, group_size * ne). The rows are (in order): output splits and output offsets.
+    :param str group_name: Name of the group to perform all-to-all on.
+    :param int major_align: Optional alignment for the major dimension of the output chunk for each expert. If not provided, the alignment is assumed to be 1. Any alignment adjustment will be reflected in the output offsets.
+
+    A 2D AllToAllv shuffle is illustrated below:
+    (world_size = 2, ne = 2, total number of experts = 4)::
+
+      Source: |       Rank 0      |       Rank 1      |
+              | c0 | c1 | c2 | c3 | d0 | d1 | d2 | d3 |
+
+      Dest  : |       Rank 0      |       Rank 1      |
+              | c0 | d0 | c1 | d1 | c2 | d2 | c3 | d3 |
+
+    where each `c_i` / `d_i` are slices of the `input` tensor, targeting expert
+    `i`, with length indicated by input splits.  That is, the 2D AllToAllv
+    shuffle achieves a transpose from rank-major order at input to expert-major
+    order at output.
+
+    If `major_align` is not 1, the output offsets of c1, c2, c3 will be
+    up-aligned to this value. For example, if c0 has length 5 and d0 has
+    length 7 (making a total of 12), and if the `major_align` is set to 16,
+    the output offset of c1 will be 16. Similar for c2 and c3. This value has
+    no effect on the offset of the minor dimension, i.e.  d0, d1, d2 and d3.
+    Note: since cutlass does not support empty bins, we set the aligned length
+    to `major_align` if it is 0. See
+    https://github.com/pytorch/pytorch/issues/152668.
+
+
+.. py:function:: all_to_all_vdev_2d_offset(Tensor input, Tensor out, Tensor in_splits_offsets, Tensor out_splits_offsets, str group_name) -> None
+
+    Perform a 2D AllToAllv shuffle operation, with input split and offset
+    information provided on device. The input offsets are not required to be
+    exact prefix sum of the input splits, i.e. paddings are allowed between the
+    split chunks. The paddings, however, will not be transferred to peer
+    ranks.
+
+    In Mixture of Experts models, this operation can be used to combine tokens
+    processed by experts on parallel ranks. This operation can be viewed as an
+    "reverse" operation to the `all_to_all_vdev_2d` operation (which shuffles
+    tokens to experts).
+
+    :param Tensor input: Input tensor to perform all-to-all on. Must be symmetric.
+    :param Tensor out: Output tensor to store the result of the all-to-all operation. Must be symmetric.
+    :param Tensor in_splits_offsets: Tensor containing the splits and offsets of data to send to each expert. Must be symmetric. Must be of size (2, group_size * ne), where `ne` is the number of experts. The rows are (in order): input splits and input offsets. The splits are in the unit of elements in the 1st dimension.
+    :param Tensor out_splits_offsets: Tensor containing the splits and offsets of data received from each peer. Must be symmetric. Must be of size (2, group_size * ne). The rows are (in order): output splits and output offsets.
+    :param str group_name: Name of the group to perform all-to-all on.
+
+```

pyrefly.toml

@@ -130,6 +130,7 @@ errors.bad-param-name-override = false
 # Mypy doesn't require that imports are explicitly imported, so be compatible with that.
 # Might be a good idea to turn this on in future.
 errors.implicit-import = false
+errors.deprecated = false # re-enable after we've fix import formatting
 permissive-ignores = true
 replace-imports-with-any = ["!sympy.printing.*", "sympy.*", "onnxscript.onnx_opset.*"]
 search-path = ["tools/experimental"]

test/distributed/elastic/agent/server/test/api_test.py

@@ -257,7 +257,10 @@ class RecordWorkerEventsTest(unittest.TestCase):
         self.assertEqual(failed_event.metadata["state"], "FAILED")
         self.assertEqual(failed_event.metadata["global_rank"], 0)
         md = json.loads(failed_event.metadata["metadata"])
-        self.assertEqual(failed_event.metadata["raw_error"], '{"message": "<NONE>"}')
+        self.assertEqual(
+            failed_event.metadata["raw_error"],
+            '{"message": "<NONE>", "errorTraits": {"category": "system_terminated_error", "retryability": "False"}}',
+        )
         self.assertEqual(md["exit_code"], [1])
         self.assertEqual(md["worker_pid"], [1000])
 

test/distributed/tensor/test_dtensor_export.py

@@ -22,7 +22,6 @@ from torch.distributed.tensor.parallel import (
     parallelize_module,
     RowwiseParallel,
 )
-from torch.nn.attention.flex_attention import create_block_mask, flex_attention
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
     parametrize,
@@ -83,46 +82,7 @@ class SimpleModelAnnotated(torch.nn.Module):
         return self.mlp_1(x)
 
 
-class FlexAttentionModel(torch.nn.Module):
-    def __init__(self, device):
-        super().__init__()
-        self.proj_q = torch.nn.Linear(16, 128, device=device)
-        self.proj_k = torch.nn.Linear(16, 128, device=device)
-        self.proj_v = torch.nn.Linear(16, 128, device=device)
-        self.proj_out = torch.nn.Linear(128, 16, device=device)
-        self.num_heads = 8
-        self.head_dim = 16
-
-    def forward(self, x, *, block_mask=None):
-        batch_size, seq_len, embed_dim = x.shape
-        # Project to Q, K, V
-        q = self.proj_q(x)
-        k = self.proj_k(x)
-        v = self.proj_v(x)
-        # After colwise parallel, q/k/v are sharded on the last dimension
-        # Get the actual size after sharding
-        hidden_size = q.shape[-1]
-        num_heads_local = hidden_size // self.head_dim
-        # Reshape to (batch, num_heads, seq_len, head_dim)
-        q = q.view(batch_size, seq_len, num_heads_local, self.head_dim).transpose(1, 2)
-        k = k.view(batch_size, seq_len, num_heads_local, self.head_dim).transpose(1, 2)
-        v = v.view(batch_size, seq_len, num_heads_local, self.head_dim).transpose(1, 2)
-        # Apply flex_attention
-        attn_output_raw = flex_attention(q, k, v, block_mask=block_mask)
-        # Reshape back to (batch, seq_len, hidden_size)
-        attn_output = (
-            attn_output_raw.transpose(1, 2)
-            .contiguous()
-            .view(batch_size, seq_len, hidden_size)
-        )
-        # Output projection
-        output = self.proj_out(attn_output)
-        return output
-
-
-def strict_export_and_aot_export_joint_with_descriptors(model, args, kwargs=None):
-    if kwargs is None:
-        kwargs = {}
+def strict_export_and_aot_export_joint_with_descriptors(model, inputs):
     # needed for stric export
     torch.utils._pytree.register_constant(DTensorSpec)
 
@@ -131,43 +91,36 @@ def strict_export_and_aot_export_joint_with_descriptors(model, args, kwargs=None
         install_free_tensors=True, inline_inbuilt_nn_modules=True
     ):
         with torch._export.utils._disable_aten_to_metadata_assertions():
-            ep = torch.export.export(model, args, kwargs, strict=True)
+            ep = torch.export.export(model, (inputs,), strict=True)
 
     # joint_gm produced here is missing the backward region, due to incompatiblility
     # between ep.module() and aot_export_joint_with_descriptors.
     # Keeping this here to show the issue.
-    return aot_export_joint_with_descriptors_alone(ep.module(), args, kwargs)
+    return aot_export_joint_with_descriptors_alone(ep.module(), inputs)
 
 
-def graph_capture_and_aot_export_joint_with_descriptors_v2(model, args, kwargs=None):
-    if kwargs is None:
-        kwargs = {}
-    gm = dynamo_graph_capture_for_export(model)(*args, **kwargs)
+def graph_capture_and_aot_export_joint_with_descriptors_v2(model, inputs):
+    gm = dynamo_graph_capture_for_export(model)(inputs)
     fake_mode = gm.meta.get("fake_mode", None)
     with tracing(TracingContext(fake_mode)):
-        return aot_export_joint_with_descriptors_alone(gm, args, kwargs)
+        return aot_export_joint_with_descriptors_alone(gm, inputs)
 
 
-def graph_capture_and_aot_export_joint_with_descriptors(model, args, kwargs=None):
-    if kwargs is None:
-        kwargs = {}
+def graph_capture_and_aot_export_joint_with_descriptors(model, inputs):
     with torch._dynamo.config.patch(install_free_tensors=True):
         # TODO: switch to use the official graph_capture API once it is ready
-        gm = _dynamo_graph_capture_for_export(model)(*args, **kwargs)
+        gm = _dynamo_graph_capture_for_export(model)(inputs)
         fake_mode = gm.meta.get("fake_mode", None)
     with tracing(TracingContext(fake_mode)):
-        return aot_export_joint_with_descriptors_alone(gm, args, kwargs)
+        return aot_export_joint_with_descriptors_alone(gm, inputs)
 
 
-def aot_export_joint_with_descriptors_alone(model, args, kwargs=None):
-    if kwargs is None:
-        kwargs = {}
+def aot_export_joint_with_descriptors_alone(model, inputs):
     with contextlib.ExitStack() as stack:
         joint_with_descriptors = aot_export_joint_with_descriptors(
             stack,
             model,
-            args,
-            kwargs,
+            (inputs,),
         )
         return joint_with_descriptors.graph_module
 
@@ -215,8 +168,8 @@ class DTensorExportTest(TestCase):
         }
         tp_model = parallelize_module(model, mesh_2d["tp"], parallelize_plan)
 
-        inp = torch.rand(20, 10, device=self.device_type)
-        inputs = (distribute_tensor(inp, mesh_2d["tp"], placements=[Replicate()]),)
+        inputs = torch.rand(20, 10, device=self.device_type)
+        inputs = distribute_tensor(inputs, mesh_2d["tp"], placements=[Replicate()])
 
         joint_gm = export_fn(tp_model, inputs)
         fw_gm, bw_gm = min_cut_rematerialization_partition(
@@ -399,10 +352,9 @@ class DTensorExportTest(TestCase):
         }
         tp_model = parallelize_module(model, mesh_2d["tp"], parallelize_plan)
 
-        inp = torch.rand(20, 10, device=self.device_type)
-        inp_dtensor = distribute_tensor(inp, mesh_2d["tp"], placements=[Replicate()])
-        torch._dynamo.mark_dynamic(inp_dtensor, 0, min=5, max=100)
-        inputs = (inp_dtensor,)
+        inputs = torch.rand(20, 10, device=self.device_type)
+        inputs = distribute_tensor(inputs, mesh_2d["tp"], placements=[Replicate()])
+        torch._dynamo.mark_dynamic(inputs, 0, min=5, max=100)
 
         joint_gm = export_fn(tp_model, inputs)
 
@@ -438,67 +390,15 @@ class DTensorExportTest(TestCase):
         z = torch.randn(16, 16)
         y_dtensor = distribute_tensor(y, device_mesh, placements=[Replicate()])
         z_dtensor = DTensor.from_local(z, device_mesh, placements=[Partial()])
-        inputs = (x_dtensor, y_dtensor, z_dtensor)
 
         # Run model to verify it works
-        output = model(*inputs)
+        output = model(x_dtensor, y_dtensor, z_dtensor)
         with torch._dynamo.config.patch(install_free_tensors=True):
             # TODO: switch to use the official graph_capture API once it is ready
-            gm = export_fn(model)(*inputs)
-        output_gm = gm(*inputs)
+            gm = export_fn(model)(x_dtensor, y_dtensor, z_dtensor)
+        output_gm = gm(x_dtensor, y_dtensor, z_dtensor)
         self.assertEqual(output, output_gm)
 
-    def test_flex_attention_dtensor_export(self):
-        device_mesh = init_device_mesh(self.device_type, mesh_shape=(self.world_size,))
-        model = FlexAttentionModel(self.device_type)
-
-        # Parallelize the model: shard on head dimension
-        # proj_q, proj_k, proj_v are colwise parallel (output is sharded on head dimension)
-        # proj_out is rowwise parallel (input is sharded, output needs reduction)
-        parallelize_plan = {
-            "proj_q": ColwiseParallel(),
-            "proj_k": ColwiseParallel(),
-            "proj_v": ColwiseParallel(),
-            "proj_out": RowwiseParallel(),
-        }
-        tp_model = parallelize_module(model, device_mesh, parallelize_plan)
-        batch_size = 4
-        seq_len = 64
-        embed_dim = 16
-        num_heads = 8
-
-        # Input tensor replicated across all devices
-        inp = torch.randn(batch_size, seq_len, embed_dim, device=self.device_type)
-        inputs = (distribute_tensor(inp, device_mesh, placements=[Replicate()]),)
-
-        def causal_mask(b, h, q_idx, kv_idx):
-            return q_idx >= kv_idx
-
-        block_mask = create_block_mask(
-            causal_mask,
-            batch_size,
-            num_heads,
-            seq_len,
-            seq_len,
-            device=self.device_type,
-        )
-
-        flex_kwargs = {"block_mask": block_mask}
-
-        joint_gm = graph_capture_and_aot_export_joint_with_descriptors(
-            tp_model, inputs, flex_kwargs
-        )
-
-        self.assertTrue(
-            _count_op(joint_gm, torch.ops.higher_order.flex_attention),
-            1,
-        )
-
-        self.assertTrue(
-            _count_op(joint_gm, torch.ops.higher_order.flex_attention_backward),
-            2,
-        )
-
 
 instantiate_parametrized_tests(DTensorExportTest)
 

test/distributed/tensor/test_math_ops.py

@@ -13,6 +13,7 @@ from torch.distributed.tensor import (
     distribute_module,
     distribute_tensor,
     DTensor,
+    Partial,
     Replicate,
     Shard,
 )
@@ -649,6 +650,41 @@ class DistMathOpsTest(DTensorTestBase):
         self.assertEqual(len(comm_counts), 1)
         self.assertEqual(comm_counts[funcol.all_gather_into_tensor], 1)
 
+    @with_comms
+    def test_vector_norm(self):
+        device_mesh = self.build_device_mesh()
+
+        grad = torch.randn(12, 8)
+
+        sharded_grad = distribute_tensor(grad, device_mesh, [Shard(0)])
+
+        # non-sharded op
+        out = torch.ops.aten.linalg_vector_norm(grad, 2)
+
+        # sharded op
+        sharded_out = torch.ops.aten.linalg_vector_norm(sharded_grad, 2)
+
+        self.assertEqual(sharded_out.full_tensor(), out)
+
+    @with_comms
+    def test_vector_norm_partial(self):
+        device_mesh = self.build_device_mesh()
+
+        rank = device_mesh.get_local_rank()
+        all_ranks = list(range(self.world_size))
+
+        local_grad = torch.tensor([rank, 1], dtype=torch.float32)
+        full_grad = torch.tensor([sum(all_ranks), self.world_size], dtype=torch.float32)
+
+        partial_grad = DTensor.from_local(local_grad, device_mesh, [Partial()])
+
+        # full result
+        out = torch.ops.aten.linalg_vector_norm(full_grad, 2)
+
+        # partial result
+        partial_out = torch.ops.aten.linalg_vector_norm(partial_grad, 2)
+        self.assertEqual(partial_out.full_tensor(), out)
+
     @with_comms
     def test_foreach_norm(self):
         device_mesh = self.build_device_mesh()
@@ -668,6 +704,33 @@ class DistMathOpsTest(DTensorTestBase):
         for o, so in zip(out, sharded_out):
             self.assertEqual(so.full_tensor(), o)
 
+    @with_comms
+    def test_foreach_norm_partial(self):
+        device_mesh = self.build_device_mesh()
+
+        rank = device_mesh.get_local_rank()
+        all_ranks = list(range(self.world_size))
+
+        local_grad0 = torch.tensor([rank, 1], dtype=torch.float32)
+        local_grad1 = torch.tensor([rank + 1, 2], dtype=torch.float32)
+
+        grad0 = torch.tensor([sum(all_ranks), self.world_size], dtype=torch.float32)
+        grad1 = torch.tensor(
+            [sum(all_ranks) + self.world_size, 2 * self.world_size], dtype=torch.float32
+        )
+
+        partial_grad0 = DTensor.from_local(local_grad0, device_mesh, [Partial()])
+        partial_grad1 = DTensor.from_local(local_grad1, device_mesh, [Partial()])
+
+        # full result
+        out = torch.ops.aten._foreach_norm([grad0, grad1], 2)
+
+        # partial result
+        partial_out = torch.ops.aten._foreach_norm([partial_grad0, partial_grad1], 2)
+
+        for o, po in zip(out, partial_out):
+            self.assertEqual(po.full_tensor(), o)
+
     @with_comms
     def test_foreach_norm_different_mesh(self):
         mesh_shape = (2, self.world_size // 2)

test/distributed/tensor/test_redistribute.py

@@ -7,6 +7,10 @@ import itertools
 import unittest
 
 import torch
+from torch.distributed._local_tensor import (
+    maybe_disable_local_tensor_mode,
+    maybe_run_for_local_tensor,
+)
 from torch.distributed.device_mesh import init_device_mesh
 from torch.distributed.tensor import (
     DeviceMesh,
@@ -29,7 +33,9 @@ from torch.testing._internal.common_utils import (
     TEST_HPU,
 )
 from torch.testing._internal.distributed._tensor.common_dtensor import (
+    create_local_tensor_test_class,
     DTensorTestBase,
+    map_local_tensor_for_rank,
     with_comms,
 )
 from torch.utils._debug_mode import DebugMode
@@ -163,7 +169,9 @@ class RedistributeTest(DTensorTestBase):
             )
 
             # make local tensor as the element of the corresponding chunked list
-            local_tensor = splitted_list[self.rank]
+            local_tensor = map_local_tensor_for_rank(
+                splitted_list, self.rank, lambda tl, r: tl[r]
+            )
             replica_tensor = distribute_tensor(local_replica, device_mesh, replica_spec)
             with comm_mode:
                 reshard_tensor = replica_tensor.redistribute(device_mesh, shard_spec)
@@ -407,7 +415,7 @@ class RedistributeTest(DTensorTestBase):
     def test_partial_to_shard(self, dtype):
         device_mesh = self.build_device_mesh()
         partial_spec = [Partial()]
-        my_rank = device_mesh.get_rank()
+        my_rank = self.rank
 
         input_sizes_and_shard_dim = [
             ((self.world_size * 3, 3), 0),
@@ -440,8 +448,13 @@ class RedistributeTest(DTensorTestBase):
                 for idx in range(self.world_size)
             ]
 
-            local_shape = list(input_size)
-            local_shape[shard_dim] = chunk_sizes[my_rank]
+            @maybe_run_for_local_tensor
+            def _compute_local_shape(rank) -> list[int]:
+                local_shape = list(input_size)
+                local_shape[shard_dim] = chunk_sizes[rank]
+                return local_shape
+
+            local_shape = _compute_local_shape(my_rank)
 
             # test partial to shard, trigger reduce_scatter
             with comm_mode:
@@ -534,10 +547,12 @@ class RedistributeTest(DTensorTestBase):
                         1,
                     )
                 else:
-                    self.assertEqual(
-                        comm_mode.get_comm_counts()[funcol.all_gather_into_tensor],
-                        1,
-                    )
+                    # TODO: Integrate local tensor with CommDebugMode
+                    if not self.is_local_tensor_enabled:
+                        self.assertEqual(
+                            comm_mode.get_comm_counts()[funcol.all_gather_into_tensor],
+                            1,
+                        )
 
         # test 2d device mesh
         mesh_2d = DeviceMesh(
@@ -586,7 +601,8 @@ class RedistributeTest(DTensorTestBase):
                     out_dt = sharded_dt.redistribute(mesh_2d, dst)
 
                 self.assertEqual(out_dt.placements, expected_dt.placements)
-                self.assertEqual(comm_mode.get_total_counts(), comm_counts_2d[idx])
+                if not self.is_local_tensor_enabled:
+                    self.assertEqual(comm_mode.get_total_counts(), comm_counts_2d[idx])
 
                 local_out_dt = out_dt.to_local()
                 local_expected_dt = expected_dt.to_local()
@@ -1027,23 +1043,27 @@ class DistributeWithDeviceOrderTest(DTensorTestBase):
     def test_ordered_distribute_all_combination(self):
         """Exhaustively test all possible sharding combinations and verify correctness"""
         torch.manual_seed(21)
-        mesh = init_device_mesh(self.device_type, (2, 2, 2))
-        input_tensor_shape = [
-            # even sharding
-            (16, 8),
-            (8, 16, 32),
-            (8, 32, 16, 16),
-            # uneven sharding with padding
-            (17, 5),
-            (13, 2, 13),
-            (33, 16, 8, 1),
-        ]
+
+        with maybe_disable_local_tensor_mode():
+            mesh = init_device_mesh(self.device_type, (2, 2, 2))
+            input_tensor_shape = [
+                # even sharding
+                (16, 8),
+                (8, 16, 32),
+                (8, 32, 16, 16),
+                # uneven sharding with padding
+                (17, 5),
+                (13, 2, 13),
+                (33, 16, 8, 1),
+            ]
 
         # 1. Verify correctness of distribute_tensor from Tensor to DTensor.
         for tensor_shape in input_tensor_shape:
             input_data = torch.randn(tensor_shape, device=self.device_type)
             tensor_rank = input_data.ndim
-            for shard_order in self.generate_shard_orders(mesh, tensor_rank):
+            with maybe_disable_local_tensor_mode():
+                shard_orders = self.generate_shard_orders(mesh, tensor_rank)
+            for shard_order in shard_orders:
                 sharded_dt = self.distribute_tensor(
                     input_data.clone(), mesh, placements=None, shard_order=shard_order
                 )
@@ -1057,7 +1077,9 @@ class DistributeWithDeviceOrderTest(DTensorTestBase):
             input_data = torch.randn(tensor_shape, device=self.device_type)
             tensor_rank = input_data.ndim
             prev_sharded_dt = None
-            for shard_order in self.generate_shard_orders(mesh, tensor_rank):
+            with maybe_disable_local_tensor_mode():
+                shard_orders = self.generate_shard_orders(mesh, tensor_rank)
+            for shard_order in shard_orders:
                 if prev_sharded_dt is None:
                     prev_sharded_dt = self.distribute_tensor(
                         input_data.clone(),
@@ -1077,26 +1099,27 @@ class DistributeWithDeviceOrderTest(DTensorTestBase):
         """Test mixing Partial in the original placements and do redistribute."""
         # This test takes 226s to complete on 8XA100...
         torch.manual_seed(21)
-        mesh = init_device_mesh(self.device_type, (2, 2, 2))
-        input_tensor_shape = [
-            # even sharding
-            (16, 8),
-            (8, 16, 32),
-            # uneven sharding with padding
-            (17, 5),
-            (13, 2, 13),
-            (33, 16, 8, 1),
-        ]
-        placement_choice = [
-            Shard(0),
-            Shard(1),
-            Shard(2),
-            Partial("sum"),
-            Partial("min"),
-            Replicate(),
-        ]
-        # pick 3 for the 3D mesh
-        partial_placement_comb = list(itertools.combinations(placement_choice, 3))
+        with maybe_disable_local_tensor_mode():
+            mesh = init_device_mesh(self.device_type, (2, 2, 2))
+            input_tensor_shape = [
+                # even sharding
+                (16, 8),
+                (8, 16, 32),
+                # uneven sharding with padding
+                (17, 5),
+                (13, 2, 13),
+                (33, 16, 8, 1),
+            ]
+            placement_choice = [
+                Shard(0),
+                Shard(1),
+                Shard(2),
+                Partial("sum"),
+                Partial("min"),
+                Replicate(),
+            ]
+            # pick 3 for the 3D mesh
+            partial_placement_comb = list(itertools.combinations(placement_choice, 3))
 
         def _is_valid_placement(placements, tensor_rank):
             # Check if placements is valid for tensor with rank `tensor_rank`
@@ -1112,7 +1135,9 @@ class DistributeWithDeviceOrderTest(DTensorTestBase):
                     continue
                 local_tensor = torch.randn(shape, device=self.device_type)
                 full_tensor = DTensor.from_local(local_tensor, mesh, placements)
-                for shard_order in self.generate_shard_orders(mesh, len(shape)):
+                with maybe_disable_local_tensor_mode():
+                    shard_orders = self.generate_shard_orders(mesh, len(shape))
+                for shard_order in shard_orders:
                     sharded_dt = self.redistribute(
                         full_tensor, mesh, placements=None, shard_order=shard_order
                     )
@@ -1163,5 +1188,18 @@ class DistributeWithDeviceOrderTest(DTensorTestBase):
         self.assertEqual(x_ordered_dt.to_local(), x_strided_dt.to_local())
 
 
+RedistributeTestWithLocalTensor = create_local_tensor_test_class(
+    RedistributeTest,
+)
+
+MultiDimRedistributeTestWithLocalTensor = create_local_tensor_test_class(
+    MultiDimRedistributeTest,
+    skipped_tests=["test_multi_dim_mesh"],
+)
+
+DistributeWithDeviceOrderTestWithLocalTensor = create_local_tensor_test_class(
+    DistributeWithDeviceOrderTest,
+)
+
 if __name__ == "__main__":
     run_tests()

test/distributed/test_device_mesh.py

@@ -462,7 +462,9 @@ class DeviceMeshTestNDim(DTensorTestBase):
         ep_mesh_2 = DeviceMesh(self.device_type, mesh_group_2)
         ep_mesh = ep_mesh_1 if self.rank < self.world_size // 2 else ep_mesh_2
         # ep_mesh is considered different from mesh_2d["TP"]
-        self.assertEqual(mesh_2d["TP"]._flatten_mesh_list, ep_mesh._flatten_mesh_list)
+        self.assertEqual(
+            mesh_2d["TP"].mesh.flatten().tolist(), ep_mesh.mesh.flatten().tolist()
+        )
         self.assertEqual(mesh_2d["TP"]._layout, ep_mesh._layout)
         self.assertEqual(mesh_2d["TP"].mesh.shape, ep_mesh.mesh.shape)
         self.assertEqual(mesh_2d["TP"].device_type, ep_mesh.device_type)
@@ -477,7 +479,7 @@ class DeviceMeshTestNDim(DTensorTestBase):
             another_mesh_1 if self.rank < self.world_size // 2 else another_mesh_2
         )
         # another_mesh is considered the same as ep_mesh
-        self.assertEqual(ep_mesh._flatten_mesh_list, another_mesh._flatten_mesh_list)
+        self.assertEqual(ep_mesh._flatten_rank_map, another_mesh._flatten_rank_map)
         self.assertEqual(ep_mesh._layout, another_mesh._layout)
         self.assertEqual(ep_mesh.mesh.shape, another_mesh.mesh.shape)
         self.assertEqual(ep_mesh.device_type, another_mesh.device_type)

test/dynamo/test_dicts.py

@@ -1076,7 +1076,22 @@ class DictTests(torch._dynamo.test_case.TestCase):
     def test_newly_constructed_default_dict(self):
         def f(x):
             d = defaultdict(list)
-            d[0] = 42
+            d[0] = [
+                42,
+            ]
+            return x + 1, d
+
+        x = torch.ones(2)
+        ref = f(x)
+        res = torch.compile(f, backend="eager", fullgraph=True)(x)
+
+        self.assertEqual(ref, res)
+
+    @unittest.expectedFailure
+    def test_newly_constructed_default_dict_with_dict(self):
+        def f(x):
+            d = defaultdict(dict, {2: {"a": 1}})
+            d[0] = {"b": 2}
             return x + 1, d
 
         x = torch.ones(2)

test/dynamo/test_guard_serialization.py

@@ -747,7 +747,7 @@ class TestGuardSerialization(TestGuardSerializationBase):
             ):
                 self._test_serialization("NN_MODULE", fn, m, x)
 
-    def test_function_match(self):
+    def test_class_match(self):
         def fn(x):
             # usage of this context manager installs a FUNCTION_MATCH guard
             with torch.no_grad():
@@ -759,9 +759,9 @@ class TestGuardSerialization(TestGuardSerializationBase):
         # we don't support FUNCTION_MATCH because it adds an ID_MATCH guard, and we don't
         # support that in serialization
         with self.assertRaisesRegex(
-            PackageError, "FUNCTION_MATCH guard cannot be serialized."
+            PackageError, "CLASS_MATCH guard cannot be serialized."
         ):
-            self._test_serialization("FUNCTION_MATCH", fn, x)
+            self._test_serialization("CLASS_MATCH", fn, x)
 
     def test_closure_match(self):
         def fn(x):
@@ -957,12 +957,12 @@ class TestGuardSerialization(TestGuardSerializationBase):
         self._test_check_fn(ref, loaded, {"x": torch.randn(3)}, True)
 
         def fn(x):
-            # usage of this context manager installs a FUNCTION_MATCH guard
+            # usage of this context manager installs a CLASS_MATCH guard
             with torch.no_grad():
                 y = x * 2
             return y
 
-        ref, loaded = self._test_serialization("FUNCTION_MATCH", fn, torch.randn(3))
+        ref, loaded = self._test_serialization("CLASS_MATCH", fn, torch.randn(3))
         self._test_check_fn(ref, loaded, {"x": torch.randn(3)}, True)
 
     def test_dispatch_key_set_match(self):
@@ -1467,7 +1467,8 @@ class TestGuardSerialization(TestGuardSerializationBase):
             torch._dynamo.optimize(
                 package=package,
                 guard_filter_fn=lambda gs: [
-                    x.guard_type not in ("CLOSURE_MATCH", "ID_MATCH") for x in gs
+                    x.guard_type not in ("CLOSURE_MATCH", "ID_MATCH", "CLASS_MATCH")
+                    for x in gs
                 ],
             )(foo)(ddp_model, x)
             self.assertEqual(len(package._codes[foo.__code__].guarded_codes), 1)

test/dynamo/test_logging.py

@@ -245,8 +245,7 @@ due to:
 Traceback (most recent call last):
   File "test_logging.py", line N, in throw
     raise AssertionError
-torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
-LoweringException: AssertionError:
+torch._inductor.exc.InductorError: LoweringException: AssertionError:
   target: aten.round.default
   args[0]: TensorBox(StorageBox(
     InputBuffer(name='primals_1', layout=FixedLayout('cpu', torch.float32, size=[1000, 1000], stride=[1000, 1]))

test/dynamo/test_package.py

@@ -326,7 +326,8 @@ def add(x, y):
 
             def guard_filter_fn(guards):
                 return [
-                    guard.guard_type not in ("CLOSURE_MATCH", "FUNCTION_MATCH")
+                    guard.guard_type
+                    not in ("CLOSURE_MATCH", "FUNCTION_MATCH", "MODULE_MATCH")
                     for guard in guards
                 ]
 

test/dynamo/test_regional_inductor.py

@@ -1,7 +1,6 @@
 # Owner(s): ["module: dynamo"]
 
 import functools
-import unittest
 
 import torch
 import torch._inductor.test_case
@@ -155,10 +154,6 @@ class RegionalInductorTests(torch._inductor.test_case.TestCase):
     def test_invoke_subgraph_inner(self, serialize):
         # Checks that the inductor regions are searched recursively.
 
-        # TODO: GraphPickler does not recompile nested subgraphs?
-        if serialize:
-            raise unittest.SkipTest("GraphPickler doesn't recompile nested subgraphs")
-
         @torch.compiler.nested_compile_region
         def gn(x):
             with fx_traceback.annotate({"compile_with_inductor": 0}):
@@ -184,12 +179,6 @@ class RegionalInductorTests(torch._inductor.test_case.TestCase):
     @requires_cuda_and_triton
     @parametrize("serialize", [False, True])
     def test_flex_attention(self, serialize):
-        if serialize:
-            # TODO: Fixed in next PR
-            raise unittest.SkipTest(
-                "FlexAttentionBackward isn't marked cacheable even though it is"
-            )
-
         def _squared(score, b, h, m, n):
             return score * score
 
@@ -229,11 +218,6 @@ class RegionalInductorTests(torch._inductor.test_case.TestCase):
     @requires_cuda_and_triton
     @parametrize("serialize", [False, True])
     def test_selective_ac_flex(self, serialize):
-        if serialize:
-            raise unittest.SkipTest(
-                "FlexAttentionBackward isn't marked cacheable even though it is"
-            )
-
         class FlexAttentionModule(torch.nn.Module):
             def __init__(self, hidden_size, num_heads):
                 super().__init__()

test/expect/TestFXAPIBackwardCompatibility.test_function_back_compat-fx_backcompat_function_signatures.expect

@@ -55,7 +55,7 @@ torch.fx.node.Node.append(self, x: 'Node') -> None
 torch.fx.node.Node.format_node(self, placeholder_names: Optional[List[str]] = None, maybe_return_typename: Optional[List[str]] = None, include_tensor_metadata: bool = False) -> Optional[str]
 torch.fx.node.Node.insert_arg(self, idx: int, arg: torch.fx.node.Argument) -> None
 torch.fx.node.Node.prepend(self, x: 'Node') -> None
-torch.fx.node.Node.replace_all_uses_with(self, replace_with: 'Node', delete_user_cb: Callable[[Node], bool] = <function <lambda>>, propagate_meta: bool = False) -> List[Node]
+torch.fx.node.Node.replace_all_uses_with(self, replace_with: 'Node', delete_user_cb: Optional[Callable[[Node], bool]] = None, propagate_meta: bool = False) -> List[Node]
 torch.fx.node.Node.replace_input_with(self, old_input: 'Node', new_input: 'Node') -> None
 torch.fx.node.Node.update_arg(self, idx: int, arg: torch.fx.node.Argument) -> None
 torch.fx.node.Node.update_kwarg(self, key: str, arg: torch.fx.node.Argument) -> None

test/inductor/test_flex_attention.py

@@ -5734,141 +5734,6 @@ BlockMask(shape=(1,s1,s2048,s2048),ssparsity=46.88%,s
 
         self.assertEqual(flex_output, sdpa_output, atol=1e-3, rtol=1e-3)
 
-    @supported_platform
-    def test_pytree_flatten_unflatten(self, device):
-        """Test that BlockMask can be correctly flattened and unflattened using class methods."""
-
-        def causal_mask(b, h, q_idx, kv_idx):
-            return q_idx >= kv_idx
-
-        # Create a BlockMask with various attributes set
-        block_mask = create_block_mask(
-            causal_mask, B=2, H=4, Q_LEN=512, KV_LEN=512, device=device
-        )
-
-        # Flatten and unflatten using class methods
-        tensors, context = block_mask._flatten()
-        reconstructed_mask = BlockMask._unflatten(tensors, context)
-
-        # Verify the reconstructed mask has the same attributes
-        self.assertEqual(reconstructed_mask.shape, block_mask.shape)
-        self.assertEqual(reconstructed_mask.sparsity(), block_mask.sparsity())
-
-        # Verify all tensor attributes are equal (using _TENSOR_ATTRS)
-        for attr_name in BlockMask._TENSOR_ATTRS:
-            original_value = getattr(block_mask, attr_name)
-            reconstructed_value = getattr(reconstructed_mask, attr_name)
-
-            if original_value is None:
-                self.assertIsNone(
-                    reconstructed_value,
-                    f"Tensor attribute {attr_name} should be None but got {reconstructed_value}",
-                )
-            else:
-                self.assertIsInstance(
-                    original_value,
-                    torch.Tensor,
-                    f"Expected {attr_name} to be a Tensor",
-                )
-                self.assertTrue(
-                    torch.equal(original_value, reconstructed_value),
-                    f"Tensor attribute {attr_name} not equal after reconstruction",
-                )
-
-        # Verify all context attributes are equal (using _CONTEXT_ATTRS)
-        for attr_name in BlockMask._CONTEXT_ATTRS:
-            original_value = getattr(block_mask, attr_name)
-            reconstructed_value = getattr(reconstructed_mask, attr_name)
-
-            self.assertEqual(
-                original_value,
-                reconstructed_value,
-                f"Context attribute {attr_name} not equal after reconstruction",
-            )
-
-    @supported_platform
-    def test_pytree_flatten_with_keys(self, device):
-        """Test that BlockMask._flatten_with_keys works correctly for tracing."""
-
-        def causal_mask(b, h, q_idx, kv_idx):
-            return q_idx >= kv_idx
-
-        block_mask = create_block_mask(
-            causal_mask, B=2, H=4, Q_LEN=512, KV_LEN=512, device=device
-        )
-
-        tensors_with_keys, context_with_keys = block_mask._flatten_with_keys()
-
-        self.assertEqual(len(tensors_with_keys), len(BlockMask._TENSOR_ATTRS))
-        self.assertEqual(len(context_with_keys), len(BlockMask._CONTEXT_ATTRS))
-
-        from torch.utils._pytree import GetAttrKey
-
-        for key, tensor in tensors_with_keys:
-            self.assertIsInstance(key, GetAttrKey)
-            self.assertIsNotNone(key)
-
-        for key, value in context_with_keys:
-            self.assertIsInstance(key, GetAttrKey)
-            self.assertIsNotNone(key)
-
-    @supported_platform
-    def test_pytree_preserves_new_attributes(self, device):
-        """
-        Test that BlockMask._TENSOR_ATTRS and _CONTEXT_ATTRS are correctly defined
-        and that flatten/unflatten preserves all attributes in these lists.
-
-        """
-
-        def causal_mask(b, h, q_idx, kv_idx):
-            return q_idx >= kv_idx
-
-        block_mask = create_block_mask(
-            causal_mask, B=2, H=4, Q_LEN=512, KV_LEN=512, device=device
-        )
-
-        # Flatten and unflatten using class methods
-        tensors, context = block_mask._flatten()
-        reconstructed_mask = BlockMask._unflatten(tensors, context)
-
-        # Verify the number of tensors and context values matches the attribute lists
-        self.assertEqual(
-            len(tensors),
-            len(BlockMask._TENSOR_ATTRS),
-            "Number of tensors should match _TENSOR_ATTRS length",
-        )
-        self.assertEqual(
-            len(context),
-            len(BlockMask._CONTEXT_ATTRS),
-            "Number of context values should match _CONTEXT_ATTRS length",
-        )
-
-        # Verify all attributes from the lists exist and are equal after reconstruction
-        for attr_name in BlockMask._TENSOR_ATTRS + BlockMask._CONTEXT_ATTRS:
-            self.assertTrue(
-                hasattr(reconstructed_mask, attr_name),
-                f"Reconstructed mask missing attribute: {attr_name}",
-            )
-            original_value = getattr(block_mask, attr_name)
-            reconstructed_value = getattr(reconstructed_mask, attr_name)
-
-            if isinstance(original_value, torch.Tensor):
-                self.assertTrue(
-                    torch.equal(original_value, reconstructed_value),
-                    f"Tensor attribute {attr_name} not equal after reconstruction",
-                )
-            elif original_value is None:
-                self.assertIsNone(
-                    reconstructed_value,
-                    f"Attribute {attr_name} should be None but got {reconstructed_value}",
-                )
-            else:
-                self.assertEqual(
-                    original_value,
-                    reconstructed_value,
-                    f"Attribute {attr_name} not equal after reconstruction",
-                )
-
 
 @large_tensor_test_class("2GB", device=test_device[0])
 class TestPagedAttention(InductorTestCase):
@@ -6753,7 +6618,7 @@ class TestLearnableBiases(InductorTestCase):
         )
         # Error in backwards
         with self.assertRaisesRegex(
-            torch._inductor.exc.LoweringException,
+            torch._inductor.exc.InductorError,
             "Using multiple indexing operations on the same tensor that requires gradients",
         ):
             self._check_outputs_and_grads(

test/inductor/test_torchinductor_opinfo.py

@@ -294,7 +294,6 @@ inductor_expected_failures_single_sample["xpu"] = {
         i32,
         i64,
     },  # align with cuda.
-    "linalg.eig": {f32, f64},
     ("linalg.pinv", "singular"): {f64},
     # could not create a primitive
     "addmv": {f64},

test/jit/test_concrete_module_type.py

@@ -0,0 +1,53 @@
+# Owner(s): ["oncall: jit"]
+
+import unittest
+
+import torch
+from torch.testing._internal.common_utils import raise_on_run_directly
+
+
+class TestConcreteModuleTypeFindSubmodule(unittest.TestCase):
+    def test_error_message_includes_submodule_name(self):
+        class ChildModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(5, 3)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        class ParentModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.existing_child = ChildModule()
+
+            def forward(self, x):
+                return self.existing_child(x)
+
+        module = ParentModule()
+        scripted_module = torch.jit.script(module)
+
+        self.assertIsNotNone(scripted_module.existing_child)
+
+        # Now try to trigger the error by accessing a non-existent submodule
+        # through the internal ConcreteModuleType mechanism. This happens
+        # when the TorchScript compiler tries to resolve submodule references.
+        class BrokenModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.parent = ParentModule()
+
+            def forward(self, x):
+                return self.parent.missing_submodule(x)
+
+        broken_module = BrokenModule()
+
+        with self.assertRaises(RuntimeError) as context:
+            torch.jit.script(broken_module)
+
+        error_msg = str(context.exception)
+        self.assertIn("missing_submodule", error_msg.lower())
+
+
+if __name__ == "__main__":
+    raise_on_run_directly("test/test_jit.py")

test/xpu/test_gemm.py

@@ -238,7 +238,7 @@ class TestBasicGEMM(TestCase):
                 )
 
     @precisionOverride({torch.float: 1e-4, torch.double: 1e-6, torch.half: 1e-1})
-    @dtypes(torch.float32, torch.half, torch.double)
+    @dtypes(torch.float32, torch.half, torch.double, torch.complex64)
     @tf32_on_and_off(0.05)
     def test_addmm(self, device, dtype):
         self._test_addmm_impl(torch.addmm, None, device, dtype)
@@ -313,6 +313,7 @@ class TestBasicGEMM(TestCase):
         torch.half,
         torch.float32,
         torch.float64,
+        torch.complex64,
     )
     @tf32_on_and_off(0.05)
     def test_mm(self, device, dtype):
@@ -416,7 +417,7 @@ class TestBasicGEMM(TestCase):
             _test_mm(n, m, p, dtype, genf)
 
     @precisionOverride({torch.half: 0.05, torch.bfloat16: 0.05})
-    @dtypes(torch.float32, torch.bfloat16, torch.half, torch.float64)
+    @dtypes(torch.float32, torch.bfloat16, torch.half, torch.float64, torch.complex64)
     @tf32_on_and_off(0.05)
     def test_bmm(self, device, dtype):
         batch_sizes = [1, 10]
@@ -533,7 +534,7 @@ class TestBasicGEMM(TestCase):
         self.assertEqual(res7, ref)
 
     @precisionOverride({torch.half: 0.05, torch.bfloat16: 0.05})
-    @dtypes(torch.float64, torch.float32, torch.bfloat16, torch.half)
+    @dtypes(torch.float64, torch.float32, torch.bfloat16, torch.half, torch.complex64)
     @tf32_on_and_off(0.005)
     def test_addbmm(self, device, dtype):
         num_batches = 2
@@ -637,7 +638,7 @@ class TestBasicGEMM(TestCase):
             self._test_addbmm_baddbmm("addbmm", b1, b2, ref, out_tensor)
 
     @precisionOverride({torch.half: 0.1, torch.bfloat16: 0.5, torch.float64: 1e-6})
-    @dtypes(torch.float64, torch.float32, torch.bfloat16, torch.half)
+    @dtypes(torch.float64, torch.float32, torch.bfloat16, torch.half, torch.complex64)
     @tf32_on_and_off(0.01)
     def test_baddbmm(self, device, dtype):
         num_batches = 10

third_party/xpu.txt

@@ -1 +1 @@
-ce9db15136c5e8ba1b51710aae574ce4791c5d73
+30dcaa83183b189b00dc827fb39234714fe4e46d

tools/autograd/gen_autograd_functions.py

@@ -863,7 +863,7 @@ static PyObject* THP${op}_${name}_getter(THPCppFunction *self, void *_unused) {
             saved_variables.append(f"{type.cpp_type()} {name};")
 
             if type in MISC_GETTER_DEFS:
-                # pyrefly: ignore  # index-error
+                # pyrefly: ignore [index-error]
                 getter_def, body = MISC_GETTER_DEFS[type]
                 getter_definitions.append(
                     getter_def.substitute(op=info.op, name=name, body=body)
@@ -1034,7 +1034,7 @@ static PyObject* THP${op}_${name}_getter(THPCppFunction *self, void *_unused) {
     unpack_ivalues = []
     for typ, name in zip(apply_functional_args_ref_types, apply_functional_args):
         typ = typ.removesuffix("&")
-        # pyrefly: ignore  # bad-argument-type
+        # pyrefly: ignore [bad-argument-type]
         unpack_ivalues.append(f"auto {name} = packed_args.unpack<{typ}>();")
 
     schema_args = [f"std::array<bool, {len(input_name_to_idx)}>"]

tools/autograd/gen_trace_type.py

@@ -182,7 +182,7 @@ def format_trace_inputs(f: NativeFunction) -> str:
             ADD_TRACE_INPUT.substitute(
                 name=f.func.arguments.out[i].name, input=f.func.arguments.out[i].name
             )
-            # pyrefly: ignore  # unbound-name
+            # pyrefly: ignore [unbound-name]
             for i in range(num_out_args)
         ]
 

tools/autograd/gen_variable_type.py

@@ -1495,7 +1495,7 @@ def emit_body(
                 else:
                     expr = f"SavedVariable({var}, {str(is_output).lower()})"
                     if foreacharg is not None and "original_selfs" not in expr:
-                        # pyrefly: ignore  # unbound-name
+                        # pyrefly: ignore [unbound-name]
                         expr = expr.replace(src_name, name_in_expr)
             elif (
                 type == BaseCType(tensorListT)
@@ -1845,14 +1845,14 @@ def emit_body(
                                 )
                             )
                         cur_derivative_conditions.append(
-                            # pyrefly: ignore  # bad-argument-type
+                            # pyrefly: ignore [bad-argument-type]
                             FW_DERIVATIVE_CHECK_TEMPLATE.substitute(
                                 req_inp=inp_name + "[i]"
                             )
                         )
                     else:
                         cur_derivative_conditions.append(
-                            # pyrefly: ignore  # bad-argument-type
+                            # pyrefly: ignore [bad-argument-type]
                             FW_DERIVATIVE_CHECK_TEMPLATE.substitute(req_inp=inp_name)
                         )
 
@@ -1923,7 +1923,7 @@ def emit_body(
                 unpacked_arguments += FW_DERIVATIVE_DEFINED_GRAD_TEMPLATE.substitute(
                     inp_name="original_self",
                     inp="original_self" + input_suffix,
-                    # pyrefly: ignore  # unbound-name
+                    # pyrefly: ignore [unbound-name]
                     zeros_fn=zeros_fn,
                 )
                 unpacked_arguments += FW_DERIVATIVE_DEFINED_PRIMAL_TEMPLATE.substitute(

tools/autograd/load_derivatives.py

@@ -95,11 +95,11 @@ def add_view_copy_derivatives(
             else:
                 break
         # prefer manually-defined derivatives if any
-        # pyrefly: ignore  # unbound-name
+        # pyrefly: ignore [unbound-name]
         if len(view_copy_differentiability_infos) > 0 and fn_schema not in infos:
-            # pyrefly: ignore  # unbound-name
+            # pyrefly: ignore [unbound-name]
             assert fn_schema is not None
-            # pyrefly: ignore  # unbound-name
+            # pyrefly: ignore [unbound-name]
             view_infos[fn_schema] = view_copy_differentiability_infos
 
     infos.update(view_infos)
@@ -401,7 +401,7 @@ def postprocess_forward_derivatives(
             for arg_name in all_arg_names:
                 if arg_name in diff_arg_names:
                     arg_name = arg_name + "_t"
-                # pyrefly: ignore  # bad-argument-type
+                # pyrefly: ignore [bad-argument-type]
                 new_args.append(arg_name)
 
             # TODO we are trolling
@@ -942,7 +942,7 @@ def saved_variables(
             + f".sym_strides(), which returned a c10::SymIntArrayRef. formula={formula}"
         )
     for nctype in nctypes:
-        # pyrefly: ignore  # bad-assignment
+        # pyrefly: ignore [bad-assignment]
         name = (
             nctype.name.name if isinstance(nctype.name, SpecialArgName) else nctype.name
         )
@@ -952,7 +952,7 @@ def saved_variables(
 
             def repl(m: re.Match[str]) -> str:
                 suffix: str = (
-                    # pyrefly: ignore  # bad-assignment
+                    # pyrefly: ignore [bad-assignment]
                     info["suffix"](m) if callable(info["suffix"]) else info["suffix"]
                 )
                 expr: str = info["expr"](name) if "expr" in info else m.group(0)

tools/code_coverage/package/tool/summarize_jsons.py

@@ -67,7 +67,7 @@ def is_intrested_file(
 
     # ignore files that are not belong to pytorch
     if platform == TestPlatform.OSS:
-        # pyrefly: ignore  # import-error
+        # pyrefly: ignore [import-error]
         from package.oss.utils import get_pytorch_folder
 
         if not file_path.startswith(get_pytorch_folder()):

tools/download_mnist.py

@@ -24,7 +24,7 @@ def report_download_progress(
     file_size: int,
 ) -> None:
     if file_size != -1:
-        # pyrefly: ignore  # no-matching-overload
+        # pyrefly: ignore [no-matching-overload]
         percent = min(1, (chunk_number * chunk_size) / file_size)
         bar = "#" * int(64 * percent)
         sys.stdout.write(f"\r0% |{bar:<64}| {int(percent * 100)}%")

tools/dynamo/gb_id_mapping.py

@@ -105,10 +105,10 @@ def extract_info_from_keyword(source: str, kw: ast.keyword) -> Any:
         evaluated_context = []
         for value in kw.value.values:
             if isinstance(value, ast.FormattedValue):
-                # pyrefly: ignore  # bad-argument-type
+                # pyrefly: ignore [bad-argument-type]
                 evaluated_context.append(f"{{{ast.unparse(value.value)}}}")
             elif isinstance(value, ast.Constant):
-                # pyrefly: ignore  # bad-argument-type
+                # pyrefly: ignore [bad-argument-type]
                 evaluated_context.append(value.value)
         return "".join(evaluated_context)
     else:
@@ -154,7 +154,7 @@ def find_unimplemented_v2_calls(
 
                         for kw in node.keywords:
                             if kw.arg in info:
-                                # pyrefly: ignore  # unsupported-operation
+                                # pyrefly: ignore [unsupported-operation]
                                 info[kw.arg] = extract_info_from_keyword(source, kw)
 
                         if info["gb_type"] is None:

tools/experimental/torchfuzz/codegen.py

@@ -219,7 +219,6 @@ class DefaultFuzzTemplate(FuzzTemplate):
                 # Neural network operations
                 "torch.nn.functional.embedding",
                 "torch.nn.functional.linear",
-                "torch.nn.functional.scaled_dot_product_attention",
                 # Activation functions
                 "torch.nn.functional.relu",
                 "torch.nn.functional.leaky_relu",

tools/experimental/torchfuzz/multi_process_fuzzer.py

@@ -72,21 +72,6 @@ IGNORE_PATTERNS: list[re.Pattern] = [
     re.compile(
         r"dimensionality of sizes \(0\) must match dimensionality of strides \(1\)"
     ),  # https://github.com/pytorch/pytorch/issues/164814
-    re.compile(
-        r"self and mat2 must have the same dtype"
-    ),  # https://github.com/pytorch/pytorch/issues/165718
-    re.compile(
-        r"free\(\): invalid next size \(fast\)"
-    ),  # TODO: figure out why sometimes heap metadata gets corrupted on program exit (checks actually pass successfully)
-    re.compile(
-        r'assert "int" in str\(indices\.get_dtype\(\)\)'
-    ),  # https://github.com/pytorch/pytorch/issues/166042
-    re.compile(
-        r'self\.shape_env\.guard_or_defer_runtime_assert\(expr, "guard_equals"\)'
-    ),  # https://github.com/pytorch/pytorch/issues/166245
-    re.compile(
-        r"assert len\(self\.stride\) == len\(order\)"
-    ),  # https://github.com/pytorch/pytorch/issues/166270
     # Add more patterns here as needed, e.g.:
     # re.compile(r"Some other error message"),
 ]
@@ -311,7 +296,7 @@ def run_multi_process_fuzzer(
                 )
 
                 def write_func(msg):
-                    # pyrefly: ignore  # missing-attribute
+                    # pyrefly: ignore [missing-attribute]
                     pbar.write(msg)
             else:
                 persist_print("Progress: (install tqdm for better progress bar)")
@@ -606,7 +591,7 @@ def run_until_failure(
                     pbar.set_postfix_str(f"{total_successful}/{total_ignored}")
 
                     def write_func(msg):
-                        # pyrefly: ignore  # missing-attribute
+                        # pyrefly: ignore [missing-attribute]
                         pbar.write(msg)
                 else:
                     pbar = None

tools/experimental/torchfuzz/operators/__init__.py

@@ -23,7 +23,6 @@ from torchfuzz.operators.nn_functional import (
     LayerNormOperator,
     LinearOperator,
     ReLUOperator,
-    ScaledDotProductAttentionOperator,
     SoftmaxOperator,
 )
 from torchfuzz.operators.registry import (
@@ -78,7 +77,6 @@ __all__ = [
     "EmbeddingOperator",
     "LinearOperator",
     "ReLUOperator",
-    "ScaledDotProductAttentionOperator",
     "SoftmaxOperator",
     "DropoutOperator",
     "LayerNormOperator",

tools/experimental/torchfuzz/operators/constant.py

@@ -111,7 +111,7 @@ class ConstantOperator(Operator):
                 ]:
                     # Clamp integer values to [0, 3] to avoid index overflow in multiplication
                     # Even with multiplication, indices should stay in reasonable range
-                    # pyrefly: ignore  # bad-argument-type
+                    # pyrefly: ignore [bad-argument-type]
                     fill_value = max(0, min(3, abs(fill_value)))
 
                 tensor_creation = (

tools/experimental/torchfuzz/operators/matrix_multiply.py

@@ -15,8 +15,14 @@ from torchfuzz.tensor_fuzzer import Spec, TensorSpec
 class MatrixMultiplyOperator(Operator):
     """Base class for matrix multiplication operations."""
 
-    def __init__(self, name: str):
+    def __init__(self, name: str, torch_op: str):
         super().__init__(name)
+        self._torch_op = torch_op
+
+    @property
+    def torch_op_name(self) -> Optional[str]:
+        """Return the torch operation name."""
+        return self._torch_op
 
     def can_produce(self, output_spec: Spec) -> bool:
         """Matrix multiply operations can produce float/complex tensors of dimension >= 2."""
@@ -41,6 +47,12 @@ class MatrixMultiplyOperator(Operator):
 
     def _get_compatible_dtype(self, output_dtype):
         """Get a compatible dtype for matrix multiplication."""
+        # For matrix multiplication, we need to be flexible with input dtypes
+        # since earlier operations may have performed type promotion.
+        # We'll let the fuzzer generate whatever dtypes result from earlier operations
+        # and rely on the operation graph to ensure compatibility.
+        # Return the output dtype as a starting point, but this may be overridden
+        # by the actual tensor specs generated by the fuzzer.
         return [output_dtype, output_dtype]
 
 
@@ -48,14 +60,9 @@ class MMOperator(MatrixMultiplyOperator):
     """Operator for matrix multiplication (torch.mm)."""
 
     def __init__(self):
-        super().__init__("mm")
+        super().__init__("mm", "torch.mm")
         self.weight = 5.0
 
-    @property
-    def torch_op_name(self) -> Optional[str]:
-        """Return the torch operation name."""
-        return "torch.mm"
-
     def can_produce(self, output_spec: Spec) -> bool:
         """MM requires exactly 2D tensors."""
         if not isinstance(output_spec, TensorSpec):
@@ -89,6 +96,7 @@ class MMOperator(MatrixMultiplyOperator):
         # Choose a random inner dimension k
         k = random.randint(1, 16)
 
+        # Get compatible dtypes
         dtypes = self._get_compatible_dtype(output_spec.dtype)
 
         # First tensor: [m, k]
@@ -133,14 +141,9 @@ class AddmmOperator(MatrixMultiplyOperator):
     """Operator for additive matrix multiplication (torch.addmm)."""
 
     def __init__(self):
-        super().__init__("addmm")
+        super().__init__("addmm", "torch.addmm")
         self.weight = 5.0
 
-    @property
-    def torch_op_name(self) -> Optional[str]:
-        """Return the torch operation name."""
-        return "torch.addmm"
-
     def can_produce(self, output_spec: Spec) -> bool:
         """Addmm requires exactly 2D tensors."""
         if not isinstance(output_spec, TensorSpec):
@@ -174,6 +177,7 @@ class AddmmOperator(MatrixMultiplyOperator):
         # Choose a random inner dimension k
         k = random.randint(1, 16)
 
+        # Get compatible dtypes
         dtypes = self._get_compatible_dtype(output_spec.dtype)
 
         # Bias tensor: [m, n] (same shape as output)
@@ -226,14 +230,9 @@ class BmmOperator(MatrixMultiplyOperator):
     """Operator for batch matrix multiplication (torch.bmm)."""
 
     def __init__(self):
-        super().__init__("bmm")
+        super().__init__("bmm", "torch.bmm")
         self.weight = 5.0
 
-    @property
-    def torch_op_name(self) -> Optional[str]:
-        """Return the torch operation name."""
-        return "torch.bmm"
-
     def can_produce(self, output_spec: Spec) -> bool:
         """Batch matrix multiply requires 3D tensors."""
         if not isinstance(output_spec, TensorSpec):
@@ -267,6 +266,7 @@ class BmmOperator(MatrixMultiplyOperator):
         # Choose a random inner dimension k
         k = random.randint(1, 16)
 
+        # Get compatible dtypes
         dtypes = self._get_compatible_dtype(output_spec.dtype)
 
         # First tensor: [b, m, k]
@@ -311,14 +311,9 @@ class MatmulOperator(MatrixMultiplyOperator):
     """Operator for general matrix multiplication (torch.matmul)."""
 
     def __init__(self):
-        super().__init__("matmul")
+        super().__init__("matmul", "torch.matmul")
         self.weight = 500.0
 
-    @property
-    def torch_op_name(self) -> Optional[str]:
-        """Return the torch operation name."""
-        return "torch.matmul"
-
     def can_produce(self, output_spec: Spec) -> bool:
         """Matmul can handle various tensor dimensions >= 1."""
         if not isinstance(output_spec, TensorSpec):
@@ -348,6 +343,7 @@ class MatmulOperator(MatrixMultiplyOperator):
         output_size = output_spec.size
         output_dims = len(output_size)
 
+        # Get compatible dtypes
         dtypes = self._get_compatible_dtype(output_spec.dtype)
 
         if output_dims == 1:

tools/experimental/torchfuzz/operators/nn_functional.py

@@ -1,6 +1,5 @@
 """Neural network functional operator implementations."""
 
-import math
 import random
 from typing import Optional
 
@@ -753,17 +752,6 @@ class GroupNormOperator(Operator):
         # GroupNorm needs at least 2 dimensions (batch, channels)
         if len(output_spec.size) < 2:
             return False
-
-        # GroupNorm requires more than 1 value per channel
-        # For shape (N, C, *), num_values_per_channel = N * prod(*)
-        # We need N * prod(*) > 1
-        batch_size = output_spec.size[0]
-        spatial_size = math.prod(output_spec.size[2:])
-        num_values_per_channel = batch_size * spatial_size
-
-        if num_values_per_channel <= 1:
-            return False
-
         return is_float_dtype(output_spec.dtype)
 
     def fuzz_inputs_specs(self, output_spec: Spec) -> list[Spec]:
@@ -963,221 +951,3 @@ class SiLUOperator(Operator):
 
         input_name = input_names[0]
         return f"{output_name} = torch.nn.functional.silu({input_name})"
-
-
-class ScaledDotProductAttentionOperator(Operator):
-    """Operator for torch.nn.functional.scaled_dot_product_attention."""
-
-    def __init__(self):
-        super().__init__("torch.nn.functional.scaled_dot_product_attention")
-
-    @property
-    def torch_op_name(self) -> Optional[str]:
-        """Return the torch operation name."""
-        return "torch.nn.functional.scaled_dot_product_attention"
-
-    def can_produce(self, output_spec: Spec) -> bool:
-        """Scaled dot product attention can produce tensor outputs with floating point dtypes."""
-        if not isinstance(output_spec, TensorSpec):
-            return False
-        # SDPA needs at least 3 dimensions (batch, seq_len, embed_dim)
-        if len(output_spec.size) < 3:
-            return False
-        return is_float_dtype(output_spec.dtype)
-
-    def fuzz_inputs_specs(self, output_spec: Spec) -> list[Spec]:
-        """Generate input specs for scaled_dot_product_attention.
-
-        SDPA requires:
-        - query: (batch, seq_len, embed_dim) or (batch, num_heads, seq_len, head_dim)
-        - key: (batch, seq_len, embed_dim) or (batch, num_heads, seq_len_kv, head_dim)
-        - value: (batch, seq_len, embed_dim) or (batch, num_heads, seq_len_kv, head_dim)
-        Output shape matches query shape.
-        """
-        if not isinstance(output_spec, TensorSpec):
-            raise ValueError(
-                "ScaledDotProductAttentionOperator can only produce TensorSpec outputs"
-            )
-
-        if len(output_spec.size) < 3:
-            raise ValueError("SDPA output must have at least 3 dimensions")
-
-        # Query has the same shape as output
-        query_spec = TensorSpec(
-            size=output_spec.size, stride=output_spec.stride, dtype=output_spec.dtype
-        )
-
-        # Key and value: match query shape for simplicity
-        # In practice, seq_len for key/value can differ, but we'll keep it simple
-        key_spec = TensorSpec(
-            size=output_spec.size, stride=output_spec.stride, dtype=output_spec.dtype
-        )
-        value_spec = TensorSpec(
-            size=output_spec.size, stride=output_spec.stride, dtype=output_spec.dtype
-        )
-
-        return [query_spec, key_spec, value_spec]
-
-    def codegen(
-        self, output_name: str, input_names: list[str], output_spec: Spec
-    ) -> str:
-        """Generate code for scaled_dot_product_attention operation."""
-        if len(input_names) != 3:
-            raise ValueError("SDPA requires exactly 3 inputs: query, key, value")
-
-        # Ensure dtype compatibility by converting all inputs to the expected output dtype
-        target_dtype = str(output_spec.dtype)
-        query_name, key_name, value_name = input_names
-        return f"{output_name} = torch.nn.functional.scaled_dot_product_attention({query_name}.to({target_dtype}), {key_name}.to({target_dtype}), {value_name}.to({target_dtype}))"
-
-
-class MultiHeadAttentionForwardOperator(Operator):
-    """Operator for torch.nn.functional.multi_head_attention_forward."""
-
-    def __init__(self):
-        super().__init__("torch.nn.functional.multi_head_attention_forward")
-
-    @property
-    def torch_op_name(self) -> Optional[str]:
-        """Return the torch operation name."""
-        return "torch.nn.functional.multi_head_attention_forward"
-
-    def can_produce(self, output_spec: Spec) -> bool:
-        """Multi-head attention forward can produce tensor outputs with floating point dtypes."""
-        if not isinstance(output_spec, TensorSpec):
-            return False
-        # MHA needs at least 3 dimensions (seq_len, batch, embed_dim)
-        if len(output_spec.size) < 3:
-            return False
-        # MHA cannot handle 0-sized dimensions (seq_len, batch, or embed_dim must be > 0)
-        if any(dim == 0 for dim in output_spec.size):
-            return False
-        return is_float_dtype(output_spec.dtype)
-
-    def fuzz_inputs_specs(self, output_spec: Spec) -> list[Spec]:
-        """Generate input specs for multi_head_attention_forward.
-
-        MHA requires:
-        - query, key, value: (seq_len, batch, embed_dim)
-        - in_proj_weight: (3*embed_dim, embed_dim) for combined QKV projection
-        - in_proj_bias: (3*embed_dim,) optional
-        - out_proj_weight: (embed_dim, embed_dim)
-        - out_proj_bias: (embed_dim,) optional
-
-        For simplicity, we'll use the combined in_proj_weight path.
-
-        IMPORTANT: The order of optional parameters matters for codegen!
-        We must ensure that when we have 6 inputs, they are in the order:
-        query, key, value, in_proj_weight, in_proj_bias, out_proj_weight
-        NOT: query, key, value, in_proj_weight, out_proj_weight, out_proj_bias
-        """
-        if not isinstance(output_spec, TensorSpec):
-            raise ValueError(
-                "MultiHeadAttentionForwardOperator can only produce TensorSpec outputs"
-            )
-
-        if len(output_spec.size) < 3:
-            raise ValueError("MHA output must have at least 3 dimensions")
-
-        # Output shape: (seq_len, batch, embed_dim)
-        seq_len, batch, embed_dim = output_spec.size[:3]
-
-        # Query, key, value have the same shape as output
-        query_spec = TensorSpec(
-            size=output_spec.size, stride=output_spec.stride, dtype=output_spec.dtype
-        )
-        key_spec = TensorSpec(
-            size=output_spec.size, stride=output_spec.stride, dtype=output_spec.dtype
-        )
-        value_spec = TensorSpec(
-            size=output_spec.size, stride=output_spec.stride, dtype=output_spec.dtype
-        )
-
-        # in_proj_weight: (3*embed_dim, embed_dim)
-        in_proj_weight_spec = TensorSpec(
-            size=(3 * embed_dim, embed_dim),
-            stride=(embed_dim, 1),
-            dtype=output_spec.dtype,
-        )
-
-        # out_proj_weight: (embed_dim, embed_dim)
-        out_proj_weight_spec = TensorSpec(
-            size=(embed_dim, embed_dim),
-            stride=(embed_dim, 1),
-            dtype=output_spec.dtype,
-        )
-
-        # For simplicity and correctness, always generate all required tensors
-        # This avoids ambiguity in the codegen about which optional parameters are present
-        # We'll use a simplified signature: query, key, value, in_proj_weight, out_proj_weight only
-        specs = [
-            query_spec,
-            key_spec,
-            value_spec,
-            in_proj_weight_spec,
-            out_proj_weight_spec,
-        ]
-
-        from typing import cast
-
-        return cast(list[Spec], specs)
-
-    def _calculate_stride(self, size):
-        """Calculate stride for a given size."""
-        if not size:
-            return ()
-        stride = []
-        current_stride = 1
-        for dim_size in reversed(size):
-            stride.append(current_stride)
-            current_stride *= dim_size
-        return tuple(reversed(stride))
-
-    def codegen(
-        self, output_name: str, input_names: list[str], output_spec: Spec
-    ) -> str:
-        """Generate code for multi_head_attention_forward operation."""
-        if len(input_names) != 5:
-            raise ValueError(
-                "MHA requires exactly 5 inputs: query, key, value, in_proj_weight, out_proj_weight"
-            )
-
-        if not isinstance(output_spec, TensorSpec):
-            raise ValueError(
-                "MultiHeadAttentionForwardOperator can only produce TensorSpec outputs"
-            )
-
-        target_dtype = str(output_spec.dtype)
-        embed_dim = output_spec.size[-1]
-
-        # Determine number of heads (must divide embed_dim evenly)
-        # Common choices: 8, 4, 2, 1
-        possible_heads = [h for h in [8, 4, 2, 1] if embed_dim % h == 0]
-        num_heads = possible_heads[0] if possible_heads else 1
-
-        query_name = input_names[0]
-        key_name = input_names[1]
-        value_name = input_names[2]
-        in_proj_weight_name = input_names[3]
-        out_proj_weight_name = input_names[4]
-
-        # Build the function call without optional biases
-        code = f"""{output_name}, _ = torch.nn.functional.multi_head_attention_forward(
-    {query_name}.to({target_dtype}),
-    {key_name}.to({target_dtype}),
-    {value_name}.to({target_dtype}),
-    {embed_dim},
-    {num_heads},
-    {in_proj_weight_name}.to({target_dtype}),
-    None,  # in_proj_bias
-    None,  # bias_k
-    None,  # bias_v
-    False,  # add_zero_attn
-    0.0,  # dropout_p (no dropout for testing)
-    {out_proj_weight_name}.to({target_dtype}),
-    None,  # out_proj_bias
-    training=False,  # Use eval mode for deterministic behavior
-    need_weights=False,  # Don't compute attention weights for performance
-)"""
-
-        return code

tools/experimental/torchfuzz/operators/registry.py

@@ -32,7 +32,6 @@ from torchfuzz.operators.nn_functional import (
     LinearOperator,
     ReLUOperator,
     RMSNormOperator,
-    ScaledDotProductAttentionOperator,
     SigmoidOperator,
     SiLUOperator,
     SoftmaxOperator,
@@ -102,7 +101,6 @@ class OperatorRegistry:
         # Neural network functional operators
         self.register(EmbeddingOperator())
         self.register(LinearOperator())
-        self.register(ScaledDotProductAttentionOperator())
 
         # Activation functions
         self.register(ReLUOperator())

tools/experimental/torchfuzz/operators/tensor_pointwise.py

@@ -1,6 +1,7 @@
 """Tensor pointwise operator implementation."""
 
 import random
+from typing import Optional
 
 import torch
 
@@ -16,10 +17,16 @@ from torchfuzz.type_promotion import (
 class PointwiseOperator(Operator):
     """Base class for element-wise pointwise operations."""
 
-    def __init__(self, name: str, symbol: str):
+    def __init__(self, name: str, torch_op: str, symbol: str):
         super().__init__(name)
+        self._torch_op = torch_op
         self.symbol = symbol
 
+    @property
+    def torch_op_name(self) -> Optional[str]:
+        """Return the torch operation name."""
+        return self._torch_op
+
     def can_produce(self, output_spec: Spec) -> bool:
         """Tensor pointwise operations can produce tensors but not scalars."""
         if isinstance(output_spec, TensorSpec) and output_spec.dtype == torch.bool:
@@ -67,7 +74,9 @@ class PointwiseOperator(Operator):
     ) -> str:
         """Generate code for pointwise operation."""
         if len(input_names) == 2:
-            return f"{output_name} = {self.torch_op_name}({input_names[0]}, {input_names[1]})"
+            return (
+                f"{output_name} = {self._torch_op}({input_names[0]}, {input_names[1]})"
+            )
         else:
             # Chain operations using symbols for readability
             expr = f" {self.symbol} ".join(input_names)
@@ -78,42 +87,26 @@ class AddOperator(PointwiseOperator):
     """Operator for element-wise addition."""
 
     def __init__(self, weight: float = 1.0):
-        super().__init__("add", "+")
+        super().__init__("add", "torch.add", "+")
         self.weight = float(weight)
 
-    @property
-    def torch_op_name(self) -> str:
-        return "torch.add"
-
 
 class MulOperator(PointwiseOperator):
     """Operator for element-wise multiplication."""
 
     def __init__(self):
-        super().__init__("mul", "*")
-
-    @property
-    def torch_op_name(self) -> str:
-        return "torch.mul"
+        super().__init__("mul", "torch.mul", "*")
 
 
 class SubOperator(PointwiseOperator):
     """Operator for element-wise subtraction."""
 
     def __init__(self):
-        super().__init__("sub", "-")
-
-    @property
-    def torch_op_name(self) -> str:
-        return "torch.sub"
+        super().__init__("sub", "torch.sub", "-")
 
 
 class DivOperator(PointwiseOperator):
     """Operator for element-wise division."""
 
     def __init__(self):
-        super().__init__("div", "/")
-
-    @property
-    def torch_op_name(self) -> str:
-        return "torch.div"
+        super().__init__("div", "torch.div", "/")

tools/flight_recorder/components/config_manager.py

@@ -78,22 +78,22 @@ class JobConfig:
     def parse_args(
         self: "JobConfig", args: Optional[Sequence[str]]
     ) -> argparse.Namespace:
-        # pyrefly: ignore  # bad-assignment
+        # pyrefly: ignore [bad-assignment]
         args = self.parser.parse_args(args)
-        # pyrefly: ignore  # missing-attribute
+        # pyrefly: ignore [missing-attribute]
         if args.selected_ranks is not None:
-            # pyrefly: ignore  # missing-attribute
+            # pyrefly: ignore [missing-attribute]
             assert args.just_print_entries, (
                 "Not support selecting ranks without printing entries"
             )
-        # pyrefly: ignore  # missing-attribute
+        # pyrefly: ignore [missing-attribute]
         if args.pg_filters is not None:
-            # pyrefly: ignore  # missing-attribute
+            # pyrefly: ignore [missing-attribute]
             assert args.just_print_entries, (
                 "Not support selecting pg filters without printing entries"
             )
-        # pyrefly: ignore  # missing-attribute
+        # pyrefly: ignore [missing-attribute]
         if args.verbose:
             logger.set_log_level(logging.DEBUG)
-        # pyrefly: ignore  # bad-return
+        # pyrefly: ignore [bad-return]
         return args

tools/flight_recorder/components/utils.py

@@ -41,7 +41,7 @@ def format_frame(frame: dict[str, str]) -> str:
 def format_frames(frames: list[dict[str, str]]) -> str:
     formatted_frames = []
     for frame in frames:
-        # pyrefly: ignore  # bad-argument-type
+        # pyrefly: ignore [bad-argument-type]
         formatted_frames.append(format_frame(frame))
     return "\n".join(formatted_frames)
 
@@ -696,7 +696,7 @@ def check_version(version_by_ranks: dict[str, str], version: str) -> None:
 
 
 def get_version_detail(version: str) -> tuple[int, int]:
-    # pyrefly: ignore  # bad-assignment
+    # pyrefly: ignore [bad-assignment]
     version = version.split(".")
     assert len(version) == 2, f"Invalid version {version}"
     major, minor = map(int, version)

tools/flight_recorder/fr_trace.py

@@ -40,17 +40,17 @@ from tools.flight_recorder.components.types import types
 
 def main(args: Optional[Sequence[str]] = None) -> None:
     config = JobConfig()
-    # pyrefly: ignore  # bad-assignment
+    # pyrefly: ignore [bad-assignment]
     args = config.parse_args(args)
-    # pyrefly: ignore  # missing-attribute
+    # pyrefly: ignore [missing-attribute]
     assert args.trace_dir, "Trace directory trace_dir is required"
-    # pyrefly: ignore  # bad-argument-type
+    # pyrefly: ignore [bad-argument-type]
     details, version = read_dir(args)
-    # pyrefly: ignore  # bad-argument-type
+    # pyrefly: ignore [bad-argument-type]
     db = build_db(details, args, version)
-    # pyrefly: ignore  # missing-attribute
+    # pyrefly: ignore [missing-attribute]
     if args.output:
-        # pyrefly: ignore  # no-matching-overload
+        # pyrefly: ignore [no-matching-overload]
         with open(args.output, "wb") as f:
             pickle.dump((types, db), f)
 

tools/gdb/pytorch-gdb.py

@@ -34,7 +34,7 @@ class TensorRepr(gdb.Command):  # type: ignore[misc, no-any-unimported]
     on it.
     """
 
-    # pyrefly: ignore  # bad-argument-type
+    # pyrefly: ignore [bad-argument-type]
     __doc__ = textwrap.dedent(__doc__).strip()
 
     def __init__(self) -> None:

tools/gen_vulkan_spv.py

@@ -118,7 +118,7 @@ def extract_filename(path: str, keep_ext: bool = True) -> Any:
 
 
 # https://gist.github.com/pypt/94d747fe5180851196eb
-# pyrefly: ignore  # invalid-inheritance
+# pyrefly: ignore [invalid-inheritance]
 class UniqueKeyLoader(Loader):
     def construct_mapping(self, node, deep=False):  # type: ignore[no-untyped-def]
         if not isinstance(node, MappingNode):
@@ -234,7 +234,7 @@ def preprocess(
         last_indent = input_indent
 
     while blank_lines != 0:
-        # pyrefly: ignore  # unbound-name
+        # pyrefly: ignore [unbound-name]
         python_lines.append(python_indent + "print(file=OUT_STREAM)")
         blank_lines -= 1
 
@@ -669,7 +669,7 @@ def generateShaderDispatchStr(shader_info: ShaderInfo, name: str) -> str:
             "    ",
         )
 
-    # pyrefly: ignore  # unbound-name
+    # pyrefly: ignore [unbound-name]
     return shader_dispatch_str
 
 
@@ -684,18 +684,18 @@ def genCppFiles(
         name = getName(spvPath).replace("_spv", "")
 
         sizeBytes, spv_bin_str = generateSpvBinStr(spvPath, name)
-        # pyrefly: ignore  # bad-argument-type
+        # pyrefly: ignore [bad-argument-type]
         spv_bin_strs.append(spv_bin_str)
 
         shader_info = getShaderInfo(srcPath)
 
         register_shader_info_strs.append(
-            # pyrefly: ignore  # bad-argument-type
+            # pyrefly: ignore [bad-argument-type]
             generateShaderInfoStr(shader_info, name, sizeBytes)
         )
 
         if shader_info.register_for is not None:
-            # pyrefly: ignore  # bad-argument-type
+            # pyrefly: ignore [bad-argument-type]
             shader_registry_strs.append(generateShaderDispatchStr(shader_info, name))
 
     spv_bin_arrays = "\n".join(spv_bin_strs)

tools/jit/gen_unboxing.py

@@ -131,14 +131,14 @@ class ComputeCodegenUnboxedKernels:
                 else:
                     arg_cpp = f"c10::IValue({arg_default})"
             args_code.append(
-                # pyrefly: ignore  # bad-argument-type
+                # pyrefly: ignore [bad-argument-type]
                 f"""c10::Argument("{arg.name}", nullptr, ::std::nullopt, {arg_cpp})"""
             )
 
         returns = f.func.returns
         returns_code = []
         for ret in returns:
-            # pyrefly: ignore  # bad-argument-type
+            # pyrefly: ignore [bad-argument-type]
             returns_code.append(f"""c10::Argument("{ret.name if ret.name else ""}")""")
         return f"""
 // aten::{schema}

tools/linter/adapters/_linter/file_linter.py

@@ -112,7 +112,7 @@ class FileLinter:
         first_results = None
         original = replacement = pf.contents
 
-        # pyrefly: ignore  # bad-assignment
+        # pyrefly: ignore [bad-assignment]
         while True:
             try:
                 results = sorted(self._lint(pf), key=LintResult.sort_key)

tools/linter/adapters/_linter/sets.py

@@ -41,7 +41,7 @@ class LineWithSets:
         t = self.tokens[i]
         after = i < len(self.tokens) - 1 and self.tokens[i + 1]
         if t.string == "Set" and t.type == token.NAME:
-            # pyrefly: ignore  # bad-return
+            # pyrefly: ignore [bad-return]
             return after and after.string == "[" and after.type == token.OP
         return (
             (t.string == "set" and t.type == token.NAME)

tools/linter/adapters/clangtidy_linter.py

@@ -19,13 +19,13 @@ from typing import NamedTuple
 # PyTorch directory root
 def scm_root() -> str:
     path = os.path.abspath(os.getcwd())
-    # pyrefly: ignore  # bad-assignment
+    # pyrefly: ignore [bad-assignment]
     while True:
         if os.path.exists(os.path.join(path, ".git")):
             return path
         if os.path.isdir(os.path.join(path, ".hg")):
             return path
-        # pyrefly: ignore  # bad-argument-type
+        # pyrefly: ignore [bad-argument-type]
         n = len(path)
         path = os.path.dirname(path)
         if len(path) == n:

tools/linter/adapters/codespell_linter.py

@@ -20,6 +20,8 @@ FORBIDDEN_WORDS = {
     "multipy",  # project pytorch/multipy is dead  # codespell:ignore multipy
 }
 
+MAX_FILE_SIZE: int = 1024 * 1024 * 1024  # 1GB in bytes
+
 
 class LintSeverity(str, Enum):
     ERROR = "error"
@@ -86,6 +88,39 @@ def run_codespell(path: Path) -> str:
 
 def check_file(filename: str) -> list[LintMessage]:
     path = Path(filename).absolute()
+
+    # Check if file is too large
+    try:
+        file_size = os.path.getsize(path)
+        if file_size > MAX_FILE_SIZE:
+            return [
+                LintMessage(
+                    path=filename,
+                    line=None,
+                    char=None,
+                    code="CODESPELL",
+                    severity=LintSeverity.WARNING,
+                    name="file-too-large",
+                    original=None,
+                    replacement=None,
+                    description=f"File size ({file_size} bytes) exceeds {MAX_FILE_SIZE} bytes limit, skipping",
+                )
+            ]
+    except OSError as err:
+        return [
+            LintMessage(
+                path=filename,
+                line=None,
+                char=None,
+                code="CODESPELL",
+                severity=LintSeverity.ERROR,
+                name="file-access-error",
+                original=None,
+                replacement=None,
+                description=f"Failed to get file size: {err}",
+            )
+        ]
+
     try:
         run_codespell(path)
     except Exception as err:
@@ -101,7 +136,7 @@ def check_dictionary(filename: str) -> list[LintMessage]:
         words_set = set(words)
         if len(words) != len(words_set):
             raise ValueError("The dictionary file contains duplicate entries.")
-        # pyrefly: ignore  # no-matching-overload
+        # pyrefly: ignore [no-matching-overload]
         uncased_words = list(map(str.lower, words))
         if uncased_words != sorted(uncased_words):
             raise ValueError(

tools/linter/adapters/grep_linter.py

@@ -16,6 +16,11 @@ from typing import NamedTuple
 
 
 IS_WINDOWS: bool = os.name == "nt"
+MAX_FILE_SIZE: int = 1024 * 1024 * 1024  # 1GB in bytes
+MAX_MATCHES_PER_FILE: int = 100  # Maximum number of matches to report per file
+MAX_ORIGINAL_SIZE: int = (
+    512 * 1024
+)  # 512KB - don't compute replacement if original is larger
 
 
 class LintSeverity(str, Enum):
@@ -25,6 +30,10 @@ class LintSeverity(str, Enum):
     DISABLED = "disabled"
 
 
+LINTER_NAME: str = ""
+ERROR_DESCRIPTION: str | None = None
+
+
 class LintMessage(NamedTuple):
     path: str | None
     line: int | None
@@ -56,72 +65,143 @@ def run_command(
         logging.debug("took %dms", (end_time - start_time) * 1000)
 
 
+def print_lint_message(
+    name: str,
+    severity: LintSeverity = LintSeverity.ERROR,
+    path: str | None = None,
+    line: int | None = None,
+    original: str | None = None,
+    replacement: str | None = None,
+    description: str | None = None,
+) -> None:
+    """
+    Create a LintMessage and print it as JSON.
+
+    Accepts the same arguments as LintMessage constructor.
+    """
+    char = None
+    code = LINTER_NAME
+    description = description or ERROR_DESCRIPTION
+    lint_message = LintMessage(
+        path, line, char, code, severity, name, original, replacement, description
+    )
+    print(json.dumps(lint_message._asdict()), flush=True)
+
+
+def group_lines_by_file(lines: list[str]) -> dict[str, list[str]]:
+    """
+    Group matching lines by filename.
+
+    Args:
+        lines: List of grep output lines in format "filename:line:content"
+
+    Returns:
+        Dictionary mapping filename to list of line remainders (without filename prefix)
+    """
+    grouped: dict[str, list[str]] = {}
+    for line in lines:
+        if not line:
+            continue
+        # Extract filename and remainder from "filename:line:content" format
+        parts = line.split(":", 1)
+        filename = parts[0]
+        remainder = parts[1] if len(parts) > 1 else ""
+        if filename not in grouped:
+            grouped[filename] = []
+        grouped[filename].append(remainder)
+    return grouped
+
+
+def check_allowlist(
+    filename: str,
+    allowlist_pattern: str,
+) -> bool:
+    """
+    Check if a file matches the allowlist pattern.
+
+    Args:
+        filename: Path to the file to check
+        allowlist_pattern: Pattern to grep for in the file
+
+    Returns:
+        True if the file should be skipped (allowlist pattern matched), False otherwise.
+        Prints error message and returns False if there was an error running grep.
+    """
+    if not allowlist_pattern:
+        return False
+
+    try:
+        proc = run_command(["grep", "-nEHI", allowlist_pattern, filename])
+    except Exception as err:
+        print_lint_message(
+            name="command-failed",
+            description=(
+                f"Failed due to {err.__class__.__name__}:\n{err}"
+                if not isinstance(err, subprocess.CalledProcessError)
+                else (
+                    "COMMAND (exit code {returncode})\n"
+                    "{command}\n\n"
+                    "STDERR\n{stderr}\n\n"
+                    "STDOUT\n{stdout}"
+                ).format(
+                    returncode=err.returncode,
+                    command=" ".join(as_posix(x) for x in err.cmd),
+                    stderr=err.stderr.decode("utf-8").strip() or "(empty)",
+                    stdout=err.stdout.decode("utf-8").strip() or "(empty)",
+                )
+            ),
+        )
+        return False
+
+    # allowlist pattern was found, abort lint
+    if proc.returncode == 0:
+        return True
+
+    return False
+
+
 def lint_file(
-    matching_line: str,
+    filename: str,
+    line_remainders: list[str],
     allowlist_pattern: str,
     replace_pattern: str,
-    linter_name: str,
     error_name: str,
-    error_description: str,
-) -> LintMessage | None:
-    # matching_line looks like:
-    #   tools/linter/clangtidy_linter.py:13:import foo.bar.baz
-    split = matching_line.split(":")
-    filename = split[0]
+) -> None:
+    """
+    Lint a file with one or more pattern matches, printing LintMessages as they're created.
 
-    if allowlist_pattern:
-        try:
-            proc = run_command(["grep", "-nEHI", allowlist_pattern, filename])
-        except Exception as err:
-            return LintMessage(
-                path=None,
-                line=None,
-                char=None,
-                code=linter_name,
-                severity=LintSeverity.ERROR,
-                name="command-failed",
-                original=None,
-                replacement=None,
-                description=(
-                    f"Failed due to {err.__class__.__name__}:\n{err}"
-                    if not isinstance(err, subprocess.CalledProcessError)
-                    else (
-                        "COMMAND (exit code {returncode})\n"
-                        "{command}\n\n"
-                        "STDERR\n{stderr}\n\n"
-                        "STDOUT\n{stdout}"
-                    ).format(
-                        returncode=err.returncode,
-                        command=" ".join(as_posix(x) for x in err.cmd),
-                        stderr=err.stderr.decode("utf-8").strip() or "(empty)",
-                        stdout=err.stdout.decode("utf-8").strip() or "(empty)",
-                    )
-                ),
-            )
+    Args:
+        filename: Path to the file being linted
+        line_remainders: List of line remainders (format: "line:content" without filename prefix)
+        allowlist_pattern: Pattern to check for allowlisting
+        replace_pattern: Pattern for sed replacement
+        error_name: Human-readable error name
+    """
+    if not line_remainders:
+        return
 
-        # allowlist pattern was found, abort lint
-        if proc.returncode == 0:
-            return None
+    should_skip = check_allowlist(filename, allowlist_pattern)
+    if should_skip:
+        return
 
+    # Check if file is too large to compute replacement
+    file_size = os.path.getsize(filename)
+    compute_replacement = replace_pattern and file_size <= MAX_ORIGINAL_SIZE
+
+    # Apply replacement to entire file if pattern is specified and file is not too large
     original = None
     replacement = None
-    if replace_pattern:
-        with open(filename) as f:
-            original = f.read()
-
+    if compute_replacement:
+        # When we have a replacement, report a single message with line=None
         try:
+            with open(filename) as f:
+                original = f.read()
+
             proc = run_command(["sed", "-r", replace_pattern, filename])
             replacement = proc.stdout.decode("utf-8")
         except Exception as err:
-            return LintMessage(
-                path=None,
-                line=None,
-                char=None,
-                code=linter_name,
-                severity=LintSeverity.ERROR,
+            print_lint_message(
                 name="command-failed",
-                original=None,
-                replacement=None,
                 description=(
                     f"Failed due to {err.__class__.__name__}:\n{err}"
                     if not isinstance(err, subprocess.CalledProcessError)
@@ -138,18 +218,36 @@ def lint_file(
                     )
                 ),
             )
+            return
 
-    return LintMessage(
-        path=split[0],
-        line=int(split[1]) if len(split) > 1 else None,
-        char=None,
-        code=linter_name,
-        severity=LintSeverity.ERROR,
-        name=error_name,
-        original=original,
-        replacement=replacement,
-        description=error_description,
-    )
+        print_lint_message(
+            path=filename,
+            name=error_name,
+            original=original,
+            replacement=replacement,
+        )
+    else:
+        # When no replacement, report each matching line (up to MAX_MATCHES_PER_FILE)
+        total_matches = len(line_remainders)
+        matches_to_report = min(total_matches, MAX_MATCHES_PER_FILE)
+
+        for line_remainder in line_remainders[:matches_to_report]:
+            # line_remainder format: "line_number:content"
+            split = line_remainder.split(":", 1)
+            line_number = int(split[0]) if split[0] else None
+            print_lint_message(
+                path=filename,
+                line=line_number,
+                name=error_name,
+            )
+
+        # If there are more matches than the limit, print an error
+        if total_matches > MAX_MATCHES_PER_FILE:
+            print_lint_message(
+                path=filename,
+                name="too-many-matches",
+                description=f"File has {total_matches} matches, only showing first {MAX_MATCHES_PER_FILE}",
+            )
 
 
 def main() -> None:
@@ -203,8 +301,24 @@ def main() -> None:
         nargs="+",
         help="paths to lint",
     )
+
+    # Check for duplicate arguments before parsing
+    seen_args = set()
+    for arg in sys.argv[1:]:
+        if arg.startswith("--"):
+            arg_name = arg.split("=")[0]
+            if arg_name in seen_args:
+                parser.error(
+                    f"argument {arg_name}: not allowed to be specified multiple times"
+                )
+            seen_args.add(arg_name)
+
     args = parser.parse_args()
 
+    global LINTER_NAME, ERROR_DESCRIPTION
+    LINTER_NAME = args.linter_name
+    ERROR_DESCRIPTION = args.error_description
+
     logging.basicConfig(
         format="<%(threadName)s:%(levelname)s> %(message)s",
         level=logging.NOTSET
@@ -215,6 +329,31 @@ def main() -> None:
         stream=sys.stderr,
     )
 
+    # Filter out files that are too large before running grep
+    filtered_filenames = []
+    for filename in args.filenames:
+        try:
+            file_size = os.path.getsize(filename)
+            if file_size > MAX_FILE_SIZE:
+                print_lint_message(
+                    path=filename,
+                    severity=LintSeverity.WARNING,
+                    name="file-too-large",
+                    description=f"File size ({file_size} bytes) exceeds {MAX_FILE_SIZE} bytes limit, skipping",
+                )
+            else:
+                filtered_filenames.append(filename)
+        except OSError as err:
+            print_lint_message(
+                path=filename,
+                name="file-access-error",
+                description=f"Failed to get file size: {err}",
+            )
+
+    # If all files were filtered out, nothing to do
+    if not filtered_filenames:
+        return
+
     files_with_matches = []
     if args.match_first_only:
         files_with_matches = ["--files-with-matches"]
@@ -223,30 +362,23 @@ def main() -> None:
     try:
         # Split the grep command into multiple batches to avoid hitting the
         # command line length limit of ~1M on my machine
-        arg_length = sum(len(x) for x in args.filenames)
+        arg_length = sum(len(x) for x in filtered_filenames)
         batches = arg_length // 750000 + 1
-        batch_size = len(args.filenames) // batches
-        for i in range(0, len(args.filenames), batch_size):
+        batch_size = len(filtered_filenames) // batches
+        for i in range(0, len(filtered_filenames), batch_size):
             proc = run_command(
                 [
                     "grep",
                     "-nEHI",
                     *files_with_matches,
                     args.pattern,
-                    *args.filenames[i : i + batch_size],
+                    *filtered_filenames[i : i + batch_size],
                 ]
             )
             lines.extend(proc.stdout.decode().splitlines())
     except Exception as err:
-        err_msg = LintMessage(
-            path=None,
-            line=None,
-            char=None,
-            code=args.linter_name,
-            severity=LintSeverity.ERROR,
+        print_lint_message(
             name="command-failed",
-            original=None,
-            replacement=None,
             description=(
                 f"Failed due to {err.__class__.__name__}:\n{err}"
                 if not isinstance(err, subprocess.CalledProcessError)
@@ -263,20 +395,19 @@ def main() -> None:
                 )
             ),
         )
-        print(json.dumps(err_msg._asdict()), flush=True)
         sys.exit(0)
 
-    for line in lines:
-        lint_message = lint_file(
-            line,
+    # Group lines by file to call lint_file once per file
+    grouped_lines = group_lines_by_file(lines)
+
+    for filename, line_remainders in grouped_lines.items():
+        lint_file(
+            filename,
+            line_remainders,
             args.allowlist_pattern,
             args.replace_pattern,
-            args.linter_name,
             args.error_name,
-            args.error_description,
         )
-        if lint_message is not None:
-            print(json.dumps(lint_message._asdict()), flush=True)
 
 
 if __name__ == "__main__":

tools/linter/adapters/newlines_linter.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 import argparse
 import json
 import logging
+import os
 import sys
 from enum import Enum
 from typing import NamedTuple
@@ -15,6 +16,7 @@ from typing import NamedTuple
 NEWLINE = 10  # ASCII "\n"
 CARRIAGE_RETURN = 13  # ASCII "\r"
 LINTER_CODE = "NEWLINE"
+MAX_FILE_SIZE: int = 1024 * 1024 * 1024  # 1GB in bytes
 
 
 class LintSeverity(str, Enum):
@@ -39,6 +41,34 @@ class LintMessage(NamedTuple):
 def check_file(filename: str) -> LintMessage | None:
     logging.debug("Checking file %s", filename)
 
+    # Check if file is too large
+    try:
+        file_size = os.path.getsize(filename)
+        if file_size > MAX_FILE_SIZE:
+            return LintMessage(
+                path=filename,
+                line=None,
+                char=None,
+                code=LINTER_CODE,
+                severity=LintSeverity.WARNING,
+                name="file-too-large",
+                original=None,
+                replacement=None,
+                description=f"File size ({file_size} bytes) exceeds {MAX_FILE_SIZE} bytes limit, skipping",
+            )
+    except OSError as err:
+        return LintMessage(
+            path=filename,
+            line=None,
+            char=None,
+            code=LINTER_CODE,
+            severity=LintSeverity.ERROR,
+            name="file-access-error",
+            original=None,
+            replacement=None,
+            description=f"Failed to get file size: {err}",
+        )
+
     with open(filename, "rb") as f:
         lines = f.readlines()
 

tools/linter/adapters/pyfmt_linter.py

@@ -12,7 +12,7 @@ from enum import Enum
 from pathlib import Path
 from typing import NamedTuple
 
-# pyrefly: ignore  # import-error
+# pyrefly: ignore [import-error]
 import isort
 import usort
 

tools/linter/adapters/s3_init.py

@@ -55,7 +55,7 @@ def report_download_progress(
     Pretty printer for file download progress.
     """
     if file_size != -1:
-        # pyrefly: ignore  # no-matching-overload
+        # pyrefly: ignore [no-matching-overload]
         percent = min(1, (chunk_number * chunk_size) / file_size)
         bar = "#" * int(64 * percent)
         sys.stdout.write(f"\r0% |{bar:<64}| {int(percent * 100)}%")

tools/linter/adapters/test_has_main_linter.py

@@ -15,10 +15,10 @@ import multiprocessing as mp
 from enum import Enum
 from typing import NamedTuple
 
-# pyrefly: ignore  # import-error
+# pyrefly: ignore [import-error]
 import libcst as cst
 
-# pyrefly: ignore  # import-error
+# pyrefly: ignore [import-error]
 import libcst.matchers as m
 
 

tools/linter/adapters/workflow_consistency_linter.py

@@ -78,7 +78,7 @@ def print_lint_message(
 
     lint_message = LintMessage(
         path=str(path),
-        # pyrefly: ignore  # unbound-name
+        # pyrefly: ignore [unbound-name]
         line=line_number,
         char=None,
         code="WORKFLOWSYNC",

tools/lite_interpreter/gen_selected_mobile_ops_header.py

@@ -73,7 +73,7 @@ def get_selected_kernel_dtypes_code(
         for kernel_tag, dtypes in selective_builder.kernel_metadata.items():
             conditions = ["scalar_type == at::ScalarType::" + x for x in dtypes]
             body_parts.append(
-                # pyrefly: ignore  # bad-argument-type
+                # pyrefly: ignore [bad-argument-type]
                 if_condition_template.substitute(
                     kernel_tag_name=kernel_tag,
                     dtype_checks=" || ".join(conditions),

tools/nightly.py

@@ -311,7 +311,7 @@ class Venv:
             python=python,
             capture_output=True,
         ).stdout
-        # pyrefly: ignore  # no-matching-overload
+        # pyrefly: ignore [no-matching-overload]
         candidates = list(map(Path, filter(None, map(str.strip, output.splitlines()))))
         candidates = [p for p in candidates if p.is_dir() and p.name == "site-packages"]
         if not candidates:
@@ -481,7 +481,7 @@ class Venv:
         cmd = [str(python), *args]
         env = popen_kwargs.pop("env", None) or {}
         check = popen_kwargs.pop("check", True)
-        # pyrefly: ignore  # no-matching-overload
+        # pyrefly: ignore [no-matching-overload]
         return subprocess.run(
             cmd,
             check=check,
@@ -533,7 +533,7 @@ class Venv:
         cmd = [str(self.bindir / "uv"), *args]
         env = popen_kwargs.pop("env", None) or {}
         check = popen_kwargs.pop("check", True)
-        # pyrefly: ignore  # no-matching-overload
+        # pyrefly: ignore [no-matching-overload]
         return subprocess.run(
             cmd,
             check=check,
@@ -941,7 +941,7 @@ def _move_single(
 
 def _copy_files(listing: list[Path], source_dir: Path, target_dir: Path) -> None:
     for src in listing:
-        # pyrefly: ignore  # bad-argument-type
+        # pyrefly: ignore [bad-argument-type]
         _move_single(src, source_dir, target_dir, shutil.copy2, "Copying")
 
 

tools/nightly_hotpatch.py

@@ -118,7 +118,7 @@ def download_patch(pr_number: int, repo_url: str, download_dir: str) -> str:
             urllib.request.urlopen(patch_url) as response,
             open(patch_file, "wb") as out_file,
         ):
-            # pyrefly: ignore  # bad-specialization
+            # pyrefly: ignore [bad-specialization]
             shutil.copyfileobj(response, out_file)
         if not os.path.isfile(patch_file):
             print(f"Failed to download patch for PR #{pr_number}")

tools/pyi/gen_pyi.py

@@ -994,7 +994,7 @@ def add_docstr_to_hint(docstr: str, hint: str) -> str:
         hint = hint.removesuffix("...").rstrip()  # remove "..."
         content = hint + "\n" + textwrap.indent(f'r"""\n{docstr}\n"""', prefix="    ")
         # Remove trailing whitespace on each line
-        # pyrefly: ignore  # no-matching-overload
+        # pyrefly: ignore [no-matching-overload]
         return "\n".join(map(str.rstrip, content.splitlines())).rstrip()
 
     # attribute or property

tools/setup_helpers/cmake.py

@@ -100,7 +100,7 @@ class CMake:
             if ver is not None:
                 eprint(f"Found {cmd} ({command}) version: {ver}", end="")
                 cmake_versions.append(f"{cmd}=={ver}")
-                # pyrefly: ignore  # unsupported-operation
+                # pyrefly: ignore [unsupported-operation]
                 if ver >= CMAKE_MINIMUM_VERSION:
                     eprint(f" (>={CMAKE_MINIMUM_VERSION})")
                     valid_cmake_versions[cmd] = ver

tools/setup_helpers/generate_linker_script.py

@@ -31,9 +31,9 @@ def gen_linker_script(
     text_line_start = text_line_start[0]
 
     # ensure that parent directory exists before writing
-    # pyrefly: ignore  # bad-assignment
+    # pyrefly: ignore [bad-assignment]
     fout = Path(fout)
-    # pyrefly: ignore  # missing-attribute
+    # pyrefly: ignore [missing-attribute]
     fout.parent.mkdir(parents=True, exist_ok=True)
 
     with open(fout, "w") as f:

tools/stats/upload_utilization_stats/upload_utilization_stats.py

@@ -60,7 +60,7 @@ class SegmentGenerator:
         df[time_col_name] = pd.to_datetime(df[time_col_name], unit="s", utc=True)
 
         # get unique cmd names
-        # pyrefly: ignore  # bad-argument-type
+        # pyrefly: ignore [bad-argument-type]
         unique_cmds_df = pd.DataFrame(df[cmd_col_name].unique(), columns=[cmd_col_name])
 
         # get all detected python cmds

tools/stats/utilization_stats_lib.py

@@ -2,7 +2,8 @@ from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Optional
 
-from dataclasses_json import DataClassJsonMixin
+#  pyrefly: ignore [missing-import]
+from dataclasses_json import DataClassJsonMixin  # type: ignore[import-not-found]
 
 
 _DATA_MODEL_VERSION = 1.5
@@ -17,7 +18,7 @@ class UtilizationStats:
 
 
 @dataclass
-class UtilizationMetadata(DataClassJsonMixin):
+class UtilizationMetadata(DataClassJsonMixin):  # type: ignore[misc, no-any-unimported]
     level: str
     workflow_id: str
     job_id: str
@@ -33,7 +34,7 @@ class UtilizationMetadata(DataClassJsonMixin):
 
 
 @dataclass
-class GpuUsage(DataClassJsonMixin):
+class GpuUsage(DataClassJsonMixin):  # type: ignore[misc, no-any-unimported]
     uuid: Optional[str] = None
     util_percent: Optional[UtilizationStats] = None
     mem_util_percent: Optional[UtilizationStats] = None
@@ -43,14 +44,14 @@ class GpuUsage(DataClassJsonMixin):
 
 
 @dataclass
-class RecordData(DataClassJsonMixin):
+class RecordData(DataClassJsonMixin):  # type: ignore[misc, no-any-unimported]
     cpu: Optional[UtilizationStats] = None
     memory: Optional[UtilizationStats] = None
     gpu_usage: Optional[list[GpuUsage]] = None
 
 
 @dataclass
-class UtilizationRecord(DataClassJsonMixin):
+class UtilizationRecord(DataClassJsonMixin):  # type: ignore[misc, no-any-unimported]
     level: str
     timestamp: int
     data: Optional[RecordData] = None
@@ -63,7 +64,7 @@ class UtilizationRecord(DataClassJsonMixin):
 # the db schema related to this is:
 # https://github.com/pytorch/test-infra/blob/main/clickhouse_db_schema/oss_ci_utilization/oss_ci_utilization_metadata_schema.sql
 @dataclass
-class OssCiSegmentV1(DataClassJsonMixin):
+class OssCiSegmentV1(DataClassJsonMixin):  # type: ignore[misc, no-any-unimported]
     level: str
     name: str
     start_at: int

tools/test/gen_operators_yaml_test.py

@@ -7,7 +7,7 @@ import unittest
 from collections import defaultdict
 from unittest.mock import Mock, patch
 
-# pyrefly: ignore  # import-error
+# pyrefly: ignore [import-error]
 from gen_operators_yaml import (
     fill_output,
     get_parser_options,
@@ -242,6 +242,6 @@ class GenOperatorsYAMLTest(unittest.TestCase):
 
         fill_output(output, options)
 
-        # pyrefly: ignore  # missing-attribute
+        # pyrefly: ignore [missing-attribute]
         for op_val in output["operators"].values():
             self.assertFalse(op_val["include_all_overloads"])

tools/test/test_selective_build.py

@@ -88,7 +88,7 @@ operators:
         self.assertTrue(selector2.is_operator_selected("aten::sub.int"))
 
         selector_legacy_v1 = SelectiveBuilder.from_legacy_op_registration_allow_list(
-            # pyrefly: ignore  # bad-argument-type
+            # pyrefly: ignore [bad-argument-type]
             ["aten::add", "aten::add.int", "aten::mul.int"],
             False,
             False,
@@ -104,7 +104,7 @@ operators:
         )
 
         selector_legacy_v1 = SelectiveBuilder.from_legacy_op_registration_allow_list(
-            # pyrefly: ignore  # bad-argument-type
+            # pyrefly: ignore [bad-argument-type]
             ["aten::add", "aten::add.int", "aten::mul.int"],
             True,
             False,
@@ -120,7 +120,7 @@ operators:
         )
 
         selector_legacy_v1 = SelectiveBuilder.from_legacy_op_registration_allow_list(
-            # pyrefly: ignore  # bad-argument-type
+            # pyrefly: ignore [bad-argument-type]
             ["aten::add", "aten::add.int", "aten::mul.int"],
             False,
             True,

tools/testing/target_determination/heuristics/historical_class_failure_correlation.py

@@ -83,9 +83,9 @@ def _rank_correlated_tests(
 ) -> list[str]:
     # Find the tests failures that are correlated with the edited files.
     # Filter the list to only include tests we want to run.
-    # pyrefly: ignore  # bad-assignment
+    # pyrefly: ignore [bad-assignment]
     tests_to_run = set(tests_to_run)
-    # pyrefly: ignore  # bad-argument-type
+    # pyrefly: ignore [bad-argument-type]
     ratings = _get_ratings_for_tests(tests_to_run)
     prioritize = sorted(ratings, key=lambda x: -ratings[x])
     return prioritize

tools/testing/upload_artifacts.py

@@ -36,13 +36,13 @@ def concated_logs() -> str:
     for log_file in glob.glob(
         f"{REPO_ROOT}/test/test-reports/**/*.log", recursive=True
     ):
-        # pyrefly: ignore  # bad-argument-type
+        # pyrefly: ignore [bad-argument-type]
         logs.append(f"=== {log_file} ===")
         with open(log_file) as f:
             # For every line, prefix with fake timestamp for log classifier
             for line in f:
                 line = line.rstrip("\n")  # Remove any trailing newline
-                # pyrefly: ignore  # bad-argument-type
+                # pyrefly: ignore [bad-argument-type]
                 logs.append(f"2020-01-01T00:00:00.0000000Z {line}")
     return "\n".join(logs)
 

torch/_C/__init__.pyi.in

@@ -2759,6 +2759,7 @@ class _NodeBase:
     ) -> None: ...
     def _update_args_kwargs(self, args: tuple[Any, ...], kwargs: dict[str, Any]): ...
     def _prepend(self, n: FxNode) -> None: ...
+    def _replace_input_with(self, old_input: FxNode, new_input: FxNode) -> None: ...
     def _remove_from_list(self) -> None: ...
     def __lt__(self, n: Self) -> _bool: ...
     def __gt__(self, n: Self) -> _bool: ...

torch/__init__.py

@@ -1703,7 +1703,7 @@ def _check(cond, message=None):  # noqa: F811
             an object that has a ``__str__()`` method to be used as the error
             message. Default: ``None``
     """
-    _check_with(RuntimeError, cond, message)  # pyrefly: ignore  # bad-argument-type
+    _check_with(RuntimeError, cond, message)  # pyrefly: ignore [bad-argument-type]
 
 
 # TODO add deprecation annotation
@@ -1753,7 +1753,7 @@ def _check_index(cond, message=None):  # noqa: F811
             an object that has a ``__str__()`` method to be used as the error
             message. Default: ``None``
     """
-    _check_with(IndexError, cond, message)  # pyrefly: ignore  # bad-argument-type
+    _check_with(IndexError, cond, message)  # pyrefly: ignore [bad-argument-type]
 
 
 def _check_value(cond, message=None):  # noqa: F811
@@ -1771,7 +1771,7 @@ def _check_value(cond, message=None):  # noqa: F811
             an object that has a ``__str__()`` method to be used as the error
             message. Default: ``None``
     """
-    _check_with(ValueError, cond, message)  # pyrefly: ignore  # bad-argument-type
+    _check_with(ValueError, cond, message)  # pyrefly: ignore [bad-argument-type]
 
 
 def _check_type(cond, message=None):  # noqa: F811
@@ -1789,7 +1789,7 @@ def _check_type(cond, message=None):  # noqa: F811
             an object that has a ``__str__()`` method to be used as the error
             message. Default: ``None``
     """
-    _check_with(TypeError, cond, message)  # pyrefly: ignore  # bad-argument-type
+    _check_with(TypeError, cond, message)  # pyrefly: ignore [bad-argument-type]
 
 
 def _check_not_implemented(cond, message=None):  # noqa: F811
@@ -2653,7 +2653,8 @@ def compile(
     if torch.compiler.is_exporting():
         warnings.warn(
             "You are calling torch.compile inside torch.export region. "
-            "To capture an useful graph, we will implicitly switch to torch.compile(backend=eager)"
+            "To capture an useful graph, we will implicitly switch to torch.compile(backend=eager)",
+            stacklevel=2,
         )
         from torch._higher_order_ops.utils import setup_compilation_env
 

torch/_custom_op/impl.py

@@ -55,6 +55,7 @@ def warn_deprecated():
         "torch._custom_op is deprecated and will be removed in PyTorch 2.6, please "
         "use the equivalent torch.library API instead.",
         DeprecationWarning,
+        stacklevel=2,
     )
 
 
@@ -101,7 +102,7 @@ def custom_op(
             lib, ns, function_schema, name, ophandle, _private_access=True
         )
 
-        result.__name__ = func.__name__  # pyrefly: ignore  # bad-assignment
+        result.__name__ = func.__name__  # pyrefly: ignore [bad-assignment]
         result.__module__ = func.__module__
         result.__doc__ = func.__doc__
 

torch/_dispatch/python.py

@@ -154,7 +154,7 @@ def make_crossref_functionalize(
                 maybe_detach, (f_args, f_kwargs)
             )
             with fake_mode:
-                f_r = op(*f_args, **f_kwargs)  # pyrefly: ignore  # invalid-param-spec
+                f_r = op(*f_args, **f_kwargs)  # pyrefly: ignore [invalid-param-spec]
         r = op._op_dk(final_key, *args, **kwargs)
 
         def desc():

torch/_dynamo/__init__.py

@@ -55,7 +55,7 @@ from .eval_frame import (
     reset_code,
 )
 
-# pyrefly: ignore  # deprecated
+# pyrefly: ignore [deprecated]
 from .external_utils import is_compiling
 from .mutation_guard import GenerationTracker
 from .pgo import reset_code_state

torch/_dynamo/_trace_wrapped_higher_order_op.py

@@ -95,7 +95,7 @@ class ModIndex(torch.autograd.Function):
     generate_vmap_rule = True
 
     @staticmethod
-    # pyrefly: ignore  # bad-override
+    # pyrefly: ignore [bad-override]
     def forward(x: Tensor, indices: list[Tensor]) -> Tensor:
         return torch.ops.aten.index(x, indices)
 
@@ -243,7 +243,7 @@ def _trace_wrapped_functionalized(ctx: Any, *args: Any, **kwargs: Any) -> Any:
 
 def autograd_function_backward_rewritten(original_backward: Any) -> Any:
     def new_backward(ctx: Any, *grads: Any) -> Any:
-        # pyrefly: ignore  # bad-assignment
+        # pyrefly: ignore [bad-assignment]
         grads = [g.contiguous() for g in grads]
         return original_backward(ctx, *grads)
 

torch/_dynamo/aot_compile.py

@@ -81,7 +81,7 @@ class AOTCompiledFunction:
             **import_sources,
             self._artifacts.backend_id: self._artifacts.compiled_fn,
         }
-        # pyrefly: ignore  # read-only
+        # pyrefly: ignore [read-only]
         self.fn = types.FunctionType(
             self._artifacts.bytecode, f_globals, closure=self._artifacts.closure
         )

torch/_dynamo/backends/cudagraphs.py

@@ -206,7 +206,7 @@ def cudagraphs(dynamo_model: torch.fx.GraphModule, dynamo_inputs: Sequence[Any])
             assert manager is not None
 
             def fn(inputs: list[Any]) -> Any:
-                # pyrefly: ignore  # missing-attribute
+                # pyrefly: ignore [missing-attribute]
                 manager.set_to_running_backward()
                 return aot_model(inputs)
 

torch/_dynamo/backends/tvm.py

@@ -78,19 +78,19 @@ def tvm(
     opt_level = options.get("opt_level", 3)
 
     if scheduler == "auto_scheduler":
-        # pyrefly: ignore  # import-error
+        # pyrefly: ignore [import-error]
         from tvm import auto_scheduler
 
         log_file = tempfile.NamedTemporaryFile()
 
-        # pyrefly: ignore  # bad-argument-type
+        # pyrefly: ignore [bad-argument-type]
         if not os.path.exists(log_file):
             tasks, task_weights = auto_scheduler.extract_tasks(
                 mod["main"], params, target
             )
             if len(tasks) != 0:
                 tuner = auto_scheduler.TaskScheduler(tasks, task_weights)
-                # pyrefly: ignore  # bad-argument-type
+                # pyrefly: ignore [bad-argument-type]
                 if not os.path.exists(log_file):
                     assert trials > 0
                     tune_option = auto_scheduler.TuningOptions(
@@ -101,9 +101,9 @@ def tvm(
                     try:
                         tuner.tune(tune_option)
                     except Exception:
-                        # pyrefly: ignore  # bad-argument-type
+                        # pyrefly: ignore [bad-argument-type]
                         if os.path.exists(log_file):
-                            # pyrefly: ignore  # bad-argument-type
+                            # pyrefly: ignore [bad-argument-type]
                             os.unlink(log_file)
                         raise
 
@@ -113,7 +113,7 @@ def tvm(
             ):
                 lib = relay.build(mod, target=target, params=params)
     elif scheduler == "meta_schedule":
-        # pyrefly: ignore  # import-error
+        # pyrefly: ignore [import-error]
         from tvm import meta_schedule as ms
 
         with tempfile.TemporaryDirectory() as work_dir:

torch/_dynamo/bytecode_analysis.py

@@ -37,7 +37,7 @@ if sys.version_info >= (3, 11):
     TERMINAL_OPCODES.add(dis.opmap["JUMP_FORWARD"])
 else:
     TERMINAL_OPCODES.add(dis.opmap["JUMP_ABSOLUTE"])
-# pyrefly: ignore  # unsupported-operation
+# pyrefly: ignore [unsupported-operation]
 if (3, 12) <= sys.version_info < (3, 14):
     TERMINAL_OPCODES.add(dis.opmap["RETURN_CONST"])
 if sys.version_info >= (3, 13):

torch/_dynamo/bytecode_transformation.py

@@ -923,7 +923,7 @@ def devirtualize_jumps(instructions: list[Instruction]) -> None:
                 inst.arg = abs(
                     int(target.offset - inst.offset - instruction_size(inst))
                 )
-                # pyrefly: ignore  # unsupported-operation
+                # pyrefly: ignore [unsupported-operation]
                 inst.arg //= 2
             inst.argval = target.offset
             inst.argrepr = f"to {target.offset}"
@@ -1375,7 +1375,7 @@ def update_offsets(instructions: Sequence[Instruction]) -> None:
     offset = 0
     for inst in instructions:
         inst.offset = offset
-        # pyrefly: ignore  # unsupported-operation
+        # pyrefly: ignore [unsupported-operation]
         offset += instruction_size(inst)
 
 

torch/_dynamo/convert_frame.py

@@ -475,7 +475,7 @@ def cprofile_wrapper(func: Callable[_P, _T]) -> Callable[_P, _T]:
         try:
             prof.enable()
             start_ts = time.time()
-            # pyrefly: ignore  # bad-argument-type
+            # pyrefly: ignore [bad-argument-type]
             retval = prof.runcall(func, *args, **kwargs)
             profile_latency = time.time() - start_ts
             prof.disable()
@@ -995,7 +995,7 @@ def get_traced_fn(mod: Any) -> tuple[FunctionType, Optional[object]]:
     if isinstance(mod, torch.nn.Module):
         mod = mod.forward
     if hasattr(mod, "__self__"):
-        # pyrefly: ignore  # missing-attribute
+        # pyrefly: ignore [missing-attribute]
         return mod.__func__, mod.__self__
     elif inspect.isfunction(mod):
         return mod, None
@@ -1140,7 +1140,7 @@ def _fullgraph_capture_frame(
         while cur_exn.__cause__ is not None:
             cur_exn.__cause__.with_traceback(None)
             cur_exn = cur_exn.__cause__
-        # pyrefly: ignore  # invalid-inheritance
+        # pyrefly: ignore [invalid-inheritance]
         raise e.with_traceback(None) from e.__cause__  # User compiler error
 
     return CaptureOutput(
@@ -1164,7 +1164,7 @@ def compile_frame(  # type: ignore[return]
     frame_state: Optional[dict[str, Union[int, FrameStateSizeEntry]]] = None,
     distributed_state: Optional[DistributedState] = None,
     package: Optional[CompilePackage] = None,
-    # pyrefly: ignore  # bad-return
+    # pyrefly: ignore [bad-return]
 ) -> DynamoOutput:
     """
     A helper function taking a frame and backend, then return the generated bytecode

torch/_dynamo/create_parameter_op.py

@@ -20,7 +20,7 @@ allowed to compute gradients on).
 
 class TracableCreateParameter(torch.autograd.Function):
     @staticmethod
-    # pyrefly: ignore  # bad-override
+    # pyrefly: ignore [bad-override]
     def forward(ctx: Any, tensor: Any, placeholder: Any) -> torch.nn.Parameter:
         assert not tensor.requires_grad
         return placeholder.set_(tensor)

torch/_dynamo/debug_utils.py

@@ -879,7 +879,7 @@ def aot_graph_input_parser(
             data_type, shape_str = match.groups()
             shape = tuple(shape_str.split(","))
             dtype = dtype_map[data_type]
-            # pyrefly: ignore  # bad-argument-type
+            # pyrefly: ignore [bad-argument-type]
             kwargs[param] = gen_tensor(shape, dtype)
 
         match = re.search(sym_shape_regex, annotation)
@@ -893,7 +893,7 @@ def aot_graph_input_parser(
             attr_name, data_type, shape_str, _ = match.groups()
             shape = tuple(shape_str.split(","))
             dtype = dtype_map[data_type]
-            # pyrefly: ignore  # bad-argument-type
+            # pyrefly: ignore [bad-argument-type]
             setattr(container, attr_name, gen_tensor(shape, dtype))
 
     return kwargs

torch/_dynamo/decorators.py

@@ -96,7 +96,7 @@ def disable(fn=None, recursive=True, *, reason=None, wrapping=True):  # type: ig
             nonrecursive_disable_wrapper._torchdynamo_disable = True  # type: ignore[attr-defined]
             nonrecursive_disable_wrapper._torchdynamo_disable_msg = reason  # type: ignore[attr-defined]
             nonrecursive_disable_wrapper._torchdynamo_orig_callable = fn  # type: ignore[attr-defined]
-            # pyrefly: ignore  # bad-return
+            # pyrefly: ignore [bad-return]
             return nonrecursive_disable_wrapper
 
         if fn is None:
@@ -316,7 +316,7 @@ def forbid_in_graph(fn: Any) -> Any:
     if isinstance(fn, (list, tuple)):
         return [forbid_in_graph(x) for x in fn]
     assert callable(fn), "forbid_in_graph applies only to callables"
-    # pyrefly: ignore  # missing-attribute
+    # pyrefly: ignore [missing-attribute]
     fn._dynamo_forbidden = True
     return fn
 
@@ -493,10 +493,16 @@ def substitute_in_graph(
         def dispatch_fn(
             self: VariableBuilder, value: Callable[_P, _R]
         ) -> PolyfilledFunctionVariable:
+            if inspect.isclass(value):
+                guard_type = GuardBuilder.CLASS_MATCH
+            elif inspect.ismodule(value):
+                guard_type = GuardBuilder.MODULE_MATCH
+            else:
+                guard_type = GuardBuilder.ID_MATCH
             return PolyfilledFunctionVariable(
                 value,
                 source=self.source,
-                **self.install_guards(GuardBuilder.FUNCTION_MATCH),
+                **self.install_guards(guard_type),
             )
 
         id_dispatch_map[id(original_fn)] = id_dispatch_map[id(wrapped)] = dispatch_fn
@@ -664,28 +670,28 @@ def mark_dynamic(
 
     if isinstance(index, int):
         if not hasattr(t, "_dynamo_dynamic_indices"):
-            # pyrefly: ignore  # missing-attribute
+            # pyrefly: ignore [missing-attribute]
             t._dynamo_dynamic_indices = set()
-            # pyrefly: ignore  # missing-attribute
+            # pyrefly: ignore [missing-attribute]
             t._dynamo_dynamic_range = set()
-            # pyrefly: ignore  # missing-attribute
+            # pyrefly: ignore [missing-attribute]
             t._dynamo_hint_overrides = {}
 
         if not hasattr(t, "_specialize_on"):
-            # pyrefly: ignore  # missing-attribute
+            # pyrefly: ignore [missing-attribute]
             t._specialize_on = {}
 
         if hint_override:
-            # pyrefly: ignore  # missing-attribute
+            # pyrefly: ignore [missing-attribute]
             t._dynamo_hint_overrides[index] = hint_override
         # TODO(voz): Should we bounds check?
-        # pyrefly: ignore  # missing-attribute
+        # pyrefly: ignore [missing-attribute]
         t._dynamo_dynamic_indices.add(index)
         t._dynamo_dynamic_range.add(_DimRange(index, min, max))  # type: ignore[arg-type]
 
         # FX tracers don't respect @forbid_in_graph and choke on the following error since it passes in proxies:
         # TypeError: 'Attribute' object does not support item assignment
-        # pyrefly: ignore  # missing-attribute
+        # pyrefly: ignore [missing-attribute]
         if isinstance(t._specialize_on, dict):
             t._specialize_on[index] = specialize_on if specialize_on is not None else []
 
@@ -710,10 +716,10 @@ def maybe_mark_dynamic(t: Any, index: Union[int, list[Any], tuple[Any]]) -> None
 
     if isinstance(index, int):
         if not hasattr(t, "_dynamo_weak_dynamic_indices"):
-            # pyrefly: ignore  # missing-attribute
+            # pyrefly: ignore [missing-attribute]
             t._dynamo_weak_dynamic_indices = set()
         # TODO(voz): Should we bounds check?
-        # pyrefly: ignore  # missing-attribute
+        # pyrefly: ignore [missing-attribute]
         t._dynamo_weak_dynamic_indices.add(index)
         return
 
@@ -765,11 +771,11 @@ def mark_static(
         # TODO: Make this configurable via a supported public API
         _apply_func_to_inner_tensors_of_same_dim(mark_static, t, index)
 
-    # pyrefly: ignore  # bad-argument-type
+    # pyrefly: ignore [bad-argument-type]
     if not isinstance(t, torch.Tensor) and issubclass(t, torch.nn.Module):
-        # pyrefly: ignore  # missing-attribute
+        # pyrefly: ignore [missing-attribute]
         t._dynamo_marked_static = True
-        # pyrefly: ignore  # bad-return
+        # pyrefly: ignore [bad-return]
         return t
 
     if not isinstance(t, torch.Tensor):

torch/_dynamo/device_interface.py

@@ -205,7 +205,7 @@ class CudaInterface(DeviceInterface):
     Event = torch.cuda.Event  # type: ignore[assignment]
     Stream = torch.cuda.Stream  # type: ignore[assignment]
 
-    # pyrefly: ignore  # bad-override
+    # pyrefly: ignore [bad-override]
     class Worker:
         @staticmethod
         def set_device(device: int) -> None:
@@ -241,7 +241,7 @@ class CudaInterface(DeviceInterface):
     set_device = staticmethod(torch.cuda.set_device)
     device_count = staticmethod(torch.cuda.device_count)
     stream = staticmethod(torch.cuda.stream)  # type: ignore[assignment]
-    # pyrefly: ignore  # bad-override
+    # pyrefly: ignore [bad-override]
     current_stream = staticmethod(torch.cuda.current_stream)
     set_stream = staticmethod(torch.cuda.set_stream)  # type: ignore[assignment]
     _set_stream_by_id = staticmethod(torch.cuda._set_stream_by_id)  # type: ignore[assignment]
@@ -302,7 +302,7 @@ class MtiaInterface(DeviceInterface):
     Event = torch.mtia.Event  # type: ignore[assignment]
     Stream = torch.mtia.Stream  # type: ignore[assignment]
 
-    # pyrefly: ignore  # bad-override
+    # pyrefly: ignore [bad-override]
     class Worker:
         @staticmethod
         def set_device(device: int) -> None:
@@ -338,7 +338,7 @@ class MtiaInterface(DeviceInterface):
     set_device = staticmethod(torch.mtia.set_device)  # type: ignore[assignment]
     device_count = staticmethod(torch.mtia.device_count)
     stream = staticmethod(torch.mtia.stream)  # type: ignore[assignment]
-    # pyrefly: ignore  # bad-override
+    # pyrefly: ignore [bad-override]
     current_stream = staticmethod(torch.mtia.current_stream)
     set_stream = staticmethod(torch.mtia.set_stream)  # type: ignore[assignment]
     _set_stream_by_id = staticmethod(torch.mtia._set_stream_by_id)  # type: ignore[assignment]
@@ -385,7 +385,7 @@ class XpuInterface(DeviceInterface):
     Event = torch.xpu.Event  # type: ignore[assignment]
     Stream = torch.xpu.Stream  # type: ignore[assignment]
 
-    # pyrefly: ignore  # bad-override
+    # pyrefly: ignore [bad-override]
     class Worker:
         @staticmethod
         def set_device(device: int) -> None:
@@ -421,7 +421,7 @@ class XpuInterface(DeviceInterface):
     set_device = staticmethod(torch.xpu.set_device)
     device_count = staticmethod(torch.xpu.device_count)  # type: ignore[has-type]
     stream = staticmethod(torch.xpu.stream)  # type: ignore[assignment]
-    # pyrefly: ignore  # bad-override
+    # pyrefly: ignore [bad-override]
     current_stream = staticmethod(torch.xpu.current_stream)
     set_stream = staticmethod(torch.xpu.set_stream)  # type: ignore[assignment]
     _set_stream_by_id = staticmethod(torch.xpu._set_stream_by_id)  # type: ignore[assignment]
@@ -464,7 +464,7 @@ class CpuDeviceProperties:
 
 
 class CpuInterface(DeviceInterface):
-    # pyrefly: ignore  # bad-override
+    # pyrefly: ignore [bad-override]
     class Event(torch.Event):
         def __init__(self, enable_timing: bool = True) -> None:
             self.time = 0.0
@@ -475,7 +475,7 @@ class CpuInterface(DeviceInterface):
         def record(self, stream: Any = None) -> None:
             self.time = time.perf_counter()
 
-    # pyrefly: ignore  # bad-override
+    # pyrefly: ignore [bad-override]
     class Worker:
         @staticmethod
         def get_device_properties(
@@ -551,7 +551,7 @@ class MpsInterface(DeviceInterface):
     def synchronize(device: torch.types.Device = None) -> None:
         torch.mps.synchronize()
 
-    # pyrefly: ignore  # bad-override
+    # pyrefly: ignore [bad-override]
     class Worker:
         @staticmethod
         def get_device_properties(device: torch.types.Device = None) -> Any:

torch/_dynamo/eval_frame.py

@@ -499,7 +499,7 @@ class OptimizedModule(torch.nn.Module):
         self._initialize()
 
     @property
-    # pyrefly: ignore  # bad-override
+    # pyrefly: ignore [bad-override]
     def training(self) -> bool:
         return self._orig_mod.training
 
@@ -918,7 +918,7 @@ class _TorchDynamoContext:
                     while cur_exn.__cause__ is not None:
                         cur_exn.__cause__.with_traceback(None)
                         cur_exn = cur_exn.__cause__
-                    # pyrefly: ignore  # invalid-inheritance
+                    # pyrefly: ignore [invalid-inheritance]
                     raise e.with_traceback(None) from e.__cause__  # User compiler error
                 except ShortenTraceback as e:
                     # Failures in the backend likely don't have useful
@@ -1048,7 +1048,7 @@ class OptimizeContext(_TorchDynamoContext):
                 compiler_fn = rebuild_ctx()
                 ctx = torch._dynamo.compiled_autograd._enable(
                     compiler_fn,
-                    # pyrefly: ignore  # bad-argument-type
+                    # pyrefly: ignore [bad-argument-type]
                     dynamic=_dynamic,
                     ignore_active_disable_ctx=False,
                 )
@@ -1113,7 +1113,7 @@ class DisableContext(_TorchDynamoContext):
             cls_obj.__call__ = self(cls_obj.__call__)
             if issubclass(cls_obj, torch.nn.Module):
                 # NN module variable tracker directly inlines the _call_impl. Disable it.
-                # pyrefly: ignore  # missing-attribute
+                # pyrefly: ignore [missing-attribute]
                 cls_obj._call_impl = self(cls_obj._call_impl)
             return cls_obj
 
@@ -2018,7 +2018,7 @@ def export(
                         path: KeyPath, t: Union[torch.Tensor, _IntWrapper, Any]
                     ) -> Any:
                         if isinstance(t, torch.Tensor):
-                            # pyrefly: ignore  # missing-attribute
+                            # pyrefly: ignore [missing-attribute]
                             return ambient_fake_mode.from_tensor(t, static_shapes=True)
                         elif isinstance(t, _IntWrapper):
                             if (
@@ -2103,11 +2103,10 @@ def export(
             )
             and not trace_rules.check(call_to_inspect)
         ):
-            # pyrefly: ignore  # unbound-name
             dim_constraints.solve()
-            # pyrefly: ignore  # unbound-name
+
             forced_specializations = dim_constraints.forced_specializations()
-            # pyrefly: ignore  # unbound-name
+
             msg = dim_constraints.prettify_results(
                 original_signature,
                 dynamic_shapes,
@@ -2128,11 +2127,10 @@ def export(
                     )
 
             # Error if we have any constraints on static values
-            # pyrefly: ignore  # unbound-name
+
             for k in shape_env.var_to_range.keys():
                 if isinstance(k, sympy.Integer):
                     constraint_violation_error = ConstraintViolationError(
-                        # pyrefly: ignore  # unbound-name
                         f"{''.join(traceback.format_list(shape_env.var_to_stack[k]))}\n"
                         "It appears that you're trying to set a constraint on a "
                         f"value which we evaluated to have a static value of {k}. "

torch/_dynamo/exc.py

@@ -384,7 +384,7 @@ def get_dynamo_observed_exception(exc_type: type[Exception]) -> type[ObservedExc
         observed_exception_map[exc_type] = type(  # type: ignore[assignment]
             f"Observed{name}Error", (ObservedException,), {}
         )
-    # pyrefly: ignore  # index-error
+    # pyrefly: ignore [index-error]
     return observed_exception_map[exc_type]
 
 

torch/_dynamo/external_utils.py

@@ -97,9 +97,9 @@ def wrap_numpy(f: Callable[_P, _R]) -> Callable[_P, _R]:
         args, kwargs = pytree.tree_map_only(
             torch.Tensor, lambda x: x.numpy(), (args, kwargs)
         )
-        # pyrefly: ignore  # invalid-param-spec
+        # pyrefly: ignore [invalid-param-spec]
         out = f(*args, **kwargs)
-        # pyrefly: ignore  # missing-attribute
+        # pyrefly: ignore [missing-attribute]
         return pytree.tree_map_only(np.ndarray, lambda x: torch.as_tensor(x), out)
 
     return wrap

torch/_dynamo/functional_export.py

@@ -290,7 +290,7 @@ class DynamoGraphTransformer(torch.fx.Transformer):
             else:
                 placeholder.node.meta["val"] = self.flat_inputs[i]
 
-            # pyrefly: ignore  # unsupported-operation
+            # pyrefly: ignore [unsupported-operation]
             self.new_input_nodes[i] = placeholder
 
     def _create_placeholder_mapping(self) -> None:
@@ -365,18 +365,18 @@ class DynamoGraphTransformer(torch.fx.Transformer):
 
         # Copy module metadata like the original implementation
         if hasattr(self.module, "meta"):
-            # pyrefly: ignore  # unsupported-operation
+            # pyrefly: ignore [unsupported-operation]
             if "dynamo_flat_name_to_original_fqn" in self.module.meta:
-                # pyrefly: ignore  # index-error
+                # pyrefly: ignore [index-error]
                 result_gm.meta["dynamo_flat_name_to_original_fqn"] = self.module.meta[
-                    # pyrefly: ignore  # index-error
+                    # pyrefly: ignore [index-error]
                     "dynamo_flat_name_to_original_fqn"
                 ]
-            # pyrefly: ignore  # unsupported-operation
+            # pyrefly: ignore [unsupported-operation]
             if "dynamo_compile_id" in self.module.meta:
-                # pyrefly: ignore  # index-error
+                # pyrefly: ignore [index-error]
                 result_gm.meta["dynamo_compile_id"] = self.module.meta[
-                    # pyrefly: ignore  # index-error
+                    # pyrefly: ignore [index-error]
                     "dynamo_compile_id"
                 ]
 
@@ -408,11 +408,10 @@ def _suggest_or_raise_constraint_violation(
             torch._ops.OpOverloadPacket | torch._ops.OpOverload,
         )
     ):
-        # pyrefly: ignore  # unbound-name
         dim_constraints.solve()
-        # pyrefly: ignore  # unbound-name
+
         forced_specializations = dim_constraints.forced_specializations()
-        # pyrefly: ignore  # unbound-name
+
         msg = dim_constraints.prettify_results(
             inspect.signature(orig_callable),  # type: ignore[attr-defined]
             dynamic_shapes,
@@ -433,11 +432,10 @@ def _suggest_or_raise_constraint_violation(
                 )
 
         # Error if we have any constraints on static values
-        # pyrefly: ignore  # unbound-name
+
         for k in shape_env.var_to_range.keys():
             if isinstance(k, sympy.Integer):
                 constraint_violation_error = ConstraintViolationError(
-                    # pyrefly: ignore  # unbound-name
                     f"{''.join(traceback.format_list(shape_env.var_to_stack[k]))}\n"
                     "It appears that you're trying to set a constraint on a "
                     f"value which we evaluated to have a static value of {k}. "

torch/_dynamo/graph_deduplication.py

@@ -456,8 +456,10 @@ def _add_mutation_dependencies(
             for user in mutated_arg.users:
                 if user is node:
                     continue
+                # pyrefly: ignore [unsupported-operation]
                 elif user < node:
                     node_to_additional_deps[node].add(user)
+                # pyrefly: ignore [unsupported-operation]
                 elif user > node:
                     node_to_additional_deps[user].add(node)
 

torch/_dynamo/graph_region_tracker.py

@@ -333,7 +333,7 @@ class GraphRegionTracker:
             if len(group) > 1:
                 region_group = []
                 min_rank = math.inf
-                # pyrefly: ignore  # bad-assignment
+                # pyrefly: ignore [bad-assignment]
                 for node in group:
                     # some nodes aren't in the topo ranking?
                     if node in topological_ranking: