[DeepSeek][Kernels] add on device, parallelized permute_indices kernel (#1077)

This PR updates the original Triton kernel in indices.py
(#1062) to move it to a more
parallelized and vectorized implementation to better support larger
runs.
This provides roughly an 8.7x perf improvement for 256 experts. 

This kernel functionally achieves the same as before:
Prepare permutation indices and the number of tokens for each expert
with any required alignment, where alignment is driven by the group gemm
requirement.
* The permutation indices are the indices of the tokens for each expert.
* The number of tokens for each expert is the sum of the number of
tokens for
such experts from all ranks. 

However, this kernel improves the parallelization by assigning one block
to each expert, and threads in each block process the tokens in chunks,
all while using vectorized loads and stores.
This results in the performance improvement while yielding the same
indices.

Testing:
~~~
python moe_kernels.py
  ~~~

(runs simple verification using cpu, previous kernel and new optimized
kernel)

Benchmark of performance improvement using fixed 256 experts across
various world sizes
~~~

Summary for fixed total experts (256), where we mimic various DeepSeek
experts across a number of ranks:
Ranks | Experts/Rank | Original (ms) | Optimized (ms) | Block Size |
Speedup | Correct

------------------------------------------------------------------------------------------
1 | 256 | 0.241 | 0.020 | 128 | 11.95x | 1
4 | 64 | 0.222 | 0.020 | 128 | 10.97x | 1
16 | 16 | 0.221 | 0.021 | 128 | 10.40x | 1
32 | 8 | 0.224 | 0.025 | 64 | 8.81x | 1
64 | 4 | 0.225 | 0.037 | 128 | 6.01x | 1
128 | 2 | 0.217 | 0.051 | 128 | 4.26x | 1
  ~~~

added unit testing: 
~~~
..Testing 256 experts per rank across 1 ranks (total: 256)
Testing 128 experts per rank across 2 ranks (total: 256)
Testing 64 experts per rank across 4 ranks (total: 256)
Testing 32 experts per rank across 8 ranks (total: 256)
Testing 16 experts per rank across 16 ranks (total: 256)
Testing 8 experts per rank across 32 ranks (total: 256)
Testing 4 experts per rank across 64 ranks (total: 256)
Testing 2 experts per rank across 128 ranks (total: 256)
.
----------------------------------------------------------------------
Ran 3 tests in 1.104s

OK
~~~

Author: lessw2020

torchtitan/experiments/deepseek_v3/moe_kernels.py

@@ -0,0 +1,247 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import triton
+import triton.language as tl
+
+
+__all__ = ["generate_permute_indices"]
+
+
+# parallelized kernel
+@triton.jit
+def _fill_indices_kernel(
+    tokens_per_expert_group_ptr,
+    start_index_values_ptr,
+    write_offsets_ptr,
+    output_ptr,
+    experts_per_rank: tl.constexpr,
+    num_ranks: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,  # Number of threads per block
+):
+    pid = tl.program_id(axis=0)
+    num_programs = tl.num_programs(axis=0)
+
+    # map programs (blocks) to the experts and loop (grid stride) if needed
+    for expert_id in range(pid, experts_per_rank, num_programs):
+
+        # read this experts write offset
+        write_offset = tl.load(write_offsets_ptr + expert_id)
+
+        # loop over all ranks
+        for r in range(num_ranks):
+            # index into tokens_per_expert_group array
+            i = r * experts_per_rank + expert_id
+
+            # load start index and number of tokens for this expert-rank pair
+            start_index = tl.load(start_index_values_ptr + i)
+            length = tl.load(tokens_per_expert_group_ptr + i)
+
+            # each thread in block processes tokens in parallel
+            offsets = tl.arange(0, BLOCK_SIZE)
+
+            # tokens are processed in chunks of BLOCK_SIZE
+            for chunk_start in range(0, length, BLOCK_SIZE):
+                chunk_offsets = chunk_start + offsets
+
+                # mask valid indices
+                mask = chunk_offsets < length
+
+                values = start_index + chunk_offsets
+
+                # destination
+                dest_indices = write_offset + chunk_offsets
+
+                # store
+                tl.store(output_ptr + dest_indices, values, mask=mask)
+
+            # update write offset for next rank
+            write_offset += length
+
+
+# ==============
+# wrapper
+# ==============
+
+
+def fill_indices_wrapper(
+    tokens_per_expert_group: torch.Tensor,
+    start_index_values: torch.Tensor,
+    write_offsets: torch.Tensor,
+    experts_per_rank: int,
+    num_ranks: int,
+    max_len: int,
+    block_size: int = 128,
+    max_blocks: int = 1024,  # cap on total number of blocks to launch
+):
+    # preallocate output
+    permuted_indices = torch.full(
+        (max_len,), -1, dtype=torch.int32, device=tokens_per_expert_group.device
+    )
+
+    # write offsets is per local expert...
+    num_blocks = min(experts_per_rank, max_blocks)
+    # grid = one block per expert unless capped and then we loop...
+    grid = (num_blocks,)
+
+    # launch kernel
+    _fill_indices_kernel[grid](
+        tokens_per_expert_group,
+        start_index_values,
+        write_offsets,
+        permuted_indices,
+        experts_per_rank,
+        num_ranks,
+        BLOCK_SIZE=block_size,
+    )
+    return permuted_indices
+
+
+# reference
+def fill_indices_cpu(
+    tokens_per_expert_group: torch.Tensor,
+    start_index_values: torch.Tensor,
+    write_offsets: torch.Tensor,
+    experts_per_rank: int,
+    num_ranks: int,
+    max_len: int,
+):
+    # We need to preallocate the output - we ignore device and force it on cpu
+    # device = tokens_per_expert_group.device
+    permuted_indices = torch.full(
+        (max_len,),
+        -1,
+        dtype=torch.int32,
+    )  # device=device)
+    # Fill the permuted indices
+    # For each local expert
+    for e in range(experts_per_rank):
+        write_start = write_offsets[e].item()
+        # For each remote rank
+        for r in range(num_ranks):
+            i = r * experts_per_rank + e
+            start_index = start_index_values[i].item()
+            length = tokens_per_expert_group[i].item()
+            # Fill in the indices
+            if length > 0:
+                end_idx = min(write_start + length, max_len)
+                permuted_indices[write_start:end_idx] = torch.arange(
+                    start_index,
+                    start_index + (end_idx - write_start),
+                    dtype=torch.int32,
+                    # device=device,
+                )
+            write_start += length
+    return permuted_indices
+
+
+def generate_permute_indices(
+    tokens_per_expert_group: torch.Tensor,
+    experts_per_rank: int,
+    num_ranks: int,
+    max_len: int,
+    alignment: int,
+    use_cpu: bool = False,
+):
+    """
+    Prepare permutation indices and the number of tokens for each expert.
+
+    Args:
+        tokens_per_expert_group: number of tokens for each expert from all ranks.
+        experts_per_rank: number of experts per rank.
+        num_ranks: number of ranks.
+        max_len: maximum length of the output index vector.
+        alignment: alignment for each returned element in `m_sizes`.
+        use_cpu: whether to use CPU implementation.
+        use_optimized: whether to use optimized Triton implementation.
+        block_size: block size for optimized implementation.
+
+    Returns:
+        permuted_indices: permutation indices.
+        m_sizes: number of tokens for each expert.
+    """
+    # prefix sum to get start index of each expert (parallel scan kernel in future?)
+    start_index_values = (
+        torch.cumsum(tokens_per_expert_group, 0) - tokens_per_expert_group
+    )
+
+    # chunk sizes for each expert
+    chunk_size_per_expert = tokens_per_expert_group.view(num_ranks, -1).sum(0)
+
+    # align the chunk sizes (cdiv)
+    m_sizes = ((chunk_size_per_expert + alignment - 1) // alignment * alignment).to(
+        torch.int32
+    )
+
+    # additional prefix sum to get write offset of each expert in permuted_indices
+    # write offsets is per local expert, not global
+    write_offsets = torch.cumsum(m_sizes, 0) - m_sizes
+
+    # Select the implementation to use
+    if use_cpu:
+        permuted_indices = fill_indices_cpu(
+            tokens_per_expert_group,
+            start_index_values,
+            write_offsets,
+            experts_per_rank,
+            num_ranks,
+            max_len,
+        )
+    else:
+        permuted_indices = fill_indices_wrapper(
+            tokens_per_expert_group,
+            start_index_values,
+            write_offsets,
+            experts_per_rank,
+            num_ranks,
+            max_len,
+        )
+
+    return permuted_indices, m_sizes
+
+
+# Below is for testing only
+
+
+def simple_test():
+    device = torch.device("cuda", 0)
+    experts_per_rank = 4
+    num_ranks = 4
+    tokens_per_expert_group = torch.full(
+        (num_ranks * experts_per_rank,), 4, dtype=torch.int32, device=device
+    )
+    max_len = 128
+    alignment = 32
+    # Use the GPU kernel
+    permuted_indices_gpu, m_sizes = generate_permute_indices(
+        tokens_per_expert_group, experts_per_rank, num_ranks, max_len, alignment
+    )
+    # Use the CPU method
+    permuted_indices_cpu, _ = generate_permute_indices(
+        tokens_per_expert_group,
+        experts_per_rank,
+        num_ranks,
+        max_len,
+        alignment,
+        use_cpu=True,
+    )
+    # Check that the results are the same
+
+    assert torch.equal(permuted_indices_gpu.cpu(), permuted_indices_cpu)
+    assert torch.equal(
+        torch.remainder(m_sizes, alignment),
+        torch.zeros(experts_per_rank, device=device),
+    )
+    # Print the results
+    print(f"{permuted_indices_gpu=}, \n{permuted_indices_cpu=}")
+    print(f"{m_sizes=}")
+    print("Success")
+    return True  # assert would have failed meaning getting here is success.
+
+
+if __name__ == "__main__":
+    simple_test()