refactor the kernel

Signed-off-by: xiaoyao0115 <[email protected]>
NVIDIA · Nov 19, 2024 · fd83104 · fd83104
1 parent 35c9836
commit fd83104
Show file tree

Hide file tree

Showing 2 changed files with 225 additions and 166 deletions.
diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
@@ -458,8 +458,6 @@ void thd_grad_correction(at::Tensor grad, const at::Tensor &grad_per_step,
 at::Tensor thd_get_partitioned_indices(const at::Tensor &cu_seqlens, int total_tokens,
                                        int world_size, int rank);
 
-__forceinline__ __device__ int binary_search(int target, int *array, int len);
-
 /***************************************************************************************************
  * multi_tensor_* kernels
  **************************************************************************************************/