fused_moegemm_tile_partitioner.hpp Source File

fused_moegemm_tile_partitioner.hpp Source File#

Composable Kernel: fused_moegemm_tile_partitioner.hpp Source 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