Skip to content

Refactor: Merge 2 MHA paths - #192

Merged
kilinchange merged 2 commits into
masterfrom
refactor/mla
Jul 30, 2026
Merged

Refactor: Merge 2 MHA paths#192
kilinchange merged 2 commits into
masterfrom
refactor/mla

Conversation

@Chamberlain0w0

@Chamberlain0w0 Chamberlain0w0 commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

从 PR #163 提出来的一部分,修改目前 MHA 实现。

  1. 原来的 TransformerConfig::attention_type = kStandard / kRoPE 不太合适,Megatron 及其他开源实现中通常把 attn_type 分为 self/cross。这块命名更改为 Megatron 中使用的 --position-embedding-type,可选值为 learned_absolute / rope / yarn / mrope / relative / none。相应地修改创建 WPE/apply rope 的相关条件判断。
  2. 删除了 CausalSelfAttention::ForwardStandardForwardWithRoPE 两条分支,合并成一个统一的 Forward。GQA 也被纳入统一路径。
  3. ApplyRotaryEmbeddingCausalSelfAttention 成员函数提到了 transformer utils.cc
  4. causal mask buffer 现在无论 learned absolute 还是 RoPE 都会初始化;如果外部没有传 mask,会 fallback 到内部 causal mask。这个对 RoPE 直接调用且不传 mask 的场景是一个小的行为统一。

@kilinchange

Copy link
Copy Markdown
Collaborator

麻烦贴一下测试结果截图。

@Chamberlain0w0

Copy link
Copy Markdown
Contributor Author

ctest:
image

@Chamberlain0w0

Copy link
Copy Markdown
Contributor Author

3025 上运行测试,3026 上运行对比:
image

image

@kilinchange
kilinchange merged commit cedff0d into master Jul 30, 2026
2 checks passed
@kilinchange
kilinchange deleted the refactor/mla branch July 30, 2026 01:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants