fixing small bug caused in gemms and softmax

adamomainz · facebook-github-bot · commit ce74f5e7413e · 2024-12-03T09:45:09.000-08:00
Summary:
gemms are failing because of logical bug here

flash_attention does not use this function and directly calls the downstream function. When testing the change for shuffled data I only tested on flash attention and didnt see this break in testing.

Reviewed By: danzimm

Differential Revision: D66706511

fbshipit-source-id: f89a7561a41d82b8c486bd6af097057b2413b7d2
diff --git a/tritonbench/utils/data_utils.py b/tritonbench/utils/data_utils.py
@@ -1,14 +1,14 @@
 from .triton_op import IS_FBCODE
 
 
-def get_production_shapes(op_name, op_type):
+def get_production_shapes(op_name, op_type, shuffle_shapes=False):
     """Gets a list of Softmax shapes for benchmarking"""
     if IS_FBCODE:
         from .fb.durin_data import productionDataLoader
 
         return [
             shape
             for shape in productionDataLoader.get_shapes_from_frozen_durin(
-                op_name, op_type
+                op_name, op_type, shuffle_shapes
             )
         ]

Original file line number	Diff line number	Diff line change
`@@ -1,14 +1,14 @@`
`1`	`1`	`from .triton_op import IS_FBCODE`
`2`	`2`
`3`	`3`
`4`		`-def get_production_shapes(op_name, op_type):`
	`4`	`+def get_production_shapes(op_name, op_type, shuffle_shapes=False):`
`5`	`5`	`"""Gets a list of Softmax shapes for benchmarking"""`
`6`	`6`	`if IS_FBCODE:`
`7`	`7`	`from .fb.durin_data import productionDataLoader`
`8`	`8`
`9`	`9`	`return [`
`10`	`10`	`shape`
`11`	`11`	`for shape in productionDataLoader.get_shapes_from_frozen_durin(`
`12`		`- op_name, op_type`
	`12`	`+ op_name, op_type, shuffle_shapes`
`13`	`13`	`)`
`14`	`14`	`]`