Op 类代码结构与功能实现分析

op.h 和 op.cpp 文件定义了 InferLLM 框架中的各种算子（Operator）类，这些算子是构建神经网络计算图的基本单元。下面对这些文件的代码结构和功能实现进行详细分析。

1. 基础定义

using OpIOs = std::vector<std::shared_ptr<Tensor>>;
constexpr static size_t PACK_SIZE = 8;

• OpIOs 是一个类型别名，表示算子的输入、输出或权重，它是一个 Tensor 指针的向量。
• PACK_SIZE 是一个常量，用于 Int4 量化权重的打包大小，设置为 8。

2. OpBase 类

OpBase 是所有算子的基类，定义了算子的基本接口和通用功能：

class OpBase {
public:
    OpBase(Device* device, const std::string& name, OpIOs inputs)
            : m_device(device), m_inputs(inputs), m_name(name) {
        for (auto input : m_inputs) {
            input->add_user();
        }
    }

    virtual void pre_execute();
    virtual void execute(WorkSpace* workspace, uint32_t nr_past) {}
    virtual void end_execute();
    virtual void deduce_output_shape();
    virtual size_t get_workspace_in_byte() { return 0; }
    virtual void load_weights(std::ifstream&){};
    virtual uint32_t nr_weights() { return 1; };
    virtual void init(OpIOs, OpIOs, WorkSpace*){};
    
    // 其他方法...
};

2.1 OpBase 的主要方法

• 构造函数：初始化设备、输入和名称，并增加输入张量的用户计数。
• pre_execute：执行前准备，包括准备权重和输出张量的数据。
• execute：执行算子的计算，这是一个虚函数，由子类实现。
• end_execute：执行后清理，减少输入张量的当前用户计数。
• deduce_output_shape：推导输出张量的形状，默认与第一个输入相同。
• get_workspace_in_byte：获取执行时需要的工作空间大小。
• load_weights：从文件加载权重。
• init：初始化算子，设置输入、输出和工作空间。

2.2 OpBase 的辅助方法

• device：获取设备指针。
• get_kernel：获取内核指针。
• set_weights：设置权重张量。
• add_outputs：添加输出张量。
• set_name：设置算子名称。
• weights/inputs/outputs/name：获取权重、输入、输出和名称。
• need_preprocess_weight：判断是否需要预处理权重。
• preprocess_weight：预处理权重，使其更适合计算内核。

3. 具体算子类

3.1 LayerNorm 类

class LayerNorm : public OpBase {
public:
    LayerNorm(
            Device* device, const std::string& name, OpIOs inputs, size_t embd,
            bool mul = true, bool bias = false, bool rms = true, float eps = 1e-5);
    void execute(WorkSpace* workspace, uint32_t nr_past) override;

private:
    bool m_mul;
    bool m_bias;
    bool m_rms;
    float m_norm_eps = 1e-5;
};

LayerNorm 实现了层归一化操作，支持 RMS 归一化和普通归一化，可选是否使用乘法和偏置。

3.1.1 LayerNorm::execute 实现

void LayerNorm::execute(WorkSpace* workspace, uint32_t nr_past) {
    std::shared_ptr<Tensor> weight = nullptr, bias = nullptr;
    int weight_idx = 0;
    if (m_mul) {
        weight = weights()[weight_idx++];
        DType weight_type = weight->dtype();
        INFER_ASSERT(
                weight_type == DType::Float32, "layer norm weights must be float32.");
    }
    if (m_bias) {
        bias = weights()[weight_idx++];
    }
    auto input = inputs()[0];
    auto output = outputs()[0];
    uint32_t seq_len = input->shape()[0];
    uint32_t embd = input->shape()[1];
    auto kernel = get_kernel();
    if (input->dtype() == DType::Float32) {
        const float* src = input->ptr<float>();
        float* dst = output->ptr<float>();
        float *weight_ptr = nullptr, *bias_ptr = nullptr;

        if (m_mul) {
            weight_ptr = weight->ptr<float>();
        }
        if (m_bias) {
            bias_ptr = bias->ptr<float>();
        }
        if (m_rms) {
            kernel->operator()<KernelID::RmsNormFloat>(
                    src, dst, seq_len, embd, m_norm_eps);
        } else {
            kernel->operator()<KernelID::NormFloat>(
                    src, dst, seq_len, embd, m_norm_eps);
        }
        if (weight_ptr) {
            kernel->operator()<KernelID::ElemwiseBroadcastDim0Src1Float>(
                    dst, weight_ptr, dst, seq_len, embd, ElemMode::Mul);
        }
        if (bias_ptr) {
            kernel->operator()<KernelID::ElemwiseBroadcastDim0Src1Float>(
                    dst, bias_ptr, dst, seq_len, embd, ElemMode::Add);
        }
    }
}

这个方法执行层归一化操作：

1. 获取权重和偏置（如果有）
2. 获取输入和输出张量
3. 根据 m_rms 标志选择 RMS 归一化或普通归一化
4. 如果有权重，执行元素级乘法
5. 如果有偏置，执行元素级加法

3.2 MatMul 类

class MatMul : public OpBase {
public:
    MatMul(Device* device, const std::string& name, OpIOs inputs,
           std::vector<size_t> shape, bool bias = false);
    virtual void deduce_output_shape() override;
    virtual bool need_preprocess_weight(Tensor* weight) override;
    virtual std::vector<size_t> preprocess_weight(
            Tensor* tensor, void* src, void* dst) override;
    virtual void execute(WorkSpace* workspace, uint32_t nr_past) override;
    size_t get_workspace_in_byte() override;

    bool m_bias = false;
    bool m_weight_packed = false;
};

MatMul 实现了矩阵乘法操作，支持不同数据类型（Float32、Int8、Int4）的权重，以及可选的偏置。

3.2.1 MatMul::execute 实现

void MatMul::execute(WorkSpace* workspace, uint32_t) {
    auto N = weights()[0]->shape()[0];
    auto K = inputs()[0]->shape()[1];
    auto M = inputs()[0]->shape()[0];
    auto src_dtype = inputs()[0]->dtype();
    auto weight_dtype = weights()[0]->dtype();
    void* p_workspace = workspace->ptr();
    uint32_t p_workspace_size = workspace->length();
    auto kernel = get_kernel();
    if (src_dtype == DType::Float32) {
        float* dst = outputs()[0]->ptr<float>();
        const float* bias = nullptr;
        if (m_bias) {
            bias = weights()[1]->ptr<float>();
        }
        const float* src = inputs()[0]->ptr<float>();
        switch (weight_dtype) {
            case DType::Int4:
                if (!m_weight_packed) {
                    kernel->operator()<KernelID::MatmulInt4Float>(
                            dst, weights()[0]->ptr(), bias, src, M, N, K, p_workspace,
                            p_workspace_size);
                } else {
                    kernel->operator()<KernelID::MatmulInt4FloatPacked>(
                            dst, weights()[0]->ptr(), bias, src, M, N * PACK_SIZE, K,
                            p_workspace, p_workspace_size);
                }
                break;
            case DType::Int8:
                kernel->operator()<KernelID::MatmulInt8Float>(
                        dst, weights()[0]->ptr(), bias, src, M, N, K, p_workspace,
                        p_workspace_size);
                break;
            case DType::Float32:
                kernel->operator()<KernelID::MatmulFloatFloat>(
                        dst, weights()[0]->ptr<float>(), bias, src, M, N, K,
                        p_workspace, p_workspace_size);
                break;
            default:
                INFER_ASSERT(0, "not support");
        }
    }
}

这个方法执行矩阵乘法操作：

1. 获取矩阵的维度（M、N、K）
2. 获取输入和权重的数据类型
3. 根据权重的数据类型选择不同的矩阵乘法内核
4. 如果权重是 Int4 类型且已打包，使用优化的打包版本

3.3 MatMulLast 类

class MatMulLast : public MatMul {
public:
    using MatMul::MatMul;
    void deduce_output_shape() override;
    void execute(WorkSpace* workspace, uint32_t nr_past) override;
    virtual bool need_preprocess_weight(Tensor*) override { return false; }
    size_t get_workspace_in_byte() override;
};

MatMulLast 是 MatMul 的子类，专门用于计算序列中最后一个 token 的矩阵乘法，这在生成阶段很常见。

3.3.1 MatMulLast::execute 实现

void MatMulLast::execute(WorkSpace* workspace, uint32_t) {
    auto N = weights()[0]->shape()[0];
    auto K = weights()[0]->shape()[1];
    //! only compute the last token
    auto M = 1;
    auto row = inputs()[0]->shape()[0];
    auto src_dtype = inputs()[0]->dtype();
    auto weight_dtype = weights()[0]->dtype();
    void* p_workspace = workspace->ptr();
    uint32_t p_workspace_size = workspace->length();
    auto kernel = get_kernel();
    if (src_dtype == DType::Float32) {
        float* dst = outputs()[0]->ptr<float>();
        const float* bias = nullptr;
        if (m_bias) {
            bias = weights()[1]->ptr<float>();
        }
        const float* src = inputs()[0]->ptr<float>() + (row - 1) * K;
        // 根据权重类型选择不同的矩阵乘法内核
        // ...
    }
}

与 MatMul::execute 类似，但只计算最后一个 token（src 指向最后一行）。

3.4 SoftMax 类

class SoftMax : public OpBase {
public:
    SoftMax(Device* device, const std::string& name, OpIOs inputs);
    void execute(WorkSpace* workspace, uint32_t nr_past) override;
};

SoftMax 实现了 softmax 操作，通常用于注意力机制中。

3.4.1 SoftMax::execute 实现

void SoftMax::execute(WorkSpace*, uint32_t) {
    auto input = inputs()[0];
    auto output = outputs()[0];
    uint32_t seq_len = input->shape()[0];
    uint32_t embd = input->shape()[1];
    auto kernel = get_kernel();
    if (output->dtype() == DType::Float32) {
        float* src = input->ptr<float>();
        float* dst = output->ptr<float>();
        kernel->operator()<KernelID::SoftmaxFloat>(src, dst, seq_len, embd);
    } else {
        //! fp16
    }
}

这个方法执行 softmax 操作，对每一行（通常是注意力分数）应用 softmax 函数。

3.5 Reshape 类

class Reshape : public OpBase {
public:
    Reshape(Device* device, const std::string& name, OpIOs inputs,
            std::vector<int> shape);
    void deduce_output_shape() override;

private:
    std::vector<int> m_target_shape;
};

Reshape 改变张量的形状，保持元素总数不变。

3.5.1 Reshape::deduce_output_shape 实现

void Reshape::deduce_output_shape() override {
    size_t len = inputs()[0]->length();
    std::vector<size_t> out_shape;
    out_shape.resize(m_target_shape.size());
    int count = 0;
    for (size_t i = 0; i < m_target_shape.size(); i++) {
        if (m_target_shape[i] != -1) {
            out_shape[i] = m_target_shape[i];
            INFER_ASSERT(len % m_target_shape[i] == 0, "Reshape error.\n");
            len = len / m_target_shape[i];
        } else {
            count++;
        }
    }
    INFER_ASSERT(count == 1, "multi -1 in Reshape param.\n");
    for (size_t i = 0; i < m_target_shape.size(); i++) {
        if (m_target_shape[i] == -1) {
            out_shape[i] = len;
        }
    }
    outputs()[0]->set_shape(out_shape, inputs()[0]->dtype());
}

这个方法推导输出张量的形状：

1. 计算输入张量的元素总数
2. 根据目标形状计算每个维度的大小
3. 如果某个维度是 -1，则自动计算该维度的大小

3.6 Elemwise 类

class Elemwise : public OpBase {
public:
    Elemwise(
            Device* device, const std::string& name, OpIOs inputs, ElemMode mode,
            float scale = -INFINITY);
    void execute(WorkSpace* workspace, uint32_t nr_past) override;

private:
    float m_scale;
    ElemMode m_mode;
};

Elemwise 实现了元素级操作，如加法、乘法、GELU 激活等。

void Elemwise::execute(WorkSpace*, uint32_t) {
    auto input = inputs()[0];
    auto output = outputs()[0];
    uint32_t seq_len = input->shape()[0];
    uint32_t embd = input->shape()[1];
    auto kernel = get_kernel();
    if (input->dtype() == DType::Float32) {
        float* src = input->ptr<float>();
        float* dst = output->ptr<float>();
        if (inputs().size() == 1) {
            if (m_mode == ElemMode::Gelu) {
                kernel->operator()<KernelID::GeluFloat>(src, dst, seq_len * embd);
            } else if (m_mode == ElemMode::Silu) {
                kernel->operator()<KernelID::SiluFloat>(src, dst, seq_len * embd);
            } else if (m_mode == ElemMode::Scale) {
                kernel->operator()<KernelID::ScaleFloat>(
                        src, dst, seq_len * embd, m_scale);
            }
        } else if (inputs().size() == 2) {
            float* src1 = inputs()[1]->ptr<float>();
            if (m_mode == ElemMode::Add) {
                kernel->operator()<KernelID::ElemwiseFloat>(
                        src, src1, dst, seq_len * embd, ElemMode::Add);
            } else if (m_mode == ElemMode::Mul) {
                kernel->operator()<KernelID::ElemwiseFloat>(
                        src, src1, dst, seq_len * embd, ElemMode::Mul);
            }
        }
    }
}

这个方法执行元素级操作：

1. 获取输入和输出张量
2. 根据操作模式（m_mode）选择不同的元素级操作内核
3. 如果只有一个输入，执行单输入操作（如 GELU、Silu、Scale）
4. 如果有两个输入，执行双输入操作（如 Add、Mul）

3.7 SpliteHalfActiveMul 类

SpliteHalfActiveMul 实现了将输入张量分成两半，对一半应用激活函数，然后与另一半相乘的操作。这在一些模型架构中很常见，如 SwiGLU。

void SpliteHalfActiveMul::execute(WorkSpace*, uint32_t) {
    auto input = inputs()[0];
    auto output = outputs()[0];
    uint32_t seq_len = input->shape()[0];
    uint32_t embd = input->shape()[1];
    uint32_t half_embd = embd / 2;
    auto kernel = get_kernel();
    if (input->dtype() == DType::Float32) {
        float* src = input->ptr<float>();
        float* dst = output->ptr<float>();
        if (m_mode == ElemMode::Silu) {
            kernel->operator()<KernelID::SplitSiluMulFloat>(
                    src, dst, seq_len, half_embd);
        } else if (m_mode == ElemMode::Gelu) {
            kernel->operator()<KernelID::SplitGeluMulFloat>(
                    src, dst, seq_len, half_embd);
        }
    }
}

这个方法执行分半激活乘法操作：

1. 获取输入和输出张量
2. 计算每半部分的嵌入维度
3. 根据激活模式选择不同的内核（SiluMul 或 GeluMul）

3.8 DiagMask 类

DiagMask 实现了对角线掩码操作，通常用于自回归模型中的注意力机制，防止模型看到未来的 token。

void DiagMask::execute(WorkSpace*, uint32_t nr_past) {
    auto input = inputs()[0];
    auto output = outputs()[0];
    uint32_t seq_len = input->shape()[0];
    uint32_t embd = input->shape()[1];
    auto kernel = get_kernel();
    if (input->dtype() == DType::Float32) {
        float* src = input->ptr<float>();
        float* dst = output->ptr<float>();
        kernel->operator()<KernelID::DiagMaskFloat>(
                src, dst, seq_len, embd, nr_past);
    }
}

这个方法执行对角线掩码操作：

1. 获取输入和输出张量
2. 调用对角线掩码内核，将未来 token 的注意力分数设为负无穷大

3.9 AttentionBase 类

AttentionBase 是注意力机制的基类，提供了注意力机制的通用功能，如 KV 缓存管理。

size_t AttentionBase::get_workspace_in_byte() {
    auto input = inputs()[0];
    uint32_t seq_len = input->shape()[0];
    uint32_t embd = input->shape()[1];
    uint32_t head_dim = embd / m_head;
    size_t workspace_size = 0;
    // Q, K, V 矩阵的大小
    workspace_size += seq_len * embd * sizeof(float);  // Q
    workspace_size += seq_len * embd * sizeof(float);  // K
    workspace_size += seq_len * embd * sizeof(float);  // V
    // QK 矩阵的大小
    workspace_size += seq_len * seq_len * m_head * sizeof(float);  // QK
    return workspace_size;
}

std::vector<size_t> AttentionBase::preprocess_weight(
        Tensor* tensor, void* src, void* dst) {
    auto kernel = get_kernel();
    size_t M = tensor->shape()[0];
    size_t K = tensor->shape()[1];
    kernel->operator()<KernelID::MatmulInt4Reorder>(src, dst, M, K);
    m_packed_weight = true;
    return std::vector<size_t>{M / PACK_SIZE, K};
}

AttentionBase 类提供了：

1. 工作空间大小计算，用于分配足够的内存进行注意力计算
2. 权重预处理，优化 Int4 量化权重的内存布局

3.10 LlamaAttention 类

LlamaAttention 实现了 Llama 模型的注意力机制，包括旋转位置编码（RoPE）。

void LlamaAttention::execute(WorkSpace* workspace, uint32_t nr_past) {
    auto input = inputs()[0];
    auto output = outputs()[0];
    uint32_t seq_len = input->shape()[0];
    uint32_t embd = input->shape()[1];
    uint32_t head_dim = embd / m_head;
    auto kernel = get_kernel();
    
    // 获取工作空间指针
    float* p_q = (float*)workspace->ptr();
    float* p_k = p_q + seq_len * embd;
    float* p_v = p_k + seq_len * embd;
    float* p_qk = p_v + seq_len * embd;
    
    // 计算 Q, K, V
    if (input->dtype() == DType::Float32) {
        float* src = input->ptr<float>();
        float* dst = output->ptr<float>();
        
        // 根据是否使用融合权重选择不同的计算路径
        if (m_fused_weights) {
            // 使用融合权重计算 Q, K, V
            // ...
        } else {
            // 使用分离权重计算 Q, K, V
            // ...
        }
        
        // 应用旋转位置编码
        kernel->operator()<KernelID::ApplyRotaryPosEmb>(
                p_q, p_k, seq_len, embd, m_head, m_rot, nr_past, m_rotary_mode);
        
        // 将 K, V 存储到缓存中
        kernel->operator()<KernelID::CopyKVCache>(
                p_k, p_v, m_kstorage->ptr<float>(), m_vstorage->ptr<float>(), seq_len,
                embd, m_kstorage->get_id());
        
        // 计算注意力分数 QK
        kernel->operator()<KernelID::BatchMatMul>(
                p_q, m_kstorage->ptr<float>(), p_qk, seq_len,
                m_kstorage->get_id() + seq_len, embd, m_head);
        
        // 应用掩码和 softmax
        kernel->operator()<KernelID::DiagMaskSoftmaxKernel>(
                p_qk, p_qk, seq_len, m_kstorage->get_id() + seq_len, m_head, nr_past);
        
        // 计算注意力输出
        kernel->operator()<KernelID::BatchMatMulTransB>(
                p_qk, m_vstorage->ptr<float>(), dst, seq_len,
                m_vstorage->get_id() + seq_len, head_dim, m_head);
    }
}

这个方法执行 Llama 注意力机制：

1. 获取工作空间指针，为 Q, K, V, QK 分配内存
2. 计算 Q, K, V 矩阵
3. 应用旋转位置编码（RoPE）
4. 将 K, V 存储到缓存中
5. 计算注意力分数 QK
6. 应用掩码和 softmax
7. 计算注意力输出

3.11 GlmAttention 类

GlmAttention 实现了 GLM 模型的注意力机制，与 Llama 注意力机制类似，但有一些特定于 GLM 的处理。

3.12 Glm2MultiQueryAttention 类

Glm2MultiQueryAttention 实现了 GLM-2 模型的多查询注意力机制，这是一种优化的注意力变体，减少了 KV 缓存的大小。

3.13 Embedding 类

Embedding 实现了嵌入层，将 token ID 转换为嵌入向量。

void Embedding::execute(WorkSpace*, uint32_t) {
    auto input = inputs()[0];
    auto output = outputs()[0];
    auto weight = weights()[0];
    uint32_t seq_len = input->shape()[0];
    auto kernel = get_kernel();
    if (output->dtype() == DType::Float32) {
        int32_t* ids = input->ptr<int32_t>();
        float* dst = output->ptr<float>();
        if (weight->dtype() == DType::Float32) {
            float* weight_ptr = weight->ptr<float>();
            kernel->operator()<KernelID::EmbeddingFloat>(
                    ids, weight_ptr, dst, seq_len, m_embd);
        } else if (weight->dtype() == DType::Float16) {
            half* weight_ptr = weight->ptr<half>();
            kernel->operator()<KernelID::EmbeddingFp16>(
                    ids, weight_ptr, dst, seq_len, m_embd);
        }
    }
}

这个方法执行嵌入操作：

1. 获取输入 token ID、输出嵌入向量和嵌入权重
2. 根据权重的数据类型选择不同的嵌入内核
3. 将 token ID 转换为嵌入向量

4. 算子执行流程

算子的典型执行流程如下：

1. 创建算子：

   auto layernorm = std::make_shared<LayerNorm>(device, "layernorm", inputs, embd);
   

2. 推导输出形状：

   layernorm->deduce_output_shape();
   

3. 执行前准备：

   layernorm->pre_execute();
   

4. 执行计算：

   layernorm->execute(workspace, nr_past);
   

5. 执行后清理：

   layernorm->end_execute();
   

5. 内存管理

算子的内存管理主要通过以下机制实现：

1. 引用计数：

   • add_user()：增加张量的用户计数
   • decrease_curr_user_count()：减少张量的当前用户计数
   • resume_user_count()：恢复张量的用户计数

2. 数据准备和回收：

   • prepare_data()：准备张量的数据，可能涉及内存分配或从文件加载
   • recall_data()：回收张量的数据，可能涉及内存释放

3. 工作空间：

   • get_workspace_in_byte()：获取执行时需要的工作空间大小
   • workspace->ptr()：获取工作空间的指针

6. 优化策略

算子实现了多种优化策略：

1. 权重预处理：

   • need_preprocess_weight()：判断是否需要预处理权重
   • preprocess_weight()：预处理权重，优化内存布局

2. KV 缓存：

   • m_kstorage 和 m_vstorage：存储注意力机制的 K 和 V 矩阵
   • prepare_data_with_length()：准备指定长度的缓存数据
   • add_id()：更新缓存的 ID

3. 数据类型优化：

   • 支持多种数据类型（Float32、Float16、Int8、Int4）
   • 为不同数据类型提供专门的计算内核

4. 内存复用：

   • 通过工作空间复用内存
   • 通过张量共享减少内存分配

总结

op.h 和 op.cpp 文件定义了 InferLLM 框架中的各种算子类，这些算子是构建神经网络计算图的基本单元。它们提供了丰富的功能，包括层归一化、矩阵乘法、softmax、元素级操作和注意力机制等。这些算子的实现考虑了性能优化，如内存复用、数据类型优化和权重预处理等。

7. 算子的设计模式

InferLLM 框架中的算子设计采用了以下设计模式：

1. 工厂模式：
   通过 Graph::make_graph 创建计算图，然后在计算图中创建各种算子。

2. 策略模式：
   通过虚函数和多态实现不同算子的不同行为，如 execute、deduce_output_shape 等。

3. 组合模式：
   算子可以组合成更复杂的计算图，形成层次结构。

4. 观察者模式：
   通过引用计数机制，张量可以观察到它的用户（算子）何时不再需要它，从而释放资源。

8. 算子的扩展性

InferLLM 框架的算子设计具有良好的扩展性：

1. 新算子添加：
   只需继承 OpBase 类并实现相应的虚函数，即可添加新的算子。

2. 新模型支持：
   通过组合现有算子或添加新算子，可以支持新的模型架构，如 Llama、GLM 等。

3. 新硬件支持：
   通过在 Kernel 类中添加新的硬件后端实现，可以支持新的硬件平台。

9. 算子的性能优化

InferLLM 框架的算子实现了多种性能优化技术：

1. 内存优化：

   • 工作空间复用
   • 张量共享
   • KV 缓存

2. 计算优化：

   • 权重预处理
   • 融合算子（如 SpliteHalfActiveMul）
   • 特化实现（如 MatMulLast）

3. 数据类型优化：

   • 支持多种数据类型（Float32、Float16、Int8、Int4）
   • 为不同数据类型提供专门的计算内核

10. 算子的调试支持

InferLLM 框架的算子提供了一些调试支持：

1. 命名机制：
   每个算子和张量都有一个名称，便于调试和跟踪。

2. 断言检查：
   使用 INFER_ASSERT 宏进行断言检查，确保算子的正确性。

3. 形状推导：
   通过 deduce_output_shape 方法推导输出张量的形状，便于检查形状是否正确。

11. 与其他组件的交互

算子与框架的其他组件有密切的交互：

1. 与 Device 的交互：
   算子通过 m_device 获取设备信息和内核实现。

2. 与 Kernel 的交互：
   算子通过 get_kernel() 获取内核实现，调用相应的计算函数。

3. 与 Tensor 的交互：
   算子管理输入、输出和权重张量，控制它们的生命周期。

4. 与 KvStorage 的交互：
   注意力算子使用 KvStorage 存储 K 和 V 矩阵，实现缓存机制。

12. 算子的使用示例

以下是一个使用算子构建简单计算图的示例：

// 创建设备和计算图
auto device = std::make_unique<CPUDevice>(KernelType::X86, 4);
auto graph = Graph::make_graph(user_config, device.get(), "llama");

// 创建输入张量
auto input = std::make_shared<Tensor>(device.get(), "input");
input->set_shape({1, 4096}, DType::Float32);

// 创建算子
auto layernorm = std::make_shared<LayerNorm>(device.get(), "layernorm", {input}, 4096);
auto matmul = std::make_shared<MatMul>(device.get(), "matmul", {layernorm->outputs()[0]}, {4096, 4096});
auto softmax = std::make_shared<SoftMax>(device.get(), "softmax", {matmul->outputs()[0]});

// 推导输出形状
layernorm->deduce_output_shape();
matmul->deduce_output_shape();
softmax->deduce_output_shape();

// 创建工作空间
size_t workspace_size = std::max({layernorm->get_workspace_in_byte(), matmul->get_workspace_in_byte(), softmax->get_workspace_in_byte()});
auto workspace = std::make_unique<WorkSpace>(device.get(), workspace_size);

// 执行计算
layernorm->pre_execute();
layernorm->execute(workspace.get(), 0);
layernorm->end_execute();

matmul->pre_execute();
matmul->execute(workspace.get(), 0);
matmul->end_execute();

softmax->pre_execute();
softmax->execute(workspace.get(), 0);
softmax->end_execute();

通过这种方式，可以构建复杂的计算图，实现大型语言模型的推理。

总的来说，op.h 和 op.cpp 文件定义了 InferLLM 框架中的算子系统，它是框架的核心组件之一，负责实现各种神经网络操作。它的设计考虑了性能、扩展性和易用性，为大型语言模型的推理提供了高效的计算支持。