[PyTorch] Fuse permute+pad and unpermute+unpad ops for FP8 optimization (#1921)

* [PyTorch] Fuse permute+pad and unpermute+unpad ops for FP8 optimization

1.Fused `moe_permute_with_probs` + `Fp8Padding` and fused `moe_unpermute` + `Fp8Unpadding`,
  that can remove the explicit padding/unpadding of moe expert, improved performance and reduced peak gpu memory usage.
2.Add tests of fused permute/pad and unpermute/unpad.

Signed-off-by: xiaoxi-wangfj <[email protected]>

* [PyTorch/Common] Fuse permute+pad and unpermute+unpad support with_merging_probs

Signed-off-by: xiaoxi-wangfj <[email protected]>

* [PyTorch]format code

Signed-off-by: xiaoxi-wangfj <[email protected]>

* [Common]perf expert_idx loaded once

Signed-off-by: xiaoxi-wangfj <[email protected]>

* fix: pad_offsets can be None

Co-authored-by: Tim Moon <[email protected]>
Signed-off-by: xiaoxi-wangfj <[email protected]>

* add padding + merging probs bwd support. Not tested

Signed-off-by: tdophung <[email protected]>

* Fix garbage initialized act grad

Signed-off-by: tdophung <[email protected]>

* all test passing for jax permutation + pad

Signed-off-by: tdophung <[email protected]>

* change tokens_per_experts APIs to num_out_tokens with conservative allocation of worst case padding for output buffer

Signed-off-by: tdophung <[email protected]>

* change test permutation to reduce test time

Signed-off-by: tdophung <[email protected]>

* triggering PR refresh

Signed-off-by: tdophung <[email protected]>

* format code

Signed-off-by: tdophung <[email protected]>

* Remove some tests cases from pytorch side. Add a separate toekn_dispatch test for sanity in case combine accidentally undo an error on dispatch in the roundtrip test. Add distinction between L0 and L2 in test cases in jax

Signed-off-by: tdophung <[email protected]>

* format code

Signed-off-by: tdophung <[email protected]>

* remove chance for inefficiency in moving between CPU and GPU, remove redundant primitive using a new static bool for padding, add assert for align size

Signed-off-by: tdophung <[email protected]>

* fix lint in jax

Signed-off-by: tdophung <[email protected]>

* account for both jax newer and older than version 0.8.2. Adjusted gpu triton binding accordingly

Signed-off-by: tdophung <[email protected]>

* format code

Signed-off-by: tdophung <[email protected]>

* fix typo

Signed-off-by: tdophung <[email protected]>

---------

Signed-off-by: xiaoxi-wangfj <[email protected]>
Signed-off-by: tdophung <[email protected]>
Co-authored-by: Tim Moon <[email protected]>
Co-authored-by: tdophung <[email protected]>

Author: 3 people

tests/jax/test_permutation.py

@@ -200,6 +200,7 @@ def _permute_kernel(
     probs_ptr,
     scale_ptr,
     permuted_scale_ptr,
+    pad_offsets_ptr,
     # sizes
     scale_hidden_dim,
     # strides
@@ -224,8 +225,11 @@ def _permute_kernel(
     hidden_size: tl.constexpr,
     PERMUTE_PROBS: tl.constexpr,
     PERMUTE_SCALE: tl.constexpr,
+    FUSION_PAD: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
+    expert_idx = 0
+
     pid_t = tl.program_id(0)
     pid_h = tl.program_id(1)
     cur_off = pid_h * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
@@ -246,18 +250,22 @@ def _permute_kernel(
         dst_row = tl.load(
             row_id_map_ptr + pid_t * stride_row_id_map_token + idx * stride_row_id_map_expert
         ).to(tl.int64)
+        if FUSION_PAD or PERMUTE_PROBS:
+            expert_idx = tl.load(
+                row_id_map_ptr
+                + pid_t * stride_row_id_map_token
+                + (num_experts + idx) * stride_row_id_map_expert
+            )
+        if FUSION_PAD:
+            pad_off = tl.load(pad_offsets_ptr + expert_idx)
+            dst_row = dst_row + pad_off
         output_off = dst_row * stride_output_token + cur_off * stride_output_hidden
         if PERMUTE_SCALE:
             permuted_scale_off = (
                 dst_row * stride_permuted_scale_token + cur_off * stride_permuted_scale_hidden
             )
             tl.store(permuted_scale_ptr + permuted_scale_off, scale, mask=mask_scale)
         if PERMUTE_PROBS:
-            expert_idx = tl.load(
-                row_id_map_ptr
-                + pid_t * stride_row_id_map_token
-                + (num_experts + idx) * stride_row_id_map_expert
-            )
             prob_off = pid_t * stride_probs_token + expert_idx * stride_probs_expert
             prob = tl.load(probs_ptr + prob_off)
             if pid_h == 0:
@@ -297,6 +305,7 @@ def _unpermute_kernel(
     row_id_map_ptr,
     merging_probs_ptr,
     permuted_probs_ptr,
+    pad_offsets_ptr,
     # strides
     stride_row_id_map_token,
     stride_row_id_map_expert,
@@ -318,10 +327,12 @@ def _unpermute_kernel(
     PROBS_LOAD_WIDTH: tl.constexpr,
     WITH_MERGING_PROBS: tl.constexpr,
     PERMUTE_PROBS: tl.constexpr,
+    FUSION_UNPAD: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
     data_type = input_ptr.dtype.element_ty
     compute_type = tl.float32
+    expert_idx = 0
 
     pid_t = tl.program_id(0)
     pid_h = tl.program_id(1)
@@ -348,15 +359,19 @@ def _unpermute_kernel(
         src_row = tl.load(
             row_id_map_ptr + pid_t * stride_row_id_map_token + idx * stride_row_id_map_expert
         ).to(tl.int64)
-        input_off = src_row * stride_input_token + current_offset * stride_input_hidden
-        inp = tl.load(input_ptr + input_off, mask=mask)
-        inp = inp.to(compute_type)
-        if WITH_MERGING_PROBS:
+        if FUSION_UNPAD or WITH_MERGING_PROBS:
             expert_idx = tl.load(
                 row_id_map_ptr
                 + pid_t * stride_row_id_map_token
                 + (num_experts + idx) * stride_row_id_map_expert
             )
+        if FUSION_UNPAD:
+            pad_off = tl.load(pad_offsets_ptr + expert_idx)
+            src_row = src_row + pad_off
+        input_off = src_row * stride_input_token + current_offset * stride_input_hidden
+        inp = tl.load(input_ptr + input_off, mask=mask)
+        inp = inp.to(compute_type)
+        if WITH_MERGING_PROBS:
             merging_prob_off = (
                 pid_t * stride_merging_probs_token + expert_idx * stride_merging_probs_expert
             )
@@ -407,6 +422,7 @@ def _unpermute_bwd_with_merging_probs_kernel(
     fwd_input_ptr,
     merging_probs_ptr,
     row_id_map_ptr,
+    pad_offsets_ptr,
     # strides
     stride_row_id_map_token,
     stride_row_id_map_expert,
@@ -427,6 +443,7 @@ def _unpermute_bwd_with_merging_probs_kernel(
     num_experts: tl.constexpr,
     hidden_size: tl.constexpr,
     PROBS_LOAD_WIDTH: tl.constexpr,
+    FUSION_UNPAD: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
     data_type = fwd_output_grad_ptr.dtype.element_ty
@@ -450,6 +467,9 @@ def _unpermute_bwd_with_merging_probs_kernel(
             + pid * stride_row_id_map_token
             + (num_experts + idx) * stride_row_id_map_expert
         )
+        if FUSION_UNPAD:
+            pad_off = tl.load(pad_offsets_ptr + expert_idx)
+            dst_row = dst_row + pad_off
         prob_grad_accum = tl.zeros((BLOCK_SIZE,), dtype=compute_type)
         current_start = 0
         while current_start < hidden_size: