Implementation:InternLM Lmdeploy Impl Simt

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

Overview

SIMT (non-tensor-core) attention implementation that uses scalar FMA instructions for QK and PV computation, supporting quantized KV caches and serving as the fallback for small GQA group sizes.

Description

This is the MMA_SIMT specialization of Impl, which avoids tensor core MMA instructions and instead performs attention via explicit dot-product loops with warp-level shuffle reductions. Each thread processes a VEC=8 element chunk of the head dimension, with T_D=8 threads covering the D dimension and T_S=4 threads covering the S dimension within each warp. The implementation supports quantized KV caches by loading raw data into DataK/DataV fragments and converting on-the-fly using ConvertKvCache with per-token scale/zero parameters from SmemLayoutKVp. Cross-warp reduction in Merge uses shared memory staging for both O and M/L. This implementation works on all architectures (SM70+) and handles CTA_H > 1 for multi-head grouping.

Usage

Selected by DecodingConfig for small GQA group sizes (Qh <= 2) on SM80, for all decoding on SM70 (Volta), and as a fallback for head dimensions (e.g., 192) not supported by the tensor core implementations.

Code Reference

Source Location

• Repository: InternLM_Lmdeploy
• File: src/turbomind/kernels/attention/impl_simt.h
• Lines: 1-674

Signature

namespace turbomind::attention {

template<class T_, class Tkv_, 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_SIMT, T_, Tkv_, CTA_H_, CTA_Q_, CTA_S_,
            WARP_H_, WARP_Q, WARP_S, HeadDim, Stages> {

    using T = T_;
    using Tkv = Tkv_;
    static constexpr int kQuantKV = !std::is_same_v<T, Tkv>;
    static constexpr int VEC = 8;
    static constexpr int T_D = 8;
    static constexpr int T_S = WARP_SIZE / T_D;

    using FragQ = Array<T, VEC>[K_M][K_K];
    using FragS = Array<float, 1>[K_M][K_N];
    using FragO = Array<float, VEC>[V_M][V_N];
    using FragM = Array<float, 1>[K_M];
    using FragL = FragM;

    using DataK = FragK_<Tkv>;
    using DataV = FragV_<Tkv>;
    using ParamK = Array<T, 2>[K_N];
    using ParamV = Array<T, 2>[V_K];

    union SharedStorage {
        T Q[SmemLayoutQ::kSize];
        struct { Array<Tkv, Stages*SmemLayoutK::kSize> KV; T KVp[...]; };
        struct { SmemM M; SmemL L; SmemO O; };
    };

    struct StateQK { ... };
    struct StatePV { ... };

    static void ComputeQK(StateQK, FragS&, int, auto&&, auto&&);
    static void ComputePV(StatePV, FragO&, int, auto&&, auto&&);
    static void Softmax<is_residue>(FragS&, FragM&, FragL&, FragO&, float);
    static void Merge(FragO&, FragM&, FragL&, float, SharedStorage&);
    static void StoreO<is_norm>(FragO&, FragL&, SharedStorage&, Func&&);
};

} // namespace turbomind::attention

Import

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

I/O Contract

Inputs

Name	Type	Required	Description
T_	typename	Yes	Compute type (half, bfloat16)
Tkv_	typename	Yes	KV cache type (half, uint8_t, uint4_t, fp8_e4m3_t, fp4_e2m1_t)
CTA_H_	int	Yes	Number of query heads per CTA
CTA_S_	int	Yes	CTA sequence tile size (typically 64)
HeadDim	int	Yes	Head dimension
Stages	int	Yes	Pipeline stages

Outputs

Name	Type	Description
FragO	Array<float,VEC>[V_M][V_N]	Accumulated output per head
FragM	Array<float,1>[K_M]	Per-head row max
FragL	Array<float,1>[K_M]	Per-head row sum

Usage Examples

// SIMT decoding for GQA group size <= 2 on SM80
using Attention = Impl<MMA_SIMT, half, half, 2, 1, 64, 2, 1, 16, 128, 3>;
using Kernel = AttentionUniversal<arch::Sm80,
    Mainloop<Sm80_CpAsync<3>, Attention>,
    GetBlockIterFactory<half, half, 64, 128>,
    DecodingCtaMap>;

Related Pages

• Environment:InternLM_Lmdeploy_CUDA_GPU_Runtime