Remove unused type error suppressions - ads_mkl/ops

Summary:
This diff was automatically generated by the Pyre per-target upgrade tool.

It removes `# pyre-fixme` or `pyrefly: ignore` comments that are no longer needed because the underlying type errors have been resolved.
Note that it will also aim to ensure type checking runs cleanly, and will add suppressions to existing type errors.

#pyreupgrade

Reviewed By: shobhitmehro

Differential Revision: D114771938

fbshipit-source-id: 793b1012fc85e21b07e0632b5f2efe39d347a56a

Author: maggiemoss

gdpa/src/activation_utils.py

@@ -27,18 +27,15 @@
 E8M0_NEUTRAL_SCALE: int = 127
 
 
-# pyre-ignore[56]
 @dsl_user_op
 def pack_4xu8_to_u32(
-    # pyre-ignore[11]
     b0: Uint8,
     b1: Uint8,
     b2: Uint8,
     b3: Uint8,
     *,
     loc: object | None = None,
     ip: object | None = None,
-    # pyre-ignore[11]
 ) -> Uint32:
     """Pack 4 Uint8 values into a Uint32 (little-endian: b0 is lowest byte)."""
     return Uint32(
@@ -71,10 +68,8 @@ def pack_4xu8_to_u32(
     )
 
 
-# pyre-ignore[56]
 @dsl_user_op
 def store_u32_shared(
-    # pyre-ignore[11]
     ptr: Int64,
     val: Uint32,
     *,
@@ -96,10 +91,8 @@ def store_u32_shared(
     )
 
 
-# pyre-ignore[56]
 @dsl_user_op
 def tanh(
-    # pyre-ignore[11]
     a: float | Float32,
     *,
     loc: object | None = None,
@@ -118,7 +111,6 @@ def tanh(
     )
 
 
-# pyre-ignore[56]
 @dsl_user_op
 def max_f32(
     a: float | Float32,
@@ -141,7 +133,6 @@ def max_f32(
     )
 
 
-# pyre-ignore[56]
 @dsl_user_op
 def step_f32(
     x: float | Float32, *, loc: object | None = None, ip: object | None = None
@@ -162,7 +153,6 @@ def step_f32(
     )
 
 
-# pyre-ignore[56]
 @dsl_user_op
 def abs_f32(
     val: Float32, *, loc: object | None = None, ip: object | None = None
@@ -181,7 +171,6 @@ def abs_f32(
     )
 
 
-# pyre-ignore[56]
 @dsl_user_op
 def max3_f32(
     a: float | Float32,
@@ -213,7 +202,6 @@ def max3_f32(
     )
 
 
-# pyre-ignore[56]
 @dsl_user_op
 def mul_cvt_relu_8x_e4m3(
     in_0: Float32,
@@ -290,7 +278,6 @@ def mul_cvt_relu_8x_e4m3(
     )
 
 
-# pyre-ignore[56]
 @dsl_user_op
 def unpack_i64_to_u32_pair(
     packed: Int64, *, loc: object | None = None, ip: object | None = None
@@ -314,7 +301,6 @@ def unpack_i64_to_u32_pair(
     return lo, hi
 
 
-# pyre-ignore[56]
 @dsl_user_op
 def cvt_relu_8x_e4m3(
     s0: Float32,
@@ -377,7 +363,6 @@ def cvt_relu_8x_e4m3(
     )
 
 
-# pyre-ignore[56]
 @dsl_user_op
 def cvt_8x_e4m3(
     s0: Float32,
@@ -440,7 +425,6 @@ def cvt_8x_e4m3(
     )
 
 
-# pyre-ignore[56]
 @dsl_user_op
 def fused_amax_to_e8m0_scale_f32(
     amax: Float32,
@@ -524,7 +508,6 @@ class Relu:
     ReLU gradient: 1 if x >= 0, else 0
     """
 
-    # pyre-ignore[11]
     def __init__(
         self,
         scale_qk: Float32,
@@ -533,19 +516,16 @@ def __init__(
         self.c_zero: object = (Float32(0.0), Float32(0.0))
         self.c_one: object = (Float32(1.0), Float32(1.0))
 
-    # pyre-ignore[56]
     @cute.jit
     def relu(self, x: tuple[Float32, Float32]) -> tuple[Float32, Float32]:
         """Apply ReLU activation: max(0, x) for packed f32x2."""
         return (max_f32(x[0], Float32(0.0)), max_f32(x[1], Float32(0.0)))
 
-    # pyre-ignore[56]
     @cute.jit
     def grad_relu(self, x: tuple[Float32, Float32]) -> tuple[Float32, Float32]:
         """Compute ReLU gradient: 1 if x >= 0, else 0 for packed f32x2."""
         return (step_f32(x[0]), step_f32(x[1]))
 
-    # pyre-ignore[56]
     @cute.jit
     def activation_and_gradient_relu(
         self, x: tuple[Float32, Float32]
@@ -556,17 +536,13 @@ def activation_and_gradient_relu(
         """
         act = self.relu(x)
         grad = self.grad_relu(x)
-        # pyre-ignore[60]
         return *act, *grad
 
-    # pyre-ignore[56]
     @cute.jit
     def relu_and_convert(
         self,
-        # pyre-ignore[11]
         acc_S_row: cute.Tensor,
         acc_S_row_converted: cute.Tensor,
-        # pyre-ignore[11]
         e2e_freq: cutlass.Constexpr[int] = 16,
         e2e_res: cutlass.Constexpr[int] = 4,
         e2e_frg_limit: cutlass.Constexpr[int] = 1,
@@ -591,7 +567,6 @@ def relu_and_convert(
                 acc_S_row_frg[None, j].load().to(acc_S_row_converted.element_type)
             )
 
-    # pyre-ignore[56]
     @cute.jit
     def grad_relu_and_convert(
         self,
@@ -631,7 +606,6 @@ def grad_relu_and_convert(
                 acc_P_row_frg[None, j].load().to(acc_P_row_f16_frg.element_type)
             )
 
-    # pyre-ignore[56]
     @cute.jit
     def relu_and_convert_blockscaled(
         self,
@@ -821,7 +795,6 @@ def relu_and_convert_blockscaled(
 
 
 class Gelu:
-    # pyre-ignore[11]
     def __init__(
         self,
         scale_qk: Float32,
@@ -870,25 +843,18 @@ def __init__(
             Float32(0.0001171766272366646),
         )
         self.c_talyor_2_mul_c2: object = (
-            # pyre-ignore[16]
             Float32(self.c_taylor_c2[0] * self.c_two[0]),
-            # pyre-ignore[16]
             Float32(self.c_taylor_c2[1] * self.c_two[1]),
         )
         self.c_taylor_4_mul_c4: object = (
-            # pyre-ignore[16]
             Float32(self.c_taylor_c4[0] * self.c_four[0]),
-            # pyre-ignore[16]
             Float32(self.c_taylor_c4[1] * self.c_four[1]),
         )
         self.c_taylor_6_mul_c6: object = (
-            # pyre-ignore[16]
             Float32(self.c_taylor_c6[0] * self.c_six[0]),
-            # pyre-ignore[16]
             Float32(self.c_taylor_c6[1] * self.c_six[1]),
         )
 
-    # pyre-ignore[56]
     @cute.jit
     def gelu_tanh(self, x: tuple[Float32, Float32]) -> tuple[Float32, Float32]:
         # x^2
@@ -906,7 +872,6 @@ def gelu_tanh(self, x: tuple[Float32, Float32]) -> tuple[Float32, Float32]:
 
         return (t_x, t_y)
 
-    # pyre-ignore[56]
     @cute.jit
     def activation_and_gradient_fast_gelu(
         self, x: tuple[Float32, Float32]
@@ -925,10 +890,8 @@ def activation_and_gradient_fast_gelu(
         grad = cute.arch.fma_packed_f32x2(tanh_1, self.c_half, term1)
         act = cute.arch.mul_packed_f32x2(half_x, tanh_1)
 
-        # pyre-ignore[60]
         return *act, *grad
 
-    # pyre-ignore[56]
     @cute.jit
     def fast_gelu(self, x: tuple[Float32, Float32]) -> tuple[Float32, Float32]:
         # TODO: Seems like the complier may be able to handle this, without explicitly calling packed_f32x2:
@@ -940,7 +903,6 @@ def fast_gelu(self, x: tuple[Float32, Float32]) -> tuple[Float32, Float32]:
         out = cute.arch.mul_packed_f32x2(out, tanh_1)
         return out
 
-    # pyre-ignore[56]
     @cute.jit
     def grad_fast_gelu(self, x: tuple[Float32, Float32]) -> tuple[Float32, Float32]:
         # TODO: Seems like the complier may be able to handle this, without explicitly calling packed_f32x2:
@@ -960,7 +922,6 @@ def grad_fast_gelu(self, x: tuple[Float32, Float32]) -> tuple[Float32, Float32]:
         out = cute.arch.fma_packed_f32x2(one_plus_tanh, self.c_half, term1)
         return out
 
-    # pyre-ignore[56]
     @cute.jit
     def activation_and_gradient_gelu_taylor_deg6(
         self, x: tuple[Float32, Float32]
@@ -989,10 +950,8 @@ def activation_and_gradient_gelu_taylor_deg6(
 
         grad = cute.arch.add_packed_f32x2(part1, part2)
         act = cute.arch.fma_packed_f32x2(x2, tmp, x_half)  # 0.5*x + x2*tmp
-        # pyre-ignore[60]
         return *act, *grad
 
-    # pyre-ignore[56]
     @cute.jit
     def grad_gelu_taylor_deg6(
         self, x: tuple[Float32, Float32]
@@ -1020,7 +979,6 @@ def grad_gelu_taylor_deg6(
 
         return grad
 
-    # pyre-ignore[56]
     @cute.jit
     def gelu_taylor_deg6(self, x: tuple[Float32, Float32]) -> tuple[Float32, Float32]:
         # 0.5*x + x2*(c2 + x2*(c4 + x2*c6))
@@ -1033,7 +991,6 @@ def gelu_taylor_deg6(self, x: tuple[Float32, Float32]) -> tuple[Float32, Float32
         out = cute.arch.fma_packed_f32x2(x2, tmp, x_half)  # 0.5*x + x2*tmp
         return out
 
-    # pyre-ignore[56]
     @cute.jit
     def gelu_taylor_deg10(self, x: tuple[Float32, Float32]) -> tuple[Float32, Float32]:
         # 0.5*x + x^2*(c2 + x^2*(c4 + x^2*(c6 + x^2*(c8 + x^2*c10))))
@@ -1048,7 +1005,6 @@ def gelu_taylor_deg10(self, x: tuple[Float32, Float32]) -> tuple[Float32, Float3
         out = cute.arch.fma_packed_f32x2(x2, tmp, x_half)  # 0.5*x + x^2*tmp
         return out
 
-    # pyre-ignore[56]
     @cute.jit
     def gelu_and_convert(
         self,
@@ -1085,7 +1041,6 @@ def gelu_and_convert(
                 acc_S_row_frg[None, j].load().to(acc_S_row_converted.element_type)
             )
 
-    # pyre-ignore[56]
     @cute.jit
     def grad_gelu_and_convert(
         self,
@@ -1135,7 +1090,6 @@ def grad_gelu_and_convert(
                 acc_P_row_frg[None, j].load().to(acc_P_row_f16_frg.element_type)
             )
 
-    # pyre-ignore[56]
     @cute.jit
     def gelu_and_convert_blockscaled(
         self,

gdpa/src/fast_math.py

@@ -29,8 +29,7 @@
 from cutlass.cutlass_dsl import dsl_user_op, T  # pyre-ignore[21]
 
 
-@cute.jit  # pyre-ignore[56]
-# pyre-ignore[11]
+@cute.jit
 def clz(x: Int32) -> Int32:
     """Count leading zeros in a 32-bit integer.
 
@@ -59,15 +58,14 @@ def find_log2(x: Int32) -> Int32:
     return a + ((x & (x - 1)) != 0)  # Round up, add 1 if not a power of 2.
 
 
-@dsl_user_op  # pyre-ignore[56]
-# pyre-ignore[11]
+@dsl_user_op
 def umulhi(
     a: Int32,
     b: Int32,
     *,
     loc: object | None = None,
     ip: object | None = None,
-) -> Uint32:  # pyre-ignore[11]
+) -> Uint32:
     """Unsigned multiply-high: returns the upper 32 bits of the 64-bit product a*b.
 
     Uses PTX mul.hi.u32 instruction.
@@ -92,14 +90,11 @@ class FastDivmod:
     with a multiply+shift on the GPU.
     """
 
-    # pyre-ignore[11]
     divisor: Int32
-    # pyre-ignore[11]
     multiplier: Uint32
-    # pyre-ignore[11]
     shift_right: Uint32
     _loc: object | None
-    _values_pos: list[int]  # pyre-ignore[13]
+    _values_pos: list[int]
 
     def __init__(
         self,
@@ -110,9 +105,9 @@ def __init__(
         loc: object | None = None,
         ip: object | None = None,
     ) -> None:
-        self.divisor = divisor  # pyre-ignore[4]
-        self.multiplier = multipler  # pyre-ignore[4]
-        self.shift_right = shift_right  # pyre-ignore[4]
+        self.divisor = divisor
+        self.multiplier = multipler
+        self.shift_right = shift_right
         self._loc = loc
 
     # called by host
@@ -132,15 +127,14 @@ def create(
         shift_right = Uint32(p - 32)
         return FastDivmod(divisor, multiplier, shift_right, loc=loc, ip=ip)
 
-    @cute.jit  # pyre-ignore[56]
+    @cute.jit
     def div(self, dividend: Int32) -> Int32:
         return (
             Int32(umulhi(dividend, self.multiplier) >> self.shift_right)
             if self.divisor != 1
             else dividend
         )
 
-    # pyre-ignore[11]
     def divmod(self, dividend: Int32) -> tuple[Int32, Int32]:
         quotient = self.div(dividend)
         remainder = dividend - quotient * self.divisor

gdpa/src/seqlen_info.py

@@ -41,10 +41,8 @@ class SeqlenInfo:
 
     def __init__(
         self,
-        # pyre-fixme[11]: Annotation `cutlass.Int32` is not defined as a type.
         batch_idx: cutlass.Int32,
         seqlen_static: cutlass.Int32,
-        # pyre-fixme[11]: Annotation `cute.Tensor` is not defined as a type.
         cu_seqlens: cute.Tensor | None = None,
         seqused: cute.Tensor | None = None,
     ) -> None: