# Copyright (c) 2026 SandAI. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from typing import Callable import torch from magi_compiler.api import magi_compile from magi_compiler.utils import magi_logger, nvtx # 先补全依赖的类定义（确保代码可独立运行） class ModelConfig: def __init__( self, hidden_size, num_layers, num_heads_q, num_heads_kv, head_dim, intermediate_size, activation_type, params_dtype=torch.float32, eps=1e-06, ): self.hidden_size = hidden_size self.num_layers = num_layers self.num_heads_q = num_heads_q self.num_heads_kv = num_heads_kv self.head_dim = head_dim self.intermediate_size = intermediate_size self.activation_type = activation_type self.params_dtype = params_dtype self.eps = eps def __repr__(self): return ( f"ModelConfig(hidden_size={self.hidden_size}, num_layers={self.num_layers}, " f"num_heads_q={self.num_heads_q}, num_heads_kv={self.num_heads_kv}, " f"head_dim={self.head_dim}, intermediate_size={self.intermediate_size}, " f"activation_type='{self.activation_type}', params_dtype={self.params_dtype}, eps={self.eps})" ) @magi_compile(dynamic_arg_dims={'x': [0]}) class CompiledTransformerModel(torch.nn.Module): def __init__(self, config: ModelConfig): super().__init__() self.mod = TransformerModel(config) def forward(self, x): return self.mod(x) class TransformerModel(torch.nn.Module): def __init__(self, config: ModelConfig): super().__init__() self.config = config self.layers = torch.nn.ModuleList([TransformerLayer(config) for _ in range(config.num_layers)]) self.final_norm = torch.nn.LayerNorm(config.hidden_size, eps=config.eps, bias=False) def forward(self, x): for layer in self.layers: x = layer(x) x = self.final_norm(x) return x @nvtx.instrument_nvtx class TransformerLayer(torch.nn.Module): def __init__(self, config: ModelConfig): super().__init__() self.attn_norm = torch.nn.LayerNorm(config.hidden_size, eps=config.eps, bias=False) self.attention = GroupedQueryAttention(config) self.mlp_norm = torch.nn.LayerNorm(config.hidden_size, eps=config.eps, bias=False) self.mlp = MLPLayer(config) def forward(self, x): x = x + self.attention(self.attn_norm(x)) x = x + self.mlp(self.mlp_norm(x)) return x class GroupedQueryAttention(torch.nn.Module): def __init__(self, config: ModelConfig): super().__init__() self.n_heads_q = config.num_heads_q # 32 self.n_heads_kv = config.num_heads_kv # 8 self.head_dim = config.head_dim # 128 self.n_rep = self.n_heads_q // self.n_heads_kv # 32//8=4 self.hidden_size = config.hidden_size # 4096 self.q_size = self.n_heads_q * self.head_dim # 32*128=4096 self.kv_size = self.n_heads_kv * self.head_dim # 8*128=1024 self.qkv_proj = torch.nn.Linear(config.hidden_size, self.q_size + 2 * self.kv_size, bias=False) self.o_proj = torch.nn.Linear(self.q_size, config.hidden_size, bias=False) def forward(self, x): qkv = self.qkv_proj(x) q, k, v = torch.split(qkv, [self.q_size, self.kv_size, self.kv_size], dim=-1) q = q.view(1, -1, self.n_heads_q, self.head_dim) k = k.view(1, -1, self.n_heads_kv, self.head_dim) v = v.view(1, -1, self.n_heads_kv, self.head_dim) if self.n_rep > 1: k = k.repeat_interleave(self.n_rep, dim=2) v = v.repeat_interleave(self.n_rep, dim=2) q = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2) out: torch.Tensor = my_attention(q, k, v) # out = q out = out.transpose(1, 2) out = out.squeeze(0) out = out.view(-1, self.q_size) out = self.o_proj(out) return out return x # 临时屏蔽注意力计算，专注测试 MLP 部分的性能 class MLPModel(torch.nn.Module): def __init__(self, config: ModelConfig): super().__init__() self.config = config self.layers = torch.nn.ModuleList([MLPLayer(config) for _ in range(config.num_layers)]) def forward(self, x): for layer in self.layers: x = layer(x) return x @nvtx.instrument_nvtx class MLPLayer(torch.nn.Module): def __init__(self, config: ModelConfig): super().__init__() self.pre_norm = torch.nn.LayerNorm(config.hidden_size, eps=config.eps, bias=False) self.fc1 = torch.nn.Linear(config.hidden_size, config.intermediate_size, bias=False) self.activation = torch.nn.GELU() self.fc2 = torch.nn.Linear(config.intermediate_size, config.hidden_size, bias=False) def forward(self, x): x = self.pre_norm(x) # x = self.fc1(x) x = self.activation(x) # x = self.fc2(x) return x @magi_compile(dynamic_arg_dims={'x': [0]}) class CompiledMiniMLP(torch.nn.Module): def __init__(self, config: ModelConfig): super().__init__() self.mod = MLPModel(config) def forward(self, x): return self.mod(x) def benchmark_func(func: Callable, warmup_steps: int = 10, run_steps: int = 10, desc: str = "测试") -> float: torch.cuda.synchronize() start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) @nvtx.instrument_nvtx def warmup(): for _ in range(warmup_steps): _ = func() torch.cuda.synchronize() # 确保预热的 CUDA 操作全部完成 warmup() total_elapsed_ms = None @nvtx.instrument_nvtx def run(): nonlocal total_elapsed_ms total_elapsed_ms = 0.0 # 总耗时（毫秒） start_event.record() for _ in range(run_steps): func() # 要求func内部所有CUDA操作都已提交并完成！ end_event.record() end_event.synchronize() # 确保结束事件已完成 total_elapsed_ms += start_event.elapsed_time(end_event) run() avg_time = total_elapsed_ms / run_steps / 1000.0 total_time = total_elapsed_ms / 1000.0 magi_logger.info("[%s] 完成！平均耗时: %.6f 秒/次 | 总耗时: %.6f 秒 (CUDA Event 精准计时)", desc, avg_time, total_time, rank=0) torch.cuda.synchronize() return avg_time @torch.library.custom_op("athena::my_attention", mutates_args=()) def my_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: return torch.nn.functional.scaled_dot_product_attention(q, k, v) @my_attention.register_fake def _(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: return torch.empty_like(q)