Implementation:Mlc ai Mlc llm InternLM2 Model

Knowledge Sources	Mlc_ai_Mlc_llm
Domains	Model_Architecture, LLM
Last Updated	2026-02-09 19:00 GMT

Overview

Implements the InternLM2 architecture for causal language modeling within the MLC LLM framework, featuring grouped-query attention with a fused QKV projection and SiLU-gated MLP.

Description

This module provides the TVM Relax-based implementation of the InternLM2 model architecture developed by Shanghai AI Laboratory. The architecture follows a standard decoder-only transformer design with several notable characteristics:

• Fused QKV projection: Uses a single wqkv linear layer that projects to all query, key, and value heads simultaneously, improving computation efficiency.
• Grouped-query attention (GQA): Supports fewer key-value heads than query heads, reducing memory requirements while maintaining model quality.
• SiLU-gated MLP: Uses a fused gate_up_proj that combines gate and up projections, followed by SiLU activation and a w2 down-projection.
• RoPE positional embeddings: Uses rotary position embeddings with configurable rope_theta.
• Configurable bias: Attention projections optionally include bias terms, controlled by the bias config parameter. The tensor parallel sharding strategy adapts accordingly to also shard bias vectors.
• RMSNorm: Uses RMSNorm for both attention and feed-forward normalization layers.

The model stack consists of InternLM2Model (token embeddings + decoder layers + final RMSNorm), wrapped by InternLM2ForCausalLM which adds the output projection head.

Usage

Use this module when compiling InternLM2 or InternLM2.5 family models for deployment with MLC LLM. The model is identified by the internlm2 model type in configuration files.

Code Reference

Source Location

• Repository: Mlc_ai_Mlc_llm
• File: python/mlc_llm/model/internlm2/internlm2_model.py

Signature

@dataclasses.dataclass
class InternLM2Config(ConfigBase):
    vocab_size: int
    hidden_size: int
    num_hidden_layers: int
    num_attention_heads: int
    num_key_value_heads: int
    rms_norm_eps: float
    intermediate_size: int
    bias: bool
    use_cache: bool
    rope_theta: int
    pad_token_id: int
    bos_token_id: int
    eos_token_id: int
    context_window_size: int = 0
    prefill_chunk_size: int = 0
    tensor_parallel_shards: int = 1
    head_dim: int = 0
    ...

class InternLM2ForCausalLM(nn.Module):
    def __init__(self, config: InternLM2Config): ...
    def embed(self, input_ids: Tensor): ...
    def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): ...
    def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): ...
    def batch_prefill(self, input_embeds, logit_positions, paged_kv_cache): ...
    def batch_decode(self, input_embeds, paged_kv_cache): ...
    def batch_verify(self, input_embeds, paged_kv_cache): ...
    def create_paged_kv_cache(self, ...): ...
    def get_default_spec(self): ...

Import

from mlc_llm.model.internlm2.internlm2_model import InternLM2Config, InternLM2ForCausalLM

I/O Contract

Primary Classes

Class	Role	Key Characteristics
InternLM2Config	Model configuration	Includes explicit bias flag, pad/bos/eos token IDs
InternLM2Attention	GQA attention	Fused wqkv projection, optional bias, wo output projection
InternLM2MLP	Gated feed-forward	SiLU-gated via fused gate_up_proj + w2 down projection
InternLM2DecoderLayer	Transformer block	attention_norm + attention + ffn_norm + feed_forward with residual connections
InternLM2Model	Core model	tok_embeddings + layers + norm
InternLM2ForCausalLM	Top-level model	Adds output linear head for causal LM

Forward Methods

Method	Input	Output
embed	Tensor[seq_len] (int32)	Tensor[seq_len, hidden_size]
prefill	Tensor[1, seq_len, hidden_size], PagedKVCache	(Tensor[1, 1, vocab_size], PagedKVCache)
decode	Tensor[1, 1, hidden_size], PagedKVCache	(Tensor[1, 1, vocab_size], PagedKVCache)
batch_prefill	Tensor[1, seq_len, hidden_size], Tensor[batch_size], PagedKVCache	(Tensor, PagedKVCache)
batch_decode	Tensor[batch_size, 1, hidden_size], PagedKVCache	(Tensor, PagedKVCache)

Tensor Parallel Sharding

Parameter	Shard Strategy	Dimension
wqkv.weight	ShardSingleDim with segs=[q, k, v]	dim=0
wqkv.bias (if bias=True)	ShardSingleDim with segs=[q, k, v]	dim=0
wo.weight	ShardSingleDim	dim=1
gate_up_proj.weight	ShardSingleDim with segs=[i, i]	dim=0
w2.weight	ShardSingleDim	dim=1

Usage Examples

# Creating an InternLM2 config
config = InternLM2Config(
    vocab_size=92544,
    hidden_size=4096,
    num_hidden_layers=32,
    num_attention_heads=32,
    num_key_value_heads=8,
    rms_norm_eps=1e-5,
    intermediate_size=14336,
    bias=False,
    use_cache=True,
    rope_theta=1000000,
    pad_token_id=2,
    bos_token_id=1,
    eos_token_id=2,
    context_window_size=32768,
)
model = InternLM2ForCausalLM(config)

Related Pages

• Implementation:Mlc_ai_Mlc_llm_Model_Preset
• Implementation:Mlc_ai_Mlc_llm_Llama_Model