update block size for standalone_cint_v4

Differential Revision: D66350756
facebookresearch · Nov 22, 2024 · ae4905f · ae4905f
1 parent 55e346c
commit ae4905f
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/generative_recommenders/ops/triton/triton_jagged.py b/generative_recommenders/ops/triton/triton_jagged.py
@@ -1166,6 +1166,8 @@ def _get_concat_2D_jagged_tritoncc_named_specs() -> List[VersionedSpec]:
                 (256, "*fp32"),
                 (512, "*bf16"),
                 (512, "*fp32"),
+                (1024, "*bf16"),
+                (1024, "*fp32"),
             ]
             for offsets_b_type in ["*i64", "*i32"]
             for IS_DENSE_A, IS_DENSE_B in [(False, False), (True, False), (False, True)]
@@ -1306,6 +1308,8 @@ def _get_concat_2D_jagged_tritoncc_named_specs() -> List[VersionedSpec]:
                 (256, "*fp32"),
                 (512, "*bf16"),
                 (512, "*fp32"),
+                (1024, "*bf16"),
+                (1024, "*fp32"),
             ]
             for offsets_b_type in ["*i64", "*i32"]
             for IS_DENSE_A, IS_DENSE_B in [(False, False), (True, False), (False, True)]
@@ -1502,7 +1506,7 @@ def _get_split_2D_jagged_tritoncc_named_specs() -> List[VersionedSpec]:
                 },
                 default_values=default_values,
             )
-            for BLOCK_D in [64, 128, 256, 512]
+            for BLOCK_D in [64, 128, 256, 512, 1024]
             for dtype in ["*bf16", "*fp32"]
             for offsets_a_type in ["*i64", "*i32"]
             for offsets_b_type in ["*i64", "*i32"]
@@ -1556,7 +1560,7 @@ def _get_split_2D_jagged_tritoncc_named_specs() -> List[VersionedSpec]:
                 default_values=default_values,
                 version="standalone_cint_v4",
             )
-            for BLOCK_D in [64, 128, 256, 512]
+            for BLOCK_D in [64, 128, 256, 512, 1024]
             for dtype in ["*bf16", "*fp32"]
             for offsets_a_type in ["*i64", "*i32"]
             for offsets_b_type in ["*i64", "*i32"]