Skip to content

[LLM:Feature] Add segment operator runtime support#4554

Merged
wangzhaode merged 1 commit into
alibaba:masterfrom
jxt1234:feature/llm_op_runtime_support
Jun 18, 2026
Merged

[LLM:Feature] Add segment operator runtime support#4554
wangzhaode merged 1 commit into
alibaba:masterfrom
jxt1234:feature/llm_op_runtime_support

Conversation

@jxt1234

@jxt1234 jxt1234 commented Jun 18, 2026

Copy link
Copy Markdown
Collaborator

Description

Module

Type

  • Feature
  • Bugfix
  • Perf
  • Refact
  • Style
  • Doc
  • Test
  • Chore

Checklist

  • Commit message follows [Module:Type] Description format
  • Code compiles without errors
  • Tested on relevant platform(s)
  • No unrelated format or style changes included

@wangzhaode

Copy link
Copy Markdown
Collaborator

这个 PR 目前看起来基于较旧的 master,diff 里会混入一些已经在最新 master 上修正过的表观回退。建议先 rebase 到最新 master,这样 review 可以只聚焦在本 PR 真正新增的 segment operator/runtime support 相关改动上。

我按当前 master...PR 重新看了一下,rebase 后大部分无关 diff 应该会自然消失;到时候再重点看剩余的 operator correctness / backend 覆盖会更准确。

@jxt1234
jxt1234 force-pushed the feature/llm_op_runtime_support branch from 20a9389 to 8ecf6b8 Compare June 18, 2026 08:55
@wangzhaode

Copy link
Copy Markdown
Collaborator

rebase 后我重新看了一遍,之前那些无关 diff 已经消失了。当前我只看到一个需要合并前确认/修掉的问题:

MUL_SILU CPU path 作为普通 BinaryOp 暴露,但实现没有处理 broadcast。

express/MathOp.cpp 里的 _MulSilu_Binary(x, y, BinaryOpOperation_MUL_SILU),shape 侧仍会走普通 BinaryOp 的 broadcast 规则;但 source/backend/cpu/CPUBinary.cpp 里的 MulSilu::onExecute 直接用 input0 的 tensor size 分片,并按相同 offset 顺序读取 input1

auto size = static_cast<CPUBackend *>(backend())->getTensorSize(input0);
...
auto inp = i0 + start * bytes;
auto inp1 = i1 + start * bytes;
...
mSilu((float *)out, (float *)inp1, realSize);
mMul((float *)out, (float *)out, (float *)inp, realSize, -1);

如果出现 [B, H] * [H][B, H] * [1] 或其他 broadcast shape,就会读过 input1 或只计算 input0 的物理大小,而不是 output 的广播后大小。

建议二选一:

  1. 复用/扩展现有 CPUBinary 的 broadcast 逻辑来实现 MUL_SILU;或
  2. 如果这个 fused op 只允许 LLM MLP 中 gate/up 同 shape 输入,就在 shape/creator 层显式限制同 shape,并补一个测试覆盖该约束。

另外 git diff --check 还有几处 trailing whitespace,主要在 MetalRope.mm 和生成的 OpenCL source map 文件里,合并前顺手清一下即可。

@wangzhaode

Copy link
Copy Markdown
Collaborator

One more thing from the second pass after the rebase:

source/backend/metal/MetalTopKV2.mm now rejects k > kTopKLocalK in the Metal creator (kTopKLocalK = 16). The previous Metal implementation accepted up to kTopKThreadNumber * kTopKLocalK (128 * 8 = 1024), and the existing test/op/TopKV2Test.cpp uses K = 300. With this change, those larger-k cases will no longer run on Metal and should fall back to CPU instead.

For the LLM top-1 use case this is probably fine, but as a generic TopKV2 backend change it looks like a Metal capability/performance regression. Could you either keep the old large-k path for k > 16, or make it explicit that this PR intentionally only supports the small-k Metal path and relies on CPU fallback for larger k?

@wangzhaode wangzhaode self-assigned this Jun 18, 2026
@wangzhaode
wangzhaode merged commit c0a23ab into alibaba:master Jun 18, 2026
9 checks passed
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.

2 participants