Pass in Compute kernel type in New plan generation (#3061)

Summary:
Pull Request resolved: #3061

This is tricky - a small edge case that can't be reproduced locally - OSS environment will expose where compute kernel may be detected as "fused" when not passed in.

This will cause the unit test to occasionally change the kernel type in resharding (due to the new plan passed in).

Passing in the kernel type in the new plan generation fixes this.

Reviewed By: aliafzal

Differential Revision: D76233463

fbshipit-source-id: 740bfb6619df2ee862b3bf84f23dec395e14bb83

Author: aporialiao

torchrec/distributed/test_utils/test_sharding.py

@@ -452,6 +452,9 @@ def dynamic_sharding_test(
 
         new_per_param_sharding = {}
 
+        assert len(sharders) == 1
+        # pyre-ignore
+        kernel_type = sharders[0]._kernel_type
         # Construct parameter shardings
         for i in range(num_tables):
             table_name = tables[i].name
@@ -466,7 +469,7 @@ def dynamic_sharding_test(
             )
             # TODO: CW sharding constructor takes in different args
             new_per_param_sharding[table_name] = sharding_type_constructor(
-                rank=new_ranks[i][0]
+                rank=new_ranks[i][0], compute_kernel=kernel_type
             )
 
         new_module_sharding_plan = construct_module_sharding_plan(