pkgs/MagiCompiler/magi_compiler/tokenflow/utils.py

# 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)