Skip to content

Refactor: Merge 2 MHA paths #192

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Chamberlain0w0 wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Refactor: Merge 2 MHA paths #192

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
edf3216
refactor: merge 2 MHA paths, rename attention_type to position_embedd…
Chamberlain0w0 Jul 22, 2026
ccdf048
fix: create mask buffer only for GPT-2 case
Chamberlain0w0 Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions example/gpt2/config.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ inline nn::TransformerConfig GPT2Config() {
.n_head = 12,
.n_kv_head = 12,
.n_embd = 768,
.attention_type = nn::AttentionType::kStandard,
.position_embedding_type = nn::PositionEmbeddingType::kLearnedAbsolute,
.activation_type = nn::MLPType::kGELU,
.norm_type = nn::NormType::kLayerNorm,
.add_bias_linear = true,
Expand All @@ -34,7 +34,8 @@ inline void SanitizeGPT2Config(const nn::TransformerConfig &c) {
CHECK_GT(c.n_embd, 0);
CHECK_EQ(c.n_embd % c.n_head, 0) << "n_embd must be divisible by n_head";
CHECK_EQ(c.n_kv_head, c.n_head) << "GPT-2 does not use GQA; n_kv_head must equal n_head";
CHECK(c.attention_type == nn::AttentionType::kStandard) << "GPT-2 requires standard attention";
CHECK(c.position_embedding_type == nn::PositionEmbeddingType::kLearnedAbsolute)
<< "GPT-2 requires learned absolute position embedding";
CHECK(c.activation_type == nn::MLPType::kGELU) << "GPT-2 requires GELU activation";
CHECK(c.norm_type == nn::NormType::kLayerNorm) << "GPT-2 requires LayerNorm";
}
Expand Down
4 changes: 2 additions & 2 deletions example/llama3/config.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ inline nn::TransformerConfig LLaMA3Config() {
.n_head = 32,
.n_kv_head = 8,
.n_embd = 2048,
.attention_type = nn::AttentionType::kRoPE,
.position_embedding_type = nn::PositionEmbeddingType::kRoPE,
.activation_type = nn::MLPType::kSwiGLU,
.norm_type = nn::NormType::kRMSNorm,
.add_bias_linear = false,
Expand All @@ -36,7 +36,7 @@ inline void SanitizeLLaMA3Config(const nn::TransformerConfig &c) {
CHECK_EQ(c.n_head % c.n_kv_head, 0) << "n_head must be divisible by n_kv_head for GQA";
CHECK_GT(c.n_embd, 0);
CHECK_EQ(c.n_embd % c.n_head, 0) << "n_embd must be divisible by n_head";
CHECK(c.attention_type == nn::AttentionType::kRoPE) << "LLaMA-3 requires RoPE attention";
CHECK(c.position_embedding_type == nn::PositionEmbeddingType::kRoPE) << "LLaMA-3 requires RoPE position embedding";
CHECK(c.activation_type == nn::MLPType::kSwiGLU) << "LLaMA-3 requires SwiGLU activation";
CHECK(c.norm_type == nn::NormType::kRMSNorm) << "LLaMA-3 requires RMSNorm";
CHECK(!c.add_bias_linear) << "LLaMA-3 has no bias in linear layers";
Expand Down
5 changes: 3 additions & 2 deletions example/mixtral/config.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ inline nn::TransformerConfig TinyMixtralConfig() {
config.n_head = 4; // Scaled down; preserves Mixtral 4:1 GQA ratio.
config.n_kv_head = 1; // Scaled down with n_head; real Mixtral uses 8 KV heads.
config.n_embd = 512; // Scaled down from Mixtral/Megatron --hidden-size 4096.
config.attention_type = nn::AttentionType::kRoPE;
config.position_embedding_type = nn::PositionEmbeddingType::kRoPE;
config.activation_type = nn::MLPType::kSwiGLU;
config.ffn_type = nn::FFNType::kMoE;
config.norm_type = nn::NormType::kRMSNorm;
Expand Down Expand Up @@ -51,7 +51,8 @@ inline void SanitizeTinyMixtralConfig(const nn::TransformerConfig &c) {
CHECK_EQ(c.n_head % c.n_kv_head, 0) << "n_head must be divisible by n_kv_head for GQA";
CHECK_GT(c.n_embd, 0);
CHECK_EQ(c.n_embd % c.n_head, 0) << "n_embd must be divisible by n_head";
CHECK(c.attention_type == nn::AttentionType::kRoPE) << "tiny Mixtral requires RoPE attention";
CHECK(c.position_embedding_type == nn::PositionEmbeddingType::kRoPE)
<< "tiny Mixtral requires RoPE position embedding";
CHECK(c.activation_type == nn::MLPType::kSwiGLU) << "tiny Mixtral requires SwiGLU activation";
CHECK(c.ffn_type == nn::FFNType::kMoE) << "tiny Mixtral requires MoE FFN";
CHECK(c.norm_type == nn::NormType::kRMSNorm) << "tiny Mixtral requires RMSNorm";
Expand Down
15 changes: 0 additions & 15 deletions infini_train/include/nn/modules/transformer/causal_self_attention.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
#pragma once

#include <memory>
#include <tuple>
#include <vector>

#include "infini_train/include/nn/modules/module.h"
Expand Down Expand Up @@ -35,20 +34,6 @@ class CausalSelfAttention : public infini_train::nn::CloneableModule<CausalSelfA
// Setup method for different attention modes
void SetupAttention(const TransformerConfig &config);

// Standard attention forward (GPT2 style: no RoPE, no GQA)
std::vector<std::shared_ptr<infini_train::Tensor>>
ForwardStandard(const std::vector<std::shared_ptr<infini_train::Tensor>> &x);

// RoPE-aware attention forward (LLaMA3 style: with RoPE, optional GQA)
std::vector<std::shared_ptr<infini_train::Tensor>>
ForwardWithRoPE(const std::vector<std::shared_ptr<infini_train::Tensor>> &x);

// RoPE helper methods
std::tuple<std::shared_ptr<infini_train::Tensor>, std::shared_ptr<infini_train::Tensor>>
ApplyRotaryEmbedding(const std::shared_ptr<infini_train::Tensor> &xq,
const std::shared_ptr<infini_train::Tensor> &xk,
const std::shared_ptr<infini_train::Tensor> &freqs_cis);

// GQA helper method
std::shared_ptr<infini_train::Tensor> RepeatKV(const std::shared_ptr<infini_train::Tensor> &x, int64_t n_rep);
};
Expand Down
18 changes: 11 additions & 7 deletions infini_train/include/nn/modules/transformer/transformer_config.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,13 @@

namespace infini_train::nn {

enum class AttentionType {
kStandard, // Standard attention
kRoPE // Rotary Position Embedding
enum class PositionEmbeddingType {
kLearnedAbsolute, // Megatron: learned_absolute
kRoPE, // Megatron: rope
kYarn, // Megatron: yarn
kMRoPE, // Megatron: mrope
kRelative, // Megatron: relative
kNone // Megatron: none
};

enum class MLPType {
Expand Down Expand Up @@ -59,10 +63,10 @@ struct TransformerConfig {
int64_t n_kv_head = 12; // Num of Key/Value heads (<= n_head, < n_head if using GQA)
int64_t n_embd = 768; // Hidden size

AttentionType attention_type = AttentionType::kStandard; // Attention mechanism type
MLPType activation_type = MLPType::kGELU; // MLP activation type
FFNType ffn_type = FFNType::kDense; // Feed-forward module type
NormType norm_type = NormType::kLayerNorm; // Normalization type
PositionEmbeddingType position_embedding_type = PositionEmbeddingType::kLearnedAbsolute; // Position embedding type
MLPType activation_type = MLPType::kGELU; // MLP activation type
FFNType ffn_type = FFNType::kDense; // Feed-forward module type
NormType norm_type = NormType::kLayerNorm; // Normalization type

bool add_bias_linear = true; // Whether to add learnable bias to all Linear layers in the Transformer block,
// including: attention QKV projection, attention output projection, MLP FC layers (and
Expand Down
6 changes: 6 additions & 0 deletions infini_train/include/nn/modules/transformer/utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,17 @@
#pragma once

#include <cstdint>
#include <memory>
#include <tuple>

#include "infini_train/include/tensor.h"

namespace infini_train {
// RoPE helper method
std::shared_ptr<Tensor> PrecomputeFreqsCis(int64_t dim, int64_t end, float theta = 10000.0f, bool use_scaled = false,
Device device = Device());

std::tuple<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>>
ApplyRotaryEmbedding(const std::shared_ptr<Tensor> &xq, const std::shared_ptr<Tensor> &xk,
const std::shared_ptr<Tensor> &freqs_cis);
} // namespace infini_train
120 changes: 18 additions & 102 deletions infini_train/src/nn/modules/transformer/causal_self_attention.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#include "infini_train/include/nn/modules/transformer/causal_self_attention.h"

#include <cmath>
#include <limits>
#include <memory>
#include <tuple>
#include <vector>
Expand All @@ -12,6 +13,7 @@
#include "infini_train/include/nn/modules/normalization.h"
#include "infini_train/include/nn/modules/sparse.h"
#include "infini_train/include/nn/modules/transformer/transformer_config.h"
#include "infini_train/include/nn/modules/transformer/utils.h"
#include "infini_train/include/nn/parallel/global.h"
#include "infini_train/include/nn/parallel/tensor_parallel.h"
#include "infini_train/include/tensor.h"
Expand Down Expand Up @@ -42,9 +44,9 @@ CausalSelfAttention::CausalSelfAttention(const TransformerConfig &config) : Clon
/*skip_bias_add=*/false,
/*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled());

// For standard attention (GPT2 style), precompute causal mask
if (config_.attention_type == AttentionType::kStandard) {
// causal mask: (1, 1, block_size, block_size)
// FIXME(zbl): Decouple causal-mask ownership from position embedding. For now, only learned-absolute models use
// this precomputed buffer; RoPE callers provide a runtime-sized mask.
if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) {
buffers_[kParamBiasName] = function::Tril(nn::function::Ones({config_.block_size, config_.block_size}))
->View({1, 1, config_.block_size, config_.block_size});
}
Expand Down Expand Up @@ -74,98 +76,6 @@ void CausalSelfAttention::SetupAttention(const TransformerConfig &config) {
}
}

std::vector<std::shared_ptr<infini_train::Tensor>>
CausalSelfAttention::Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) {
if (config_.attention_type == AttentionType::kRoPE) {
return ForwardWithRoPE(x);
} else {
return ForwardStandard(x);
}
}

std::vector<std::shared_ptr<infini_train::Tensor>>
CausalSelfAttention::ForwardStandard(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) {
auto tp_world_size = parallel::global::GetTensorParallelSize();

const auto B = x[0]->Dims()[0]; // bs
const int64_t head_dim = n_embd_ / n_head_; // per-head dim (global)
const int64_t local_C = n_embd_ / tp_world_size; // per-rank hidden

// (B, T, C) -> ColumnParallelLinear(C, 3*C) -> (B, T, 3 * local_C)
// -> Split -> (3, B, T, local_C)
auto qkv = (*modules_[kCAttnLayerName])(x)[0]->Split(local_C, 2);

// (B, T, local_C)
auto q = qkv[0];
auto k = qkv[1];
auto v = qkv[2];

// NOTE(zbl): Acquire full T after AllGather is performed in ColumnParallelLinear
const auto T = q->Dims()[1];

// View to multi-head: local_n_head * head_dim == local_C
// (B, T, local_C) -> (B, T, h_l, Dh) -> (B, h_l, T, Dh)
k = k->View({B, T, local_n_head_, head_dim})->Transpose(1, 2);
q = q->View({B, T, local_n_head_, head_dim})->Transpose(1, 2);
v = v->View({B, T, local_n_head_, head_dim})->Transpose(1, 2);

// (B, h_l, T, T)
auto att = q->Matmul(k->Transpose(-2, -1)) * (1.0 / std::sqrt(head_dim));
// (1, 1, T, T)
auto mask = buffers_[kParamBiasName]->Slice({0, 0, 0, 0}, {1, 1, T, T}, {1, 1, 1, 1});
// (1, 1, T, T) -> eq 0 -> (1, 1, T, T) -> masked_fill -> (B, h_l, T, T)
att = att->MaskedFill(mask == 0, -std::numeric_limits<float>::infinity());
// (B, h_l, T, T)
att = nn::function::Softmax(att, -1);
// (B, h_l, T, Dh)
auto y = att->Matmul(v);
// (B, h_l, T, Dh) -> (B, T, h_l, Dh) -> (B, T, local_C)
y = y->Transpose(1, 2)->Contiguous()->View({B, T, local_C});

// Get full tensor
// (B, T, local_C) -> RowParallelLinear(n_embd, n_embd) -> (B, T, C)
y = (*modules_[kCProjLayerName])({y})[0];
// (B, T, C) == (bs, seq_len, n_embd)
return {y};
}

// RoPE helper methods
std::tuple<std::shared_ptr<infini_train::Tensor>, std::shared_ptr<infini_train::Tensor>>
CausalSelfAttention::ApplyRotaryEmbedding(const std::shared_ptr<infini_train::Tensor> &xq,
const std::shared_ptr<infini_train::Tensor> &xk,
const std::shared_ptr<infini_train::Tensor> &freqs_cis) {
// Reshape freqs_cis for broadcasting
const auto &x_shape = xq->Dims(); // (B, T, H, D)
const int64_t T = x_shape[1];
const int64_t D = x_shape[3];

std::vector<int64_t> target_shape = {1, T, 1, D / 2, 2};
auto cos_sin = freqs_cis->View(target_shape); // -> (1, T, 1, D/2, 2)

auto cos = cos_sin->Slice(-1, 0, 1, 1)->Squeeze(-1); // (1, T, 1, D/2)
auto sin = cos_sin->Slice(-1, 1, 2, 1)->Squeeze(-1); // (1, T, 1, D/2)

auto slice_pair = [](const std::shared_ptr<Tensor> &x) {
auto even = x->Slice(-1, 0, x->Dims().back(), 2);
auto odd = x->Slice(-1, 1, x->Dims().back(), 2);
return std::make_pair(even, odd);
};

auto [q_even, q_odd] = slice_pair(xq);
auto q_rotated_left = q_even * cos - q_odd * sin;
auto q_rotated_right = q_even * sin + q_odd * cos;
auto q_rotated
= nn::function::Stack(std::vector<std::shared_ptr<Tensor>>{q_rotated_left, q_rotated_right}, -1)->Flatten(-2);

auto [k_even, k_odd] = slice_pair(xk);
auto k_rotated_left = k_even * cos - k_odd * sin;
auto k_rotated_right = k_even * sin + k_odd * cos;
auto k_rotated
= nn::function::Stack(std::vector<std::shared_ptr<Tensor>>{k_rotated_left, k_rotated_right}, -1)->Flatten(-2);

return {q_rotated, k_rotated};
}

std::shared_ptr<infini_train::Tensor> CausalSelfAttention::RepeatKV(const std::shared_ptr<infini_train::Tensor> &x,
int64_t n_rep) {
const auto &shape = x->Dims();
Expand All @@ -179,21 +89,23 @@ std::shared_ptr<infini_train::Tensor> CausalSelfAttention::RepeatKV(const std::s
}

std::vector<std::shared_ptr<infini_train::Tensor>>
CausalSelfAttention::ForwardWithRoPE(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) {
CausalSelfAttention::Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) {
const auto B = x[0]->Dims()[0]; // bs
const auto C = x[0]->Dims()[2]; // n_embd

const auto tp_size = nn::parallel::global::GetTensorParallelSize();

const auto C_local = C / tp_size;
const auto H_local = n_head_ / tp_size;
const auto H_local = local_n_head_;
const auto KV_local = n_kv_head_ / tp_size;
const auto D = head_dim_; // n_embd / n_head

const auto freqs_cis = x.size() > 1 ? x[1] : nullptr;
const auto start_pos = x.size() > 2 ? x[2] : nullptr;
const auto mask = x.size() > 3 ? x[3] : nullptr;
CHECK(freqs_cis != nullptr) << "freqs_cis is null.";
if (config_.position_embedding_type == PositionEmbeddingType::kRoPE) {
CHECK(freqs_cis != nullptr) << "freqs_cis is null.";
}

// (B, T, C) -> (B, T, (H + 2 * n_kv_head) * D)
auto qkv = (*modules_[kCAttnLayerName])({x[0]})[0];
Expand All @@ -211,10 +123,10 @@ CausalSelfAttention::ForwardWithRoPE(const std::vector<std::shared_ptr<infini_tr
// v: (B, T, KV_local, D)
auto v = qkv->Slice(2, q_size_local + kv_size_local, q_size_local + 2 * kv_size_local)->View({B, T, KV_local, D});

// -> RoPE on q, k
// q: (B, T, H_local, D)
// k: (B, T, KV_local, D)
std::tie(q, k) = ApplyRotaryEmbedding(q, k, freqs_cis);
if (config_.position_embedding_type == PositionEmbeddingType::kRoPE) {
// q: (B, T, H_local, D), k: (B, T, KV_local, D)
std::tie(q, k) = ApplyRotaryEmbedding(q, k, freqs_cis);
}

// TODO(zbl): use kv cache during inference
// if (use_kv_) { ... }
Expand Down Expand Up @@ -242,6 +154,10 @@ CausalSelfAttention::ForwardWithRoPE(const std::vector<std::shared_ptr<infini_tr
if (mask) {
// mask: (1, 1, T, T)
att = att->MaskedFill(mask, std::numeric_limits<float>::lowest());
} else {
// fallback causal mask: (1, 1, T, T)
auto causal_mask = buffers_[kParamBiasName]->Slice({0, 0, 0, 0}, {1, 1, T, T}, {1, 1, 1, 1});
att = att->MaskedFill(causal_mask == 0, -std::numeric_limits<float>::infinity());
}
// (B, H_local, T, T)
att = nn::function::Softmax(att, -1);
Expand Down
22 changes: 13 additions & 9 deletions infini_train/src/nn/modules/transformer/transformer.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,11 @@ TransformerFirstStage::TransformerFirstStage(const TransformerConfig &config)
modules_[kWTELayerName] = std::make_shared<parallel::VocabParallelEmbedding>(
config_.vocab_size, config_.n_embd, parallel::global::GetSequenceParallelEnabled());

// LLaMA3 use RoPE, so they don't need position embedding
if (config_.activation_type == MLPType::kGELU) {
// Only learned absolute position embedding uses a trainable WPE table.
if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) {
modules_[kWPELayerName] = std::make_shared<Embedding>(config_.block_size, config_.n_embd);
} else if (config_.position_embedding_type != PositionEmbeddingType::kRoPE) {
LOG(FATAL) << "Unsupported position embedding type";
}
}

Expand All @@ -45,7 +47,7 @@ std::vector<std::shared_ptr<Tensor>> TransformerFirstStage::Forward(const std::v
// (B, T) -> Embedding(V_local, C) -> (B, T, C)
auto tok_emb = (*modules_[kWTELayerName])({x1});

// Add position embedding only for models that use absolute position encoding
// Add position embedding only for models that use learned absolute position encoding.
if (modules_.contains(kWPELayerName)) {
// (T_local)
// NOTE(zbl): Slice pos sequence when SP is enabled
Expand All @@ -66,7 +68,7 @@ std::vector<std::shared_ptr<Tensor>> TransformerFirstStage::Forward(const std::v
// (B, T, C)
return {tok_emb[0] + pos_emb[0]};
} else {
// For RoPE-based models (LLaMA3), no position embedding needed
// For RoPE-based models (LLaMA3), no absolute position embedding is needed.
// (B, T, C)
return tok_emb;
}
Expand Down Expand Up @@ -133,8 +135,8 @@ TransformerChunk::TransformerChunk(const TransformerConfig &config, int start_la
std::vector<std::shared_ptr<Tensor>> TransformerChunk::Forward(const std::vector<std::shared_ptr<Tensor>> &x) {
auto x1 = x[0];

// Check if we need to pass RoPE parameters (for LLaMA3 style models)
if (config_.attention_type == AttentionType::kRoPE) {
// Check if we need to pass RoPE parameters (for LLaMA3 style models).
if (config_.position_embedding_type == PositionEmbeddingType::kRoPE) {
// For RoPE models, we need to prepare freqs_cis and potentially other parameters
const auto device = x1->GetDevice();

Expand All @@ -161,9 +163,11 @@ std::vector<std::shared_ptr<Tensor>> TransformerChunk::Forward(const std::vector
for (auto &h : *std::dynamic_pointer_cast<nn::ModuleList>(modules_[kHLayerName])) {
x1 = (*h)({x1, freqs_view, start_pos_ptr, mask})[0];
}
} else {
// Standard attention (GPT2 style)
} else if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) {
// Learned absolute position embedding models (GPT-2 style).
for (auto &h : *std::dynamic_pointer_cast<nn::ModuleList>(modules_[kHLayerName])) { x1 = (*h)({x1})[0]; }
} else {
LOG(FATAL) << "Unsupported position embedding type";
}

return {x1};
Expand Down Expand Up @@ -217,7 +221,7 @@ TransformerModel::TransformerModel(const TransformerConfig config)
modules_[kPPFirstStageName] = std::make_shared<TransformerFirstStage>(config_);
transformer[TransformerFirstStage::kWTELayerName]
= modules_[kPPFirstStageName]->mutable_module(TransformerFirstStage::kWTELayerName);
if (config_.attention_type == AttentionType::kStandard) {
if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) {
transformer[TransformerFirstStage::kWPELayerName]
= modules_[kPPFirstStageName]->mutable_module(TransformerFirstStage::kWPELayerName);
}
Expand Down
Loading
Loading