import torch
from torch import nn
from torch.nn import attention
import time

# 设置随机种子
torch.manual_seed(42)


MAX_TOKEN_LENGTH = 2048
class MultiQueryAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super(MultiQueryAttention, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        assert (
            self.head_dim * num_heads == self.embed_dim
        ), "embed_dim must be divisible by num_heads"

        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, self.head_dim)
        self.v_proj = nn.Linear(embed_dim, self.head_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        
        self.register_buffer('attention_mask', torch.tril(torch.ones(MAX_TOKEN_LENGTH, MAX_TOKEN_LENGTH)))
        
        self.o_proj = nn.Linear(embed_dim, embed_dim)
        
    def forward(self, x, mask=None):
        batch_size, seq_len , _ = x.size()
        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,  1, self.head_dim)
        k = k.transpose(1, 2)
        v = self.v_proj(x).view(batch_size,   seq_len, 1, self.head_dim)
        v = v.transpose(1, 2)
        
        attention_weight = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        # print('attention_weight shape:', attention_weight)
        
        if mask is not None:
            attention_mask_temp = self.attention_mask[:seq_len, :seq_len]
            attention_weight = attention_weight.masked_fill(attention_mask_temp == 0, float('-inf'))
            
        attention_probs = nn.Softmax(dim=-1)(attention_weight)
        attention_output = torch.matmul(attention_probs, v)
        attention_output = attention_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)
        
        return self.out_proj(attention_output)
    
    
class MultiQueryAttentionEinsum(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super(MultiQueryAttentionEinsum, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        assert (
            self.head_dim * num_heads == self.embed_dim
        ), "embed_dim must be divisible by num_heads"

        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, self.head_dim)
        self.v_proj = nn.Linear(embed_dim, self.head_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        
        self.register_buffer('attention_mask', torch.tril(torch.ones(MAX_TOKEN_LENGTH, MAX_TOKEN_LENGTH)))
        
        self.o_proj = nn.Linear(embed_dim, embed_dim)
    def forward_einsum(self, x, mask=None):
        batch_size, seq_len , _ = x.size()
        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
        
        k = self.k_proj(x)
        # k = k.transpose(1, 2)
        v = self.v_proj(x)
        
        attention_weight = torch.einsum('bihj,bmj->bhim',q,k) / (self.head_dim ** 0.5)
        attention_weight = attention_weight[:, 0, :, :]
        # print('attention_weight shape (einsum):', attention_weight)
        if mask is not None:
            attention_mask_temp = self.attention_mask[:seq_len, :seq_len]
            attention_weight = attention_weight.masked_fill(attention_mask_temp == 0, float('-inf'))
            
        attention_probs = nn.Softmax(dim=-1)(attention_weight)
        # attention_output = torch.matmul(attention_probs, v)
        attention_output = torch.einsum('bkj,bjn->bkn',attention_probs , v)
        
        attention_output = attention_output.view(batch_size, seq_len, 1, self.head_dim)
        
        # Expand to num_heads
        attention_output = attention_output.expand(-1, -1, self.num_heads, -1)
        
        attention_output = attention_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)
        
        
        return self.out_proj(attention_output)
    
    
    # 测试代码
batch_size = 64
seq_len = 512
d_model = 128
num_heads = 8  # 8个Query头，但只有1个KV头

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
model = MultiQueryAttention(d_model, num_heads).to(device)
model_einsum = MultiQueryAttentionEinsum(d_model, num_heads).to(device)
x = torch.randn(batch_size, seq_len, d_model).to(device)

print(f"Input shape: {x.shape}")

# Test standard forward method
print("model forward",model.forward(x))
print("einsum model forward",model_einsum.forward_einsum(x))


# Run each forward 10 times and measure times
n_runs = 10
model_times = []
for i in range(n_runs):
    if device.type == "cuda":
        torch.cuda.synchronize()
    t0 = time.time()
    output = model(x)
    if device.type == "cuda":
        torch.cuda.synchronize()
    model_times.append(time.time() - t0)

einsum_times = []
for i in range(n_runs):
    if device.type == "cuda":
        torch.cuda.synchronize()
    t0 = time.time()
    einsum_output = model_einsum.forward_einsum(x)
    if device.type == "cuda":
        torch.cuda.synchronize()
    einsum_times.append(time.time() - t0)

mean_model = sum(model_times) / n_runs
mean_einsum = sum(einsum_times) / n_runs

print(f"Output shape: {output.shape}")
print(f"Einsum output shape: {einsum_output.shape}")

print(f"\nModel times (s): {model_times}")
print(f"Einsum times (s): {einsum_times}")

print(f"\nModel mean time: {mean_model:.6f} s")
print(f"Einsum mean time: {mean_einsum:.6f} s")
print(f"Speedup (model / einsum): {mean_model/mean_einsum:.2f}x" if mean_einsum > 0 else "N/A")

from networkx.algorithms.centrality import group
from openai.types import batch
from regex import P


MAX_TOKEN_LENGTH = 2048
class GroupQueryAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, group_size):
        super(GroupQueryAttention, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.group_size = group_size
        
        
        self.num_kv_heads = group_size
        # 计算每组有多少个 Query 头 (Group Size)
        self.num_heads_per_group = num_heads // self.num_kv_heads 
        
        assert (
            self.head_dim * num_heads == self.embed_dim
        ), "embed_dim must be divisible by num_heads"

        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, self.head_dim*group_size)
        self.v_proj = nn.Linear(embed_dim, self.head_dim*group_size)
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        
        self.register_buffer('attention_mask', torch.tril(torch.ones(MAX_TOKEN_LENGTH, MAX_TOKEN_LENGTH)))
        
        self.o_proj = nn.Linear(embed_dim, embed_dim)
        
    def forward(self, x, mask=None):
        batch_size, seq_len , _ = x.size()
        num_heads_every_group = self.num_heads // self.group_size
        q = self.q_proj(x).view(batch_size, seq_len, self.group_size, num_heads_every_group, self.head_dim)
        
        k = self.k_proj(x).view(batch_size,  seq_len,  self.group_size, self.head_dim)
        v = self.v_proj(x).view(batch_size,  seq_len, self.group_size, self.head_dim)
        
        attention_weight = torch.einsum('bignd,bjgd->bgnij', q, k) / (self.head_dim ** 0.5)
        # print('attention_weight shape:', attention_weight)
        
        if mask is not None:
            attention_mask_temp = self.attention_mask[:seq_len, :seq_len]
            attention_weight = attention_weight.masked_fill(attention_mask_temp == 0, float('-inf'))
            
        attention_probs = nn.Softmax(dim=-1)(attention_weight)
        attention_output = torch.einsum('bgnij,bjgd->bgnid',attention_probs , v)
        # print("attention_output: ",attention_output)
        attention_output = attention_output.view(batch_size, self.num_heads,  seq_len, self.head_dim).transpose(1,2)
        attention_output = attention_output.contiguous().view(batch_size, seq_len, self.num_heads*self.head_dim)
        
        # attention_output = attention_output.reshape(batch_size, seq_len, self.embed_dim)
        
        return self.out_proj(attention_output)
    
    def forward_gemini(self, x, mask=None):
        batch_size, seq_len, _ = x.size()
        
        # 1. 投影并重塑维度
        # Q: [Batch, Seq, Num_KV_Heads, Group_Size, Head_Dim]
        # g: num_kv_heads (组数)
        # r: num_heads_per_group (每组内的重复次数)
        q = self.q_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.num_heads_per_group, self.head_dim)
        
        # K, V: [Batch, Seq, Num_KV_Heads, Head_Dim]
        # 相比 MQA，这里多了 'g' 维度，但没有 'r' 维度
        k = self.k_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
        v = self.v_proj(x).view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
        
        # 2. 计算 Attention Scores
        # Einsum 逻辑:
        # b: batch
        # l: query seq len
        # s: key seq len
        # g: num_kv_heads (groups) -> Q 和 K 在此维度必须对齐
        # r: num_heads_per_group  -> K 在此维度广播 (K 没有 r 维度)
        # d: head_dim
        # Q(b,l,g,r,d) @ K(b,s,g,d) -> Scores(b,g,r,l,s)
        attention_weight = torch.einsum('blgrd, bsgd -> bgrls', q, k) / (self.head_dim ** 0.5)
        # print('attention_weight shape:', attention_weight)
        
        # 3. Masking
        if mask is not None:
            # Mask 形状通常是 [Seq, Seq]
            attention_mask_temp = self.attention_mask[:seq_len, :seq_len]
            # 自动广播 mask 到 [Batch, Groups, Group_Size, Seq, Seq]
            attention_weight = attention_weight.masked_fill(attention_mask_temp == 0, float('-inf'))
            
        attention_probs = nn.Softmax(dim=-1)(attention_weight)
        
        # 4. 计算加权和
        # Scores(b,g,r,l,s) @ V(b,s,g,d) -> Output(b,l,g,r,d)
        # 同样，V 在 r 维度上被广播
        attention_output = torch.einsum('bgrls, bsgd -> blgrd', attention_probs, v)
        # print("attention_output: ",attention_output)
        # 5. 合并维度
        # [Batch, Seq, Groups, Group_Size, Head_Dim] -> [Batch, Seq, Num_Heads, Head_Dim] -> [Batch, Seq, Embed_Dim]
        attention_output = attention_output.reshape(batch_size, seq_len, self.embed_dim)
        
        return self.out_proj(attention_output)
    
    # 测试代码
batch_size = 64
seq_len = 512
d_model = 128
num_heads = 8  # 8个Query头，但只有1个KV头
group_size = 2

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
model = GroupQueryAttention(d_model, num_heads, group_size).to(device)

model_output = model(x)
print(f"Group Query Attention output shape: {model_output[0,0,0]}")

forward_gemini_output = model.forward_gemini(x)
print(f"Group Query Attention Gemini output shape: {forward_gemini_output[0,0,0]}")

import torch
import torch.nn as nn

MAX_TOKEN_LENGTH = 2048

class GroupedQueryAttentionEinsum():
    def __init__(self, embed_dim, num_heads, num_kv_heads):
        """
        Args:
            embed_dim: 模型嵌入维度
            num_heads: Query 的总头数
            num_kv_heads: Key/Value 的头数 (GQA 分组数)
        """
        super(GroupedQueryAttentionEinsum, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = embed_dim // num_heads
        
        # 计算每组有多少个 Query 头 (Group Size)
        self.num_heads_per_group = num_heads // num_kv_heads

        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
        assert num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads"

        # Q 投影: 保持原样，覆盖所有 Query 头
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        
        # K, V 投影: 维度由 num_kv_heads 决定
        self.k_proj = nn.Linear(embed_dim, self.num_kv_heads * self.head_dim)
        self.v_proj = nn.Linear(embed_dim, self.num_kv_heads * self.head_dim)
        
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        
        self.register_buffer('attention_mask', torch.tril(torch.ones(MAX_TOKEN_LENGTH, MAX_TOKEN_LENGTH)))
        

    
    
# 测试代码
batch_size = 64
seq_len = 512
d_model = 128
num_heads = 8  # 8个Query头，但只有1个KV头
group_size = 2

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
model = GroupedQueryAttentionEinsum(d_model, num_heads, group_size).to(device)

model_output = model(x)
print(f"Group Query Attention output shape: {model_output}")

from IPython.display import HTML, display

html_code = '''
<!DOCTYPE html>
<html>
<head>
    <style>
        .container { max-width: 900px; margin: 0 auto; font-family: sans-serif; padding: 20px; }
        h2 { color: #333; border-bottom: 2px solid #2196f3; padding-bottom: 10px; }
        .box { padding: 10px 15px; margin: 5px; border-radius: 5px; font-family: monospace; font-size: 12px; text-align: center; }
        .input { background: #e3f2fd; border: 2px solid #2196f3; }
        .proj { background: #fff3e0; border: 2px solid #ff9800; }
        .norm { background: #f3e5f5; border: 2px solid #9c27b0; }
        .attn { background: #e8f5e9; border: 2px solid #4caf50; }
        .rope { background: #fce4ec; border: 2px solid #e91e63; }
        .arrow { font-size: 20px; color: #666; }
        .flex { display: flex; align-items: center; flex-wrap: wrap; justify-content: center; }
        .row { flex-direction: row; }
        .code { background: #263238; color: #aed581; padding: 15px; border-radius: 5px; font-family: monospace; margin: 10px 0; }
        .highlight { background: #ffeb3b; padding: 2px 5px; }
        table { border-collapse: collapse; width: 100%; margin: 10px 0; }
        th, td { border: 1px solid #ddd; padding: 8px; text-align: center; }
        th { background: #2196f3; color: white; }
    </style>
</head>
<body>
<div class="container">
    <h2>🔮 Gemma3 Attention 完整流程</h2>
    
    <!-- 完整流程图 -->
    <div class="flex row">
        <div class="box input">Hidden States<br>[B, Seq, Hidden]</div>
        <span class="arrow">→</span>
        <div class="box proj">Q/K/V 投影<br>Linear</div>
        <span class="arrow">→</span>
        <div class="box norm">Q_norm, K_norm<br>(RMSNorm)</div>
        <span class="arrow">→</span>
        <div class="box rope">RoPE<br>旋转位置编码</div>
        <span class="arrow">→</span>
        <div class="box attn">Attention<br>(GQA + Sliding)</div>
        <span class="arrow">→</span>
        <div class="box proj">O 投影<br>Linear</div>
    </div>

    <hr>
    
    <h2>🔑 GQA 核心: repeat_kv</h2>
    <table>
        <tr>
            <th>配置</th>
            <th>num_attention_heads</th>
            <th>num_key_value_heads</th>
            <th>n_rep</th>
        </tr>
        <tr>
            <td>Gemma3-3b</td>
            <td>8</td>
            <td>2</td>
            <td>4</td>
        </tr>
        <tr>
            <td>gemma3-2b</td>
            <td>8</td>
            <td>4</td>
            <td>2</td>
        </tr>
        <tr>
            <td>MQA</td>
            <td>8</td>
            <td>1</td>
            <td>8</td>
        </tr>
    </table>
    
    <div class="code">
# repeat_kv 核心逻辑
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    batch, num_kv_heads, seqlen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    # 复制KV头n_rep次
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_kv_heads, n_rep, seqlen, head_dim)
    return hidden_states.reshape(batch, num_kv_heads * n_rep, seqlen, head_dim)
    </div>

    <hr>
    
    <h2>🪟 Sliding Window Attention</h2>
    <p><strong>原理：</strong>每个位置i只能attend到[i-window+1, i]范围内的token，超出窗口的被mask掉</p>
    
    <table>
        <tr>
            <th>Window=5</th>
            <th>K0</th>
            <th>K1</th>
            <th>K2</th>
            <th>K3</th>
            <th>K4</th>
            <th>K5</th>
            <th>K6</th>
            <th>K7</th>
            <th>K8</th>
        </tr>
        <tr>
            <td><strong>Q0</strong></td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
        </tr>
        <tr>
            <td><strong>Q1</strong></td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
        </tr>
        <tr>
            <td><strong>Q2</strong></td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
        </tr>
        <tr>
            <td><strong>Q3</strong></td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
        </tr>
        <tr>
            <td><strong>Q4</strong></td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#4caf50">✓</td>
            <td style="background:#e0e0e0">×</td>
            <td style="background:#e0e0e0">×</td>
        </tr>
    </table>
    
    <p style="color: #666; font-size: 12px;">
    ✓ = 可以attend | × = 被mask (attn_weights = -inf → softmax后≈0)
    </p>

    <hr>
    
    <h2>📊 Gemma3 整体架构</h2>
    <svg width="600" height="320" viewBox="0 0 600 320">
        <!-- 输入框 -->
        <rect x="50" y="20" width="80" height="25" rx="4" fill="#e3f2fd" stroke="#2196f3" stroke-width="2"/>
        <text x="90" y="37" text-anchor="middle" font-size="11">input_ids</text>
        
        <rect x="140" y="20" width="90" height="25" rx="4" fill="#fff3e0" stroke="#ff9800" stroke-width="2"/>
        <text x="185" y="37" text-anchor="middle" font-size="11">pixel_values</text>
        
        <!-- Vision路径 -->
        <line x1="185" y1="45" x2="185" y2="65" stroke="#666" stroke-width="1.5"/>
        <rect x="140" y="65" width="90" height="25" rx="4" fill="#fff3e0" stroke="#ff9800" stroke-width="2"/>
        <text x="185" y="82" text-anchor="middle" font-size="11">Vision Tower</text>
        
        <line x1="185" y1="90" x2="185" y2="105" stroke="#666" stroke-width="1.5"/>
        <rect x="140" y="105" width="90" height="25" rx="4" fill="#fff3e0" stroke="#ff9800" stroke-width="2"/>
        <text x="185" y="122" text-anchor="middle" font-size="11">Projector</text>
        
        <!-- Text Embed -->
        <line x1="90" y1="45" x2="90" y2="65" stroke="#666" stroke-width="1.5"/>
        <rect x="50" y="65" width="80" height="25" rx="4" fill="#e3f2fd" stroke="#2196f3" stroke-width="2"/>
        <text x="90" y="82" text-anchor="middle" font-size="11">Embed</text>
        
        <!-- 合并 -->
        <line x1="90" y1="130" x2="185" y2="130" stroke="#666" stroke-width="2"/>
        <line x1="185" y1="130" x2="90" y2="130" marker-end="url(#arrow)" stroke="#666"/>
        
        <rect x="40" y="140" width="180" height="35" rx="4" fill="#f5f5f5" stroke="#333" stroke-width="2"/>
        <text x="130" y="162" text-anchor="middle" font-size="13"><strong>inputs_embeds</strong></text>
        
        <!-- Decoder Layers -->
        <line x1="130" y1="175" x2="130" y2="195" stroke="#666" stroke-width="2"/>
        
        <rect x="60" y="195" width="140" height="28" rx="4" fill="#e8f5e9" stroke="#4caf50" stroke-width="2"/>
        <text x="130" y="213" text-anchor="middle" font-size="11">Gemma3DecoderLayer × N</text>
        
        <rect x="60" y="230" width="140" height="28" rx="4" fill="#e8f5e9" stroke="#4caf50" stroke-width="2"/>
        <text x="130" y="248" text-anchor="middle" font-size="11">Gemma3DecoderLayer × N</text>
        
        <line x1="130" y1="258" x2="130" y2="275" stroke="#666" stroke-width="2"/>
        
        <!-- RMSNorm + 输出 -->
        <rect x="85" y="275" width="90" height="22" rx="4" fill="#f3e5f5" stroke="#9c27b0" stroke-width="2"/>
        <text x="130" y="290" text-anchor="middle" font-size="10">RMSNorm</text>
        
        <line x1="130" y1="297" x2="130" y2="315" stroke="#666" stroke-width="2"/>
        <rect x="80" y="305" width="100" height="25" rx="4" fill="#e3f2fd" stroke="#2196f3" stroke-width="2"/>
        <text x="130" y="322" text-anchor="middle" font-size="11">last_hidden_state</text>
        
        <defs>
            <marker id="arrow" markerWidth="8" markerHeight="8" refX="0" refY="3" orient="auto">
                <path d="M0,0 L0,6 L8,3 z" fill="#666"/>
            </marker>
        </defs>
        
        <!-- 图例 -->
        <rect x="350" y="60" width="15" height="15" fill="#e3f2fd" stroke="#2196f3"/>
        <text x="375" y="72" font-size="11">Text Input</text>
        
        <rect x="350" y="85" width="15" height="15" fill="#fff3e0" stroke="#ff9800"/>
        <text x="375" y="97" font-size="11">Vision</text>
        
        <rect x="350" y="110" width="15" height="15" fill="#e8f5e9" stroke="#4caf50"/>
        <text x="375" y="122" font-size="11">Decoder Layer</text>
        
        <text x="350" y="150" font-size="14"><strong>Gemma3Model:</strong></text>
        <text x="350" y="170" font-size="11" fill="#666">1. Vision Tower → feature extraction</text>
        <text x="350" y="190" font-size="11" fill="#666">2. Projector → 映射到text空间</text>
        <text x="350" y="210" font-size="11" fill="#666">3. 合并 token 和 image embeds</text>
        <text x="350" y="230" font-size="11" fill="#666">4. N层Decoder处理</text>
        <text x="350" y="250" font-size="11" fill="#666">5. RMSNorm + 输出</text>
    </svg>

</div>
</body>
</html>
'''

display(HTML(html_code))