Implementation:InternLM Lmdeploy Impl 16816

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

Overview

Attention implementation using m16n8k16 tensor core MMA instructions (Ampere+), providing QK and PV compute stages with multi-stage pipeline support and swizzled shared memory layouts.

Description

This is the MMA_16816 specialization of Impl, targeting SM80 (Ampere) and newer GPUs. It defines fragment types for Q, K, S (scores), P (probabilities), V, and O in the m16n8k16 MMA atom layout. Shared memory layouts use SmemLayoutV2 with 3-bit swizzle patterns to avoid bank conflicts. The implementation provides: TransformQ to load Q from shared memory via LDSM instructions; StateQK/ComputeQK for the Q*K^T dot product; StatePV/ComputePV for the P*V accumulation; and a Merge step for cross-warp L reduction. It supports both 2-stage and 3-stage pipelining through the Stages template parameter.

Usage

Selected by AttentionConfig for SM80 prefill attention. Combined with a mainloop (e.g., Sm80_CpAsync<2>) and wrapped in AttentionUniversal.

Code Reference

Source Location

• Repository: InternLM_Lmdeploy
• File: src/turbomind/kernels/attention/impl_16816.h
• Lines: 1-303

Signature

namespace turbomind::attention {

template<class T_, int CTA_H_, int CTA_Q_, int CTA_S_,
         int WARP_H, int WARP_Q, int WARP_S, int HeadDim, int Stages>
struct Impl<MMA_16816, T_, T_, CTA_H_, CTA_Q_, CTA_S_,
            WARP_H, WARP_Q, WARP_S, HeadDim, Stages>
    : Impl_m16k8<T_, WARP_H, WARP_Q, WARP_S, HeadDim> {

    using T = T_;
    using Tkv = T_;
    static constexpr int kHeadDim = HeadDim;
    static constexpr int CTA_H = CTA_H_;
    static constexpr int CTA_Q = CTA_Q_;
    static constexpr int CTA_S = CTA_S_;
    static constexpr int OP_K = 16;

    using FragQ = Array<T, 8>[K_K][K_M];
    using FragK = Array<T, 4>[K_K][K_N];
    using FragP = Array<T, 8>[V_M][V_K];
    using FragV = Array<T, 4>[V_K][V_N];

    union SharedStorage { ... };
    struct StateQK { ... };
    struct StatePV { ... };

    static void TransformQ(T* smem_Q, FragQ& frag_Q);
    static void ComputeQK(StateQK, FragS&, int, auto&&, auto&&);
    static void ComputePV(StatePV, FragO&, int, auto&&, auto&&);
    static void Merge(FragO&, FragM&, FragL&, float, auto&);
};

} // namespace turbomind::attention

Import

#include "src/turbomind/kernels/attention/impl_16816.h"

I/O Contract

Inputs

Name	Type	Required	Description
T_	typename	Yes	Data type (half or bfloat16)
CTA_Q_	int	Yes	CTA tile size in query dimension (typically 64)
CTA_S_	int	Yes	CTA tile size in sequence/key dimension (typically 64)
HeadDim	int	Yes	Head dimension (64, 128, 192, 256)
Stages	int	Yes	Number of pipeline stages (2 or 3)

Outputs

Name	Type	Description
FragO	Array<float,4>[V_M][V_N]	Accumulated output fragment in registers
FragM	Array<float,2>[V_M]	Running row-wise maximum for online softmax
FragL	Array<float,2>[V_M]	Running row-wise sum for online softmax

Usage Examples

// Used through AttentionConfig:
using Attention = Impl<MMA_16816, half, half, 1, 64, 64, 1, 16, 64, 128, 2>;
// Provides StateQK, StatePV, ComputeQK, ComputePV, Softmax, etc.

Related Pages

• Environment:InternLM_Lmdeploy_CUDA_GPU_Runtime