/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/fused_moe/kernel/fused_moegemm_tile_partitioner.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/fused_moe/kernel/fused_moegemm_tile_partitioner.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/fused_moe/kernel/fused_moegemm_tile_partitioner.hpp Source File

Go to the documentation of this file.

 // SPDX-License-Identifier: MIT
 // Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
  
 #pragma once
  
 namespace ck_tile {
  
 template <typename BlockShape_>
 struct FusedMoeGemmTilePartitioner_Linear
 {
     //  FusedMoeGemmShape
     using BlockShape = ck_tile::remove_cvref_t<BlockShape_>;
  
     static constexpr const char* name = "lin";
  
     CK_TILE_DEVICE auto operator()(ck_tile::index_t /*num_sorted_tiles*/,
                                    ck_tile::index_t /*intermediate_size*/)
     {
         index_t i_n = blockIdx.x;
         index_t i_m = blockIdx.y;
  
         return ck_tile::make_tuple(i_m, i_n);
     }
  
     CK_TILE_HOST static constexpr auto GridSize(index_t max_tokens, index_t intermediate_size)
     {
         // TODO: this may need tuning
         index_t ms = ck_tile::integer_divide_ceil(max_tokens, BlockShape::Block_M0);
         index_t ns = ck_tile::integer_divide_ceil(intermediate_size, BlockShape::Block_N0);
         return dim3(ns, ms, 1);
     }
 };
 } // namespace ck_tile