weighs added to logs from production data!!

Summary: Weights are now being pulled from production data + a bit of refactoring

Reviewed By: nmacchioni

Differential Revision: D65637832

fbshipit-source-id: 078aaed94f340f8bc00adbe21e433bc57f810b6f

Author: adamomainz

run.py

@@ -83,6 +83,10 @@ def _run(args: argparse.Namespace, extra_args: List[str]) -> BenchmarkOperatorRe
                 "device": args.device,
                 "logging_group": args.logging_group,
             }
+            if args.production_shapes:
+                from tritonbench.utils.fb.durin_data import productionDataLoader
+
+                kwargs["weights_loader"] = productionDataLoader
 
             if "hardware" in args:
                 kwargs["hardware"] = args.hardware

tritonbench/operators/flash_attention/operator.py

@@ -427,7 +427,6 @@ def get_ctx_vals():
         shapes = self.__additional_example_input(ctx_vals)
         requires_grad = True
         for shape in shapes:
-            print(shape)
             BATCH, H, N_CTX, D_HEAD = shape
             q = torch.randn(
                 (BATCH, H, N_CTX, D_HEAD),
@@ -459,9 +458,14 @@ def __additional_example_input(self, standard_shapes: Generator) -> Generator:
         ]
         shapes = chain(standard_shapes, llama_shapes)
         if self.add_production_shapes:
-            from ...utils.fb.durin_data import get_shapes_from_frozen_durin
+            from ...utils.fb.durin_data import productionDataLoader
 
-            shapes = chain(shapes, get_shapes_from_frozen_durin("attention"))
+            shapes = chain(
+                shapes,
+                productionDataLoader.get_shapes_from_frozen_durin(
+                    self.name, "attention"
+                ),
+            )
         return shapes
 
     @register_x_val(label="(Batch, Heads, SeqLen, Dhead)")