Enable Pyrefly for deeplearning/fbgemm

fbgemm_gpu/bench/jagged_tensor_benchmark.py

@@ -1089,6 +1089,7 @@ def keyed_jagged_index_select_dim1_ref(
         )
     if baseline:
         time_ref, _ = benchmark_torch_function(
+            # pyrefly: ignore [unbound-name]
             functools.partial(output_ref.backward, retain_graph=True),
             (grad,),
             iters=iters,

fbgemm_gpu/bench/merge_embeddings_benchmark.py

@@ -94,6 +94,7 @@ def generate_requests(
         rs.append(
             get_table_batched_offsets_from_dense(all_indices.view(T, B, L), gpu_num)
         )
+    # pyrefly: ignore [bad-return]
     return rs
 
 

fbgemm_gpu/bench/sparse_ops_benchmark.py

@@ -1402,6 +1402,7 @@ def batch_group_index_select_bwd(
         )
 
     if timeline:
+        # pyrefly: ignore [missing-attribute]
         prof.export_chrome_trace("index_select_fwd_trace.json")
 
     grads = [torch.rand_like(out) for out in out_pyt]
@@ -1442,6 +1443,7 @@ def batch_group_index_select_bwd(
         )
 
     if timeline:
+        # pyrefly: ignore [missing-attribute]
         prof.export_chrome_trace("index_select_bwd_trace.json")
 
     logging.info(
@@ -1722,10 +1724,13 @@ def permute_1d_sparse_data_bench(
     embedding tables.
     """
     if index_dtype == "int":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.int32
     elif index_dtype == "int64":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.int64
     elif index_dtype == "float":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.float32
     else:
         raise RuntimeError(f"Does not support data type {index_dtype}")
@@ -1744,8 +1749,10 @@ def permute_1d_sparse_data_bench(
     total_indices = int(lengths.sum().item())
     # Generate indices
     if index_dtype == torch.float32:
+        # pyrefly: ignore [no-matching-overload]
         indices = torch.rand(total_indices, dtype=index_dtype, device=device)
     else:
+        # pyrefly: ignore [no-matching-overload]
         indices = torch.randint(
             low=0,
             high=2**31 - 1,
@@ -1881,14 +1888,19 @@ def permute_2d_sparse_data_bench(
     systems to reorder embedding tables.
     """
     if index_dtype == "int":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.int32
     elif index_dtype == "int64":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.int64
     elif index_dtype == "float":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.float32
     elif index_dtype == "bf16":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.bfloat16
     elif index_dtype == "fp16":
+        # pyrefly: ignore [bad-assignment]
         index_dtype = torch.float16
     else:
         raise RuntimeError(f"Does not support data type {index_dtype}")
@@ -1909,9 +1921,12 @@ def permute_2d_sparse_data_bench(
         * emb_dim
     )
     total_indices = int(lengths.sum().item())
+    # pyrefly: ignore [missing-attribute]
     if index_dtype.is_floating_point:
+        # pyrefly: ignore [no-matching-overload]
         indices = torch.rand(total_indices, dtype=index_dtype, device=device)
     else:
+        # pyrefly: ignore [no-matching-overload]
         indices = torch.randint(
             low=0,
             high=2**31 - 1,

fbgemm_gpu/bench/tbe/split_table_batched_embeddings_benchmark.py

@@ -2047,6 +2047,7 @@ def device_from_files(  # noqa C901
     _Ds: list[int]
     emb_loc: list[EmbeddingLocation]
     compute_device: list[ComputeDevice]
+    # pyrefly: ignore [bad-assignment]
     _Es, _Ds, emb_loc, compute_device = zip(*embedding_specs)
 
     # Determine location suffix for trace URL based on embedding locations

fbgemm_gpu/bench/tbe/tbe_cache_benchmark.py

@@ -272,12 +272,14 @@ def lxu_cache_lookup(
     tbe: nn.Module = IntNBitTableBatchedEmbeddingBagsCodegen(
         embedding_specs, cache_load_factor=cache_load_factor
     )
+    # pyrefly: ignore [not-callable]
     tbe.fill_random_weights()
 
     # Imitate execution flow by performing prefetching once.
     indices, offsets = create_request(
         num_tables, num_embeddings, batch, avg_pooling_factor
     )
+    # pyrefly: ignore [not-callable]
     tbe.prefetch(indices, offsets)
 
     linearized_indices = torch.ops.fbgemm.linearize_cache_indices(
@@ -333,6 +335,7 @@ def lru_cache_populate_byte(
     cc: nn.Module = IntNBitTableBatchedEmbeddingBagsCodegen(
         embedding_specs, cache_load_factor=cache_load_factor
     )
+    # pyrefly: ignore [not-callable]
     cc.fill_random_weights()
 
     warm_up_requests = []
@@ -389,6 +392,7 @@ def populate(linear_indices: Tensor) -> None:
     replay_cc: nn.Module = IntNBitTableBatchedEmbeddingBagsCodegen(
         embedding_specs, cache_load_factor=cache_load_factor
     )
+    # pyrefly: ignore [not-callable]
     replay_cc.fill_random_weights()
 
     replay_timestep: int = 1
@@ -464,6 +468,7 @@ def lfu_cache_populate_byte(
         cache_load_factor=cache_load_factor,
         cache_algorithm=CacheAlgorithm.LFU,
     )
+    # pyrefly: ignore [not-callable]
     cc.fill_random_weights()
 
     warm_up_requests = []
@@ -517,6 +522,7 @@ def populate(linear_indices: Tensor) -> None:
         cache_load_factor=cache_load_factor,
         cache_algorithm=CacheAlgorithm.LFU,
     )
+    # pyrefly: ignore [not-callable]
     replay_cc.fill_random_weights()
 
     def replay_populate(linear_indices: Tensor) -> None:

fbgemm_gpu/bench/tbe/tbe_inference_benchmark.py

@@ -874,6 +874,7 @@ def nbit_device_with_spec(  # noqa C901
                 TBERequest(
                     req.indices.cpu().int(),
                     req.offsets.cpu().int(),
+                    # pyrefly: ignore [missing-attribute]
                     req.per_sample_weigths.cpu() if req.per_sample_weights else None,
                 )
                 for req in requests
@@ -1197,6 +1198,7 @@ def nbit_uvm(
     )
 
     if T_gpu > 0:
+        # pyrefly: ignore [unbound-name]
         nparams_byte = sum(w.numel() for (w, _) in emb_mixed.split_embedding_weights())
         logging.info(
             f"{weights_precision} Embedding tables: {E * T_gpu + E_uvm * T_uvm} rows, {nparams_byte / param_size_multiplier / 1.0e9: .2f} GParam, "
@@ -1647,6 +1649,7 @@ def nbit_cache(  # noqa C901
             )
             for d in Ds
         ],
+        # pyrefly: ignore [not-callable]
         record_cache_metrics=RecordCacheMetrics(
             record_cache_miss_counter, record_tablewise_cache_miss
         ),

fbgemm_gpu/bench/tbe/tbe_training_benchmark.py

@@ -639,6 +639,7 @@ def device_with_speclist(  # noqa C901
     # pyre-ignore[53]
     def _kineto_trace_handler(p: profile, phase: str) -> None:
         p.export_chrome_trace(
+            # pyrefly: ignore [missing-attribute]
             benchconfig.trace_url.format(
                 emb_op_type=emb_op_type, phase=phase, ospid=os.getpid()
             )

fbgemm_gpu/codegen/genscript/generate_backward_split.py

@@ -26,6 +26,7 @@
     from .optimizer_args import annotation_dict, OptimizerArgsSet
     from .scripts_argsparse import args
 except ImportError:
+    # pyrefly: ignore [missing-import]
     from optimizers import *
 
     # pyre-ignore[21]
@@ -118,6 +119,7 @@ def generate_backward_split_gpu(**kwargs: Any) -> None:
         ]:
             BackwardSplitGenerator.render_backward_templates(
                 template_filepath,
+                # pyrefly: ignore [bad-argument-type]
                 optimizer,
                 filename_format,
                 kwargs,
@@ -141,6 +143,7 @@ def generate_backward_split_gpu(**kwargs: Any) -> None:
             ]:
                 BackwardSplitGenerator.render_backward_templates(
                     template_filepath,
+                    # pyrefly: ignore [bad-argument-type]
                     optimizer,
                     filename_format,
                     kwargs,
@@ -241,7 +244,7 @@ def generate_backward_split_cpu(**kwargs: Any) -> None:
         if kwargs.get("has_cpu_support"):
             CodeTemplate.load(
                 "training/backward/embedding_backward_split_cpu_approx_template.cpp"
-                if "approx" in optimizer
+                if "approx" in optimizer  # pyrefly: ignore [not-iterable]
                 else "training/backward/embedding_backward_split_cpu_template.cpp"
             ).write(f"gen_embedding_backward_{optimizer}_split_cpu.cpp", **kwargs)
 

fbgemm_gpu/codegen/genscript/generate_embedding_optimizer.py

@@ -36,6 +36,7 @@ def generate_embedding_optimizer(**kwargs: Any) -> None:
 
         optimizer = kwargs.get("optimizer")
         kwargs["optimizer_class_name"] = "".join(
+            # pyrefly: ignore [missing-attribute]
             [optim.capitalize() for optim in optimizer.split("_")]
         )
         kwargs["args"] = kwargs["args"].cuda

fbgemm_gpu/codegen/genscript/jinja_environment.py

@@ -58,20 +58,26 @@
 # BT_block_size * 4 * 4 * 32 * (max_D // 128) <= 64 * 1024 (V100) or 96 * 1024 (A100)
 # Since BT_block_size >= 1, max_D <= 16K (V100) or 24K (A100).
 # Note that if we increase max_D, it will increase the compilation time significantly.
+# pyrefly: ignore [unsupported-operation]
 env.globals["max_embedding_dim"] = 2048
 
 # Max embedding dimension for legacy embedding kernels. TBE v2 can support
 # larger max embedding dimension.
+# pyrefly: ignore [unsupported-operation]
 env.globals["legacy_max_embedding_dim"] = 1024
 
 # An optimization for ROCm
+# pyrefly: ignore [unsupported-operation]
 env.globals["items_per_warp"] = 128 if args.is_rocm is False else 256
 
 # The fixed max vectors per thread for different kernels.  The numbers were
 # derived from empirical studies
+# pyrefly: ignore [unsupported-operation]
 env.globals["fixed_max_vecs_per_thread"] = {"backward": 2, "backward_indice_weights": 6}
 
+# pyrefly: ignore [unsupported-operation]
 env.globals["dense"] = False
+# pyrefly: ignore [unsupported-operation]
 env.globals["is_rocm"] = args.is_rocm
 
 
@@ -343,15 +349,24 @@ def compute_global_weight_decay(is_global_weight_decay_kernel: bool) -> str:
 ################################################################################
 
 env.globals["generate_optimized_grad_sum_loop_access"] = (
+    # pyrefly: ignore [unsupported-operation]
     generate_optimized_grad_sum_loop_access
 )
+# pyrefly: ignore [unsupported-operation]
 env.globals["get_max_vecs_template_configs"] = get_max_vecs_template_configs
+# pyrefly: ignore [unsupported-operation]
 env.globals["dispatch_optimal_kernel"] = dispatch_optimal_kernel
+# pyrefly: ignore [unsupported-operation]
 env.globals["dispatch_non_vec_blocking_kernel"] = dispatch_non_vec_blocking_kernel
+# pyrefly: ignore [unsupported-operation]
 env.globals["dispatch_vec_blocking_kernel"] = dispatch_vec_blocking_kernel
+# pyrefly: ignore [unsupported-operation]
 env.globals["is_valid_forward_config"] = is_valid_forward_config
+# pyrefly: ignore [unsupported-operation]
 env.globals["has_experimental_support"] = has_experimental_support
+# pyrefly: ignore [unsupported-operation]
 env.globals["is_valid_gwd_config"] = is_valid_gwd_config
+# pyrefly: ignore [unsupported-operation]
 env.globals["compute_global_weight_decay"] = compute_global_weight_decay
 
 ################################################################################

fbgemm_gpu/experimental/gemm/triton_gemm/fp4_quantize.py

@@ -5650,4 +5650,5 @@ def quantize_nvfp4_naive(
         )
     )
 
+    # pyrefly: ignore [bad-return]
     return xqs, x_scales

fbgemm_gpu/experimental/gemm/triton_gemm/grouped_gemm.py

@@ -337,6 +337,7 @@ def _fbgemm_grouped_gemm(
             if USE_TMA_STORE:
                 c_desc_ptr = tl.make_tensor_descriptor(
                     c_ptr + M_start_offset * N,
+                    # pyrefly: ignore [bad-argument-type]
                     shape=[m_size, n_size],
                     # pyre-ignore
                     strides=[n_size, 1],
@@ -347,6 +348,7 @@ def _fbgemm_grouped_gemm(
                 tl.static_assert(K % BLOCK_SIZE_K == 0)
                 a_desc_ptr = tl.make_tensor_descriptor(
                     a_ptr + M_start_offset * K,
+                    # pyrefly: ignore [bad-argument-type]
                     shape=[m_size, K],
                     # pyre-ignore
                     strides=[K, 1],
@@ -374,7 +376,9 @@ def _fbgemm_grouped_gemm(
                     m_offset = (tile_m_idx * BLOCK_SIZE_M).to(tl.int32)
                     n_offset = (tile_n_idx * BLOCK_SIZE_N).to(tl.int32)
                     for k_offset in range(0, K, BLOCK_SIZE_K):
+                        # pyrefly: ignore [bad-argument-type, unbound-name]
                         a = a_desc_ptr.load([m_offset, k_offset])
+                        # pyrefly: ignore [bad-argument-type, unbound-name]
                         b = b_desc_ptr.load([n_offset, k_offset])
                         if USE_FAST_ACCUM:
                             accumulator = tl.dot(a, b.T, accumulator)
@@ -512,6 +516,7 @@ def _fbgemm_grouped_gemm_ws(
             if USE_TMA_STORE:
                 c_desc_ptr = tl.make_tensor_descriptor(
                     c_ptr + M_start_offset * N,
+                    # pyrefly: ignore [bad-argument-type]
                     shape=[m_size, N],
                     # pyre-ignore
                     strides=[N, 1],
@@ -659,6 +664,7 @@ def _fbgemm_grouped_gemm_fp8_rowwise(
             if USE_TMA_STORE:
                 c_desc_ptr = tl.make_tensor_descriptor(
                     c_ptr + M_start_offset * N,
+                    # pyrefly: ignore [bad-argument-type]
                     shape=[m_size, n_size],
                     # pyre-ignore
                     strides=[n_size, 1],
@@ -668,6 +674,7 @@ def _fbgemm_grouped_gemm_fp8_rowwise(
             if USE_TMA_LOAD:
                 a_desc_ptr = tl.make_tensor_descriptor(
                     a_ptr + M_start_offset * K,
+                    # pyrefly: ignore [bad-argument-type]
                     shape=[m_size, K],
                     # pyre-ignore
                     strides=[K, 1],
@@ -695,7 +702,9 @@ def _fbgemm_grouped_gemm_fp8_rowwise(
                     m_offset = (tile_m_idx * BLOCK_SIZE_M).to(tl.int32)
                     n_offset = (tile_n_idx * BLOCK_SIZE_N).to(tl.int32)
                     for k_offset in range(0, K, BLOCK_SIZE_K):
+                        # pyrefly: ignore [bad-argument-type, unbound-name]
                         a = a_desc_ptr.load([m_offset, k_offset])
+                        # pyrefly: ignore [bad-argument-type, unbound-name]
                         b = b_desc_ptr.load([n_offset, k_offset])
                         if USE_FAST_ACCUM:
                             accumulator = tl.dot(a, b.T, accumulator)
@@ -835,6 +844,7 @@ def _fbgemm_grouped_gemm_fp8_rowwise_ws(
             if USE_TMA_STORE:
                 c_desc_ptr = tl.make_tensor_descriptor(
                     c_ptr + M_start_offset * N,
+                    # pyrefly: ignore [bad-argument-type]
                     shape=[m_size, N],
                     # pyre-ignore
                     strides=[N, 1],

fbgemm_gpu/experimental/gen_ai/gen_ai/moe/layers.py

@@ -1070,6 +1070,7 @@ def _route(
         scores = torch.sigmoid(scores)
         assert scores.shape == (B * T, self.E)
 
+        # pyrefly: ignore [not-callable]
         token_counts, expert_indices, token_indices = index_shuffling(
             scores,  # num_tokens
         )

fbgemm_gpu/experimental/gen_ai/src/quantize/common/scripts/make_heuristic.py

@@ -88,6 +88,7 @@ def get_kernel_assignment(file_path: str, threshold: float) -> dict[ProblemShape
         for row in file:
             problem_shape_str, kernel, time_ms_str = row.strip().split(",")
             shape_tuple = tuple(int(x) for x in problem_shape_str.split("_"))
+            # pyrefly: ignore [bad-argument-type]
             problem_shape = ProblemShape.from_tuple(shape_tuple)
             time_ms = float(time_ms_str)
 
@@ -110,6 +111,7 @@ def get_kernel_assignment(file_path: str, threshold: float) -> dict[ProblemShape
 
     # Prefer kernels that are used more often across all problem shapes
     while len(kernel_assignment) < len(kernel_candidates):
+        # pyrefly: ignore [no-matching-overload]
         top_kernel = sorted(kernel_count.keys(), key=kernel_count.get, reverse=True)[0]
         for problem_shape, candidates in kernel_candidates.items():
             if problem_shape not in kernel_assignment and top_kernel in candidates:

fbgemm_gpu/experimental/gen_ai/test/attention/bert_padding.py

@@ -27,6 +27,7 @@ def forward(ctx, input, indices):
         ).reshape(-1, *other_shape)
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def backward(ctx, grad_output):
         (indices,) = ctx.saved_tensors
         assert grad_output.ndim >= 2
@@ -63,6 +64,7 @@ def forward(ctx, values, indices, first_axis_dim):
         return output
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def backward(ctx, grad_output):
         (indices,) = ctx.saved_tensors
         # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
@@ -88,6 +90,7 @@ def forward(ctx, input, indices):
         return output, input.detach()
 
     @staticmethod
+    # pyrefly: ignore [bad-override]
     def backward(ctx, grad_output, grad_residual):
         (indices,) = ctx.saved_tensors
         assert grad_output.ndim >= 2