深入解讀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