深入解读Qwen3技术报告（三）：深入剖析Qwen3模型架构

重磅推荐专栏：
《大模型AIGC》
《课程大纲》
《知识星球》

本专栏致力于探索和讨论当今最前沿的技术趋势和应用领域，包括但不限于ChatGPT和Stable Diffusion等。我们将深入研究大型模型的开发和应用，以及与之相关的人工智能生成内容（AIGC）技术。通过深入的技术解析和实践经验分享，旨在帮助读者更好地理解和应用这些领域的最新进展

3. 深入剖析Qwen3模型架构

大型语言模型的架构设计直接决定了其性能上限和计算效率。Qwen3在继承前代模型优势的基础上，引入了多项架构创新，使其在保持高性能的同时，实现了更高的计算效率和更强的扩展性。本章将深入剖析Qwen3的模型架构，包括基础组件、密集模型与MoE模型的设计差异，以及分词器实现等技术细节，帮助读者理解Qwen3强大能力背后的架构秘密。

3.1 Qwen3的基础架构组件

与大多数现代大型语言模型一样，Qwen3的基础架构建立在Transformer解码器的基础上，但引入了多项改进和创新。下面我们将逐一解析Qwen3的核心架构组件。

3.1.1 整体架构概览

Qwen3的整体架构遵循自回归语言模型的典型设计，由多层Transformer解码器堆叠而成。每一层包含以下主要组件：

自注意力机制（Self-Attention）：允许模型关注输入序列中的不同位置，捕捉长距离依赖关系。
前馈网络（Feed-Forward Network）：由两个线性变换和一个非线性激活函数组成，增强模型的表示能力。
层归一化（Layer Normalization）：稳定训练过程，加速收敛。
残差连接（Residual Connection）：缓解梯度消失问题，便于训练深层网络。

上图展示了Qwen3的整体架构和单个Transformer层的内部结构。与传统Transformer不同，Qwen3采用了层前归一化（Pre-Layer Normalization）设计，即在每个子层（自注意力和前馈网络）之前应用层归一化，而不是之后。这种设计有助于稳定训练过程，特别是对于深层模型。

3.1.2 改进的自注意力机制

Qwen3在自注意力机制上引入了多项改进，以提高性能和效率：

分组查询注意力（Grouped Query Attention, GQA）：

传统的多头注意力机制（Multi-Head Attention, MHA）为每个注意力头分配独立的查询（Q）、键（K）和值（V）投影。而GQA则让多个查询头共享同一组键值头，显著减少了参数量和计算量，同时保持了性能。

# 传统多头注意力与GQA的对比实现
class MultiHeadAttention(nn.Module):def __init__(self, d_model, num_heads):super().__init__()self.d_model = d_modelself.num_heads = num_headsself.head_dim = d_model // num_heads# 传统MHA: 每个头有独立的QKV投影self.q_proj = nn.Linear(d_model, d_model)self.k_proj = nn.Linear(d_model, d_model)self.v_proj = nn.Linear(d_model, d_model)self.out_proj = nn.Linear(d_model, d_model)def forward(self, x):batch_size, seq_len, _ = x.shape# 投影并分头q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)k = self.k_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)v = self.v_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)# 计算注意力scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)attn_weights = F.softmax(scores, dim=-1)attn_output = torch.matmul(attn_weights, v)# 合并头并投影attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)output = self.out_proj(attn_output)return outputclass GroupedQueryAttention(nn.Module):def __init__(self, d_model, num_query_heads, num_kv_heads):super().__init__()self.d_model = d_modelself.num_query_heads = num_query_headsself.num_kv_heads = num_kv_heads  # 通常 num_kv_heads < num_query_headsself.head_dim = d_model // num_query_heads# GQA: 查询头数量多于键值头数量self.q_proj = nn.Linear(d_model, d_model)self.k_proj = nn.Linear(d_model, self.num_kv_heads * self.head_dim)self.v_proj = nn.Linear(d_model, self.num_kv_heads * self.head_dim)self.out_proj = nn.Linear(d_model, d_model)def forward(self, x):batch_size, seq_len, _ = x.shape# 投影q = self.q_proj(x).view(batch_size, seq_len, self.num_query_heads, self.head_dim).transpose(1, 2)k = self.k_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)v = self.v_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)# 键值重复以匹配查询头数量# 每个键值头被多个查询头共享if self.num_kv_heads != self.num_query_heads:# 计算每个键值头需要复制的次数kv_repeat_factor = self.num_query_heads // self.num_kv_heads# 复制键值头k = k.repeat_interleave(kv_repeat_factor, dim=1)v = v.repeat_interleave(kv_repeat_factor, dim=1)# 计算注意力scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)attn_weights = F.softmax(scores, dim=-1)attn_output = torch.matmul(attn_weights, v)# 合并头并投影attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)output = self.out_proj(attn_output)return output

在Qwen3中，不同规模的模型采用了不同的GQA配置：

模型	查询头数量	键值头数量	头数比例
Qwen3-0.6B	8	4	2:1
Qwen3-1.7B	16	4	4:1
Qwen3-4B	32	8	4:1
Qwen3-8B	32	8	4:1
Qwen3-14B	40	10	4:1
Qwen3-32B	64	8	8:1
Qwen3 MoE系列	128	16	8:1

这种设计大大减少了模型的参数量和计算量，同时保持了性能。

QK-Norm：

Qwen3移除了传统Transformer中的QKV-bias（查询、键、值投影的偏置项），并引入了QK-Norm技术。QK-Norm对查询和键向量进行归一化，使得注意力分数的分布更加稳定，有助于提高模型的稳定性和性能。

# QK-Norm实现
def attention_with_qk_norm(q, k, v, scale=1.0):"""带QK-Norm的注意力计算参数:q: 查询张量 [batch_size, num_heads, seq_len, head_dim]k: 键张量 [batch_size, num_heads, seq_len, head_dim]v: 值张量 [batch_size, num_heads, seq_len, head_dim]scale: 缩放因子返回:注意力输出"""# 对查询和键进行L2归一化q_normalized = F.normalize(q, p=2, dim=-1)k_normalized = F.normalize(k, p=2, dim=-1)# 计算注意力分数scores = torch.matmul(q_normalized, k_normalized.transpose(-2, -1)) * scale# 应用softmaxattn_weights = F.softmax(scores, dim=-1)# 计算输出output = torch.matmul(attn_weights, v)return output

QK-Norm的引入使得Qwen3在处理长序列和复杂任务时更加稳定，减少了训练过程中的异常情况。

旋转位置编码（Rotary Positional Embedding, RoPE）：

Qwen3采用了RoPE作为位置编码方法，它通过对查询和键向量应用旋转变换，将位置信息直接编码到自注意力计算中。与传统的位置编码相比，RoPE具有更好的相对位置感知能力和外推性。

# RoPE实现
def apply_rotary_pos_emb(q, k, cos, sin, position_ids):"""应用旋转位置编码参数:q, k: 查询和键张量cos, sin: 余弦和正弦位置编码position_ids: 位置ID返回:应用了RoPE的查询和键"""# 获取位置编码cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]# 将q和k分成实部和虚部q_real, q_imag = q[..., ::2], q[..., 1::2]k_real, k_imag = k[..., ::2], k[..., 1::2]# 应用复数乘法的旋转q_rotated_real = q_real * cos - q_imag * sinq_rotated_imag = q_real * sin + q_imag * cosk_rotated_real = k_real * cos - k_imag * sink_rotated_imag = k_real * sin + k_imag * cos# 重新组合实部和虚部q_rotated = torch.stack([q_rotated_real, q_rotated_imag], dim=-1).flatten(-2)k_rotated = torch.stack([k_rotated_real, k_rotated_imag], dim=-1).flatten(-2)return q_rotated, k_rotated

在Qwen3中，为了支持更长的上下文，RoPE的基频从传统的10,000增加到了1,000,000，并结合YARN（Yet Another RoPE extensioN）技术，使模型能够处理长达128K tokens的超长序列。

3.1.3 前馈网络与激活函数

Qwen3的前馈网络（Feed-Forward Network, FFN）采用了标准的两层结构，但在激活函数和设计细节上有所创新：

SwiGLU激活函数：

Qwen3采用了SwiGLU（Swish-Gated Linear Unit）作为前馈网络的激活函数，它是对传统GLU（Gated Linear Unit）的改进，使用Swish函数替代了Sigmoid函数，提供了更好的梯度流和性能。

# SwiGLU激活函数实现
class SwiGLU(nn.Module):def __init__(self, in_features, hidden_features, out_features):super().__init__()self.gate_proj = nn.Linear(in_features, hidden_features, bias=False)self.up_proj = nn.Linear(in_features, hidden_features, bias=False)self.down_proj = nn.Linear(hidden_features, out_features, bias=False)def forward(self, x):# 计算门控值gate = self.gate_proj(x)# 应用Swish激活: x * sigmoid(beta * x)，这里beta=1gate = gate * torch.sigmoid(gate)# 上投影up = self.up_proj(x)# 门控机制hidden = gate * up# 下投影return self.down_proj(hidden)

SwiGLU激活函数使得Qwen3在训练过程中收敛更快，并在各种任务上取得更好的性能。

隐藏层维度：

Qwen3的前馈网络隐藏层维度通常是模型维度的4倍，这种设计提供了足够的模型容量，使模型能够学习复杂的模式和关系。

模型模型维度 FFN隐藏维度比例
Qwen3-0.6B 1024 4096 4:1
Qwen3-1.7B 1536 6144 4:1
Qwen3-4B 2560 10240 4:1
Qwen3-8B 3584 14336 4:1
Qwen3-14B 4608 18432 4:1
Qwen3-32B 6144 24576 4:1
Qwen3 MoE系列 8192 32768 4:1
这种设计在保持模型表达能力的同时，也考虑了计算效率和训练稳定性。

模型	模型维度	FFN隐藏维度	比例
Qwen3-0.6B	1024	4096	4:1
Qwen3-1.7B	1536	6144	4:1
Qwen3-4B	2560	10240	4:1
Qwen3-8B	3584	14336	4:1
Qwen3-14B	4608	18432	4:1
Qwen3-32B	6144	24576	4:1
Qwen3 MoE系列	8192	32768	4:1

3.2 密集模型与MoE模型的架构差异

Qwen3系列包含两种类型的模型：密集模型（Dense Models）和混合专家模型（Mixture-of-Experts Models, MoE）。这两种模型在架构上有显著差异，下面我们将详细比较它们的设计特点。

3.2.1 密集模型架构

密集模型是传统的Transformer架构，其中所有参数在每次前向传播中都会被激活。Qwen3的密集模型从0.6B到32B，覆盖了从轻量级到大型的多个规模。

密集模型的主要特点包括：

参数高效利用：所有参数都参与每次计算，充分利用了模型容量。
实现简单：架构相对简单，易于实现和部署。
训练稳定：训练过程相对稳定，超参数调整相对简单。
推理速度快：在相同参数量下，推理速度通常快于MoE模型。

以下是Qwen3密集模型的主要配置参数：

模型	层数	模型维度	注意力头数	参数量
Qwen3-0.6B	24	1024	8	0.6B
Qwen3-1.7B	24	1536	16	1.7B
Qwen3-4B	32	2560	32	4.0B
Qwen3-8B	32	3584	32	8.0B
Qwen3-14B	40	4608	40	14.0B
Qwen3-32B	60	6144	64	32.0B

3.2.2 MoE模型架构

MoE模型是一种稀疏激活的神经网络架构，其核心思想是"专家分工"。在Qwen3的MoE模型中，每个Transformer层的前馈网络被替换为多个"专家"（Expert）网络，但在处理每个输入时，只激活其中的一部分专家。

上图展示了MoE层的基本结构，其中路由器（Router）负责决定每个输入应该由哪些专家处理，然后将多个专家的输出加权合并得到最终结果。

Qwen3的MoE模型具有以下特点：

细粒度专家分割：每个MoE层包含多个专家，每个专家都是一个完整的前馈网络。
Top-k路由：对于每个输入token，路由器选择k个最相关的专家进行处理（Qwen3中k=2）。
无共享专家：与一些包含共享专家的MoE模型不同，Qwen3的每个专家都有独特的参数，增强了专业化程度。
全局批次负载平衡：为了确保各个专家的工作负载均衡，Qwen3引入了全局批次负载平衡损失。

# MoE层的详细实现
class MoELayer(nn.Module):def __init__(self, input_dim, hidden_dim, output_dim, num_experts=10, top_k=2):"""初始化MoE层参数:input_dim: 输入维度hidden_dim: 隐藏层维度output_dim: 输出维度num_experts: 专家数量top_k: 每次激活的专家数量"""super().__init__()self.input_dim = input_dimself.hidden_dim = hidden_dimself.output_dim = output_dimself.num_experts = num_expertsself.top_k = top_k# 专家路由器self.router = nn.Linear(input_dim, num_experts, bias=False)# 创建专家网络self.experts = nn.ModuleList([SwiGLU(input_dim, hidden_dim, output_dim) for _ in range(num_experts)])def forward(self, x):"""前向传播参数:x: 输入张量 [batch_size, seq_len, input_dim]返回:输出张量 [batch_size, seq_len, output_dim]负载平衡损失"""batch_size, seq_len, _ = x.shape# 计算路由分数router_logits = self.router(x)  # [batch_size, seq_len, num_experts]# 添加噪声以打破平局if self.training:router_logits += torch.randn_like(router_logits) * 1e-2# 选择top-k专家router_probs = F.softmax(router_logits, dim=-1)top_k_probs, top_k_indices = torch.topk(router_probs, self.top_k, dim=-1)# 归一化概率top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)# 准备输出output = torch.zeros_like(x)# 计算专家使用计数expert_counts = torch.zeros(self.num_experts, device=x.device)# 对每个token应用选定的专家for i in range(batch_size):for j in range(seq_len):for k in range(self.top_k):expert_idx = top_k_indices[i, j, k].item()prob = top_k_probs[i, j, k].item()# 更新专家使用计数expert_counts[expert_idx] += 1# 应用专家expert_output = self.experts[expert_idx](x[i:i+1, j:j+1, :])output[i, j] += prob * expert_output.squeeze(0).squeeze(0)# 计算负载平衡损失# 理想情况下，每个专家应该处理相同数量的tokentotal_tokens = batch_size * seq_len * self.top_kideal_count = total_tokens / self.num_experts# 计算专家使用率expert_usage = expert_counts / total_tokens# 计算负载平衡损失load_balancing_loss = torch.sum((expert_usage - (1.0 / self.num_experts))**2)return output, load_balancing_loss

Qwen3的MoE模型配置如下：

模型	层数	模型维度	注意力头数	专家数量	激活专家数	总参数量	激活参数量
Qwen3-30B-A3B	60	8192	128	10	1	30B	3B
Qwen3-72B-A7B	80	8192	128	10	1	72B	7B
Qwen3-235B-A22B	120	8192	128	10	2	235B	22B

MoE架构的主要优势在于：

参数效率：虽然总参数量很大，但每次推理只激活一小部分，大大提高了参数效率。
性能提升：在相同激活参数量下，MoE模型通常比密集模型性能更好。
专业化能力：不同专家可以专注于不同类型的输入，提高模型处理多样化任务的能力。

例如，Qwen3-235B-A22B虽然总参数量达到235B，但每次推理只激活22B参数，这使得它能够在保持高性能的同时，控制计算资源需求。

3.3 分词器实现与多语言支持

分词器（Tokenizer）是大型语言模型的重要组成部分，它负责将原始文本转换为模型可以处理的token序列。Qwen3的分词器设计考虑了多语言支持和处理效率，下面我们将详细介绍其实现。

3.3.1 字节级字节对编码（Byte-level BPE）

Qwen3采用了字节级字节对编码（Byte-level Byte-Pair Encoding, BBPE）作为分词方法。BBPE的核心思想是将文本视为字节序列，然后应用BPE算法学习常见的字节组合，形成词汇表。

BBPE的主要优势包括：

通用性：能够处理任何语言和字符，包括表情符号、特殊符号等。
无未知词：由于基于字节，任何输入都可以被编码，不会出现未知词（UNK）。
效率：编码和解码过程高效，适合实时应用。

# BBPE分词器的简化实现
class ByteLevelBPETokenizer:def __init__(self, vocab_file, merges_file):"""初始化BBPE分词器参数:vocab_file: 词汇表文件路径merges_file: 合并规则文件路径"""# 加载词汇表with open(vocab_file, 'r', encoding='utf-8') as f:self.vocab = json.load(f)# 创建词汇表的反向映射self.ids_to_tokens = {v: k for k, v in self.vocab.items()}# 加载合并规则self.merges = {}with open(merges_file, 'r', encoding='utf-8') as f:for i, line in enumerate(f):if i == 0:  # 跳过版本信息continuepair = line.strip().split()if len(pair) == 2:self.merges[tuple(pair)] = idef encode(self, text):"""将文本编码为token ID参数:text: 输入文本返回:token ID列表"""# 将文本转换为字节bytes_str = text.encode('utf-8')# 初始化为单个字节的tokentokens = [bytes([b]) for b in bytes_str]# 应用合并规则while len(tokens) > 1:pairs = self._get_pairs(tokens)if not pairs:break# 找到优先级最高的合并对best_pair = min(pairs, key=lambda pair: self.merges.get(pair, float('inf')))if best_pair not in self.merges:break# 执行合并tokens = self._merge(tokens, best_pair)# 将token转换为IDids = [self.vocab.get(token, self.vocab['<unk>']) for token in tokens]return idsdef decode(self, ids):"""将token ID解码为文本参数:ids: token ID列表返回:解码后的文本"""# 将ID转换为tokentokens = [self.ids_to_tokens.get(id, '<unk>') for id in ids]# 将token连接并解码为文本bytes_str = b''.join([t.encode('latin1') for t in tokens])text = bytes_str.decode('utf-8', errors='replace')return textdef _get_pairs(self, tokens):"""获取相邻token对"""pairs = set()for i in range(len(tokens) - 1):pairs.add((tokens[i], tokens[i + 1]))return pairsdef _merge(self, tokens, pair):"""合并指定的token对"""new_tokens = []i = 0while i < len(tokens):if i < len(tokens) - 1 and (tokens[i], tokens[i + 1]) == pair:new_tokens.append(tokens[i] + tokens[i + 1])i += 2else:new_tokens.append(tokens[i])i += 1return new_tokens

3.3.2 多语言支持扩展

Qwen3将支持的语言从Qwen2.5的29种扩展到了119种语言和方言，这一扩展主要通过以下方式实现：

扩大词汇表：Qwen3的词汇表包含15.2万个token，比Qwen2.5的15.1万略有增加，但通过更高效的编码方式支持了更多语言。
多语言训练数据：在预训练阶段使用了覆盖119种语言的大规模文本语料，使模型能够学习各种语言的模式和规则。
语言识别与切换：Qwen3能够自动识别输入文本的语言，并在不同语言之间无缝切换，这对于多语言对话和翻译任务尤为重要。

以下是Qwen3支持的部分语言列表：

语言类别	包含语言
欧洲语系	英语、法语、德语、西班牙语、意大利语、葡萄牙语、俄语、荷兰语、瑞典语、波兰语、捷克语、丹麦语、挪威语、芬兰语、希腊语、匈牙利语、罗马尼亚语等
亚洲语系	中文（简体和繁体）、日语、韩语、阿拉伯语、希伯来语、泰语、越南语、印尼语、马来语、菲律宾语、乌尔都语、印地语、孟加拉语等
非洲语系	斯瓦希里语、豪萨语、约鲁巴语、伊博语、祖鲁语、索马里语等
其他语系	土耳其语、波斯语、库尔德语等

这种广泛的语言支持使Qwen3能够服务全球用户，促进跨语言交流和知识共享。

3.3.3 特殊token与控制码

Qwen3的分词器包含多种特殊token和控制码，用于标记序列的开始和结束、控制生成行为等：

基本特殊token：
- <bos>: 序列开始标记
- <eos>: 序列结束标记
- <pad>: 填充标记
- <unk>: 未知词标记
控制码：
- 思考模式控制：用于切换思考模式和非思考模式
- 语言控制：用于指定生成特定语言的文本
- 格式控制：用于控制生成文本的格式（如代码、表格等）

这些特殊token和控制码增强了Qwen3的灵活性和可控性，使用户能够更精确地控制模型的行为。

3.4 长上下文处理技术

Qwen3能够处理长达128K tokens的超长上下文，这一能力主要通过以下技术实现：

3.4.1 RoPE基频调整

Qwen3将RoPE的基频从传统的10,000增加到了1,000,000，这一调整使得模型能够更好地区分远距离位置，提高了长序列处理能力。

# RoPE基频调整实现
def get_rope_embeddings(dim, max_seq_len, base=1000000.0):"""生成RoPE位置编码参数:dim: 嵌入维度max_seq_len: 最大序列长度base: 基频（Qwen3使用1000000）返回:余弦和正弦位置编码"""# 计算频率inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))# 生成位置序列t = torch.arange(max_seq_len).float()# 计算外积freqs = torch.outer(t, inv_freq)# 计算余弦和正弦cos = torch.cos(freqs)sin = torch.sin(freqs)return cos, sin

3.4.2 YARN扩展

YARN（Yet Another RoPE extensioN）是一种改进的RoPE扩展方法，它通过调整位置编码的插值方式，使模型能够处理超出训练长度的序列。

# YARN实现
def yarn_rope_embeddings(dim, max_train_len, max_infer_len, base=1000000.0, scale=1.0):"""使用YARN方法生成扩展的RoPE位置编码参数:dim: 嵌入维度max_train_len: 最大训练序列长度max_infer_len: 最大推理序列长度base: 基频scale: 缩放因子返回:扩展的余弦和正弦位置编码"""# 计算原始频率inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))# 应用YARN缩放if max_infer_len > max_train_len:# 计算缩放比例yarn_scale = scale * math.log(max_infer_len / max_train_len) / math.log(2)# 缩放频率inv_freq = inv_freq * (1 - yarn_scale * torch.arange(0, dim, 2).float() / dim)# 生成位置序列t = torch.arange(max_infer_len).float()# 计算外积freqs = torch.outer(t, inv_freq)# 计算余弦和正弦cos = torch.cos(freqs)sin = torch.sin(freqs)return cos, sin

3.4.3 Dual Chunk Attention (DCA)

DCA是一种高效的注意力机制，它将长序列分成多个块，并在块内和块间分别计算注意力，大大提高了处理长序列的效率。

# DCA实现
class DualChunkAttention(nn.Module):def __init__(self, d_model, num_heads, chunk_size=4096):"""初始化DCA参数:d_model: 模型维度num_heads: 注意力头数chunk_size: 块大小"""super().__init__()self.d_model = d_modelself.num_heads = num_headsself.chunk_size = chunk_sizeself.head_dim = d_model // num_headsself.q_proj = nn.Linear(d_model, d_model, bias=False)self.k_proj = nn.Linear(d_model, d_model, bias=False)self.v_proj = nn.Linear(d_model, d_model, bias=False)self.out_proj = nn.Linear(d_model, d_model, bias=False)def forward(self, x, attention_mask=None):"""前向传播参数:x: 输入张量 [batch_size, seq_len, d_model]attention_mask: 注意力掩码返回:输出张量 [batch_size, seq_len, d_model]"""batch_size, seq_len, _ = x.shape# 投影q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)k = self.k_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)v = self.v_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)# 计算块数num_chunks = (seq_len + self.chunk_size - 1) // self.chunk_size# 准备输出output = torch.zeros_like(q)# 块内注意力for i in range(num_chunks):# 计算块的起止位置start_idx = i * self.chunk_sizeend_idx = min(start_idx + self.chunk_size, seq_len)# 提取当前块q_chunk = q[:, :, start_idx:end_idx, :]k_chunk = k[:, :, start_idx:end_idx, :]v_chunk = v[:, :, start_idx:end_idx, :]# 计算块内注意力scores = torch.matmul(q_chunk, k_chunk.transpose(-2, -1)) / math.sqrt(self.head_dim)# 应用掩码（如果有）if attention_mask is not None:mask_chunk = attention_mask[:, start_idx:end_idx, start_idx:end_idx]scores = scores + mask_chunk.unsqueeze(1)# 应用softmaxattn_weights = F.softmax(scores, dim=-1)# 计算输出chunk_output = torch.matmul(attn_weights, v_chunk)# 保存块输出output[:, :, start_idx:end_idx, :] = chunk_output# 块间注意力（简化实现）if num_chunks > 1:# 为每个块创建一个代表性tokenrep_tokens = torch.zeros(batch_size, self.num_heads, num_chunks, self.head_dim, device=x.device)for i in range(num_chunks):start_idx = i * self.chunk_sizeend_idx = min(start_idx + self.chunk_size, seq_len)# 使用平均池化创建代表性tokenrep_tokens[:, :, i, :] = q[:, :, start_idx:end_idx, :].mean(dim=2)# 计算块间注意力rep_scores = torch.matmul(rep_tokens, rep_tokens.transpose(-2, -1)) / math.sqrt(self.head_dim)rep_attn_weights = F.softmax(rep_scores, dim=-1)# 应用块间注意力for i in range(num_chunks):for j in range(num_chunks):if i != j:  # 跳过自身块start_i = i * self.chunk_sizeend_i = min(start_i + self.chunk_size, seq_len)start_j = j * self.chunk_sizeend_j = min(start_j + self.chunk_size, seq_len)# 计算块间注意力权重weight = rep_attn_weights[:, :, i, j].unsqueeze(-1).unsqueeze(-1)# 计算块间注意力cross_attn = torch.matmul(q[:, :, start_i:end_i, :].unsqueeze(3),k[:, :, start_j:end_j, :].transpose(-2, -1).unsqueeze(2)) / math.sqrt(self.head_dim)cross_attn_weights = F.softmax(cross_attn, dim=-1)cross_output = torch.matmul(cross_attn_weights, v[:, :, start_j:end_j, :].unsqueeze(2))# 加权并添加到输出output[:, :, start_i:end_i, :] += weight * cross_output.squeeze(3)# 重塑并投影输出output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)output = self.out_proj(output)return output