Implementation:InternLM Lmdeploy Gemm Epilogue

Knowledge Sources	InternLM_Lmdeploy
Domains	GPU_Kernels, GEMM
Last Updated	2026-02-07 15:00 GMT

Overview

Implements the GEMM epilogue pipeline that rearranges MMA accumulator fragments through shared memory, applies scaling/combination operations, supports split-K reduction, gated SiLU activation, and stores the final output.

Description

The Epilogue_ template orchestrates the post-MMA phase of a GEMM kernel. Key components:

• Rearrange: Transfers MMA accumulator fragments (FragC) through shared memory using RearrangeC to convert from MMA register layout to a thread-contiguous tile layout suitable for global memory stores.

• MatrixCombination_v3: Implements D = alpha * C_accum + beta * C_input using the C matrix from global memory, with support for flat, blocked, and indexed striding modes.

• ChannelCombination_v3: Applies per-channel scale and bias to the output (used for channel-wise quantization).

• Split-K Reduction: When multiple CTAs collaborate on the K dimension, uses semaphore-based barriers (sem_wait/sem_post) and a partials buffer to accumulate partial results across splits.

• GatedActivation<Silu>: Applies gated SiLU activation silu(x[2i]) * x[2i+1], halving the output dimension.

• EpilogueParam: Bundles output matrix parameters, partial accumulation buffers, barrier pointers, combination parameters, and the SiLU activation flag.

Usage

Instantiated as part of GemmUniversal and invoked after the mainloop completes. The epilogue handles all post-accumulation work before writing to global memory.

Code Reference

Source Location

• Repository: InternLM_Lmdeploy
• File: src/turbomind/kernels/gemm/epilogue.h

Signature

template<class Tc_, int M, int N, int TM_, int TN_, int THREADS,
         class RearrangeC, class OperandC, Striding mode_C, bool SplitK_>
struct Epilogue_ {
    void operator()(FragC& frag_C, const int4& tile_offset,
                    const int2& extents, int splits, int tile_id,
                    bool is_last, const EpilogueParam& param,
                    SharedStorage& storage);
};

struct EpilogueParam {
    MatrixParam c, partials;
    int* locks;
    MatrixCombination_v3 combine_mat;
    bool silu_act;
};

Import

#include "src/turbomind/kernels/gemm/epilogue.h"

I/O Contract

Inputs

Name	Type	Required	Description
frag_C	FragC&	Yes	MMA accumulator fragments from the mainloop
tile_offset	int4	Yes	(tile_m, tile_n, split_id, group_id)
extents	int2	Yes	Valid (M, N) extents within this tile
splits	int	Yes	Total number of split-K partitions
param	EpilogueParam	Yes	Output buffers, scaling, activation config

Outputs

Name	Type	Description
D matrix	global memory	Final GEMM output written to global memory
partials	global memory	Intermediate partial sums for split-K (if not last)

Usage Examples

// Called at the end of GemmUniversal::operator()
Epilogue epilogue{};
epilogue(frag_C, tile_offset, extents, splits, tile_id, is_last, epi_param, storage.epilogue);

Related Pages

• Environment:InternLM_Lmdeploy_CUDA_GPU_Runtime