Repository navigation
[Refactor] Rename layer boundary protocols, contracts and helpers - #41750
Merged
Merged
Conversation
ch-wan
requested review from
1am9trash,
BBuf,
Edwardf0t1,
Fridge003,
HaiShaw,
Jiminator,
JustinTong0323,
YAMY1234,
Ying1123,
fzyzcjy,
hnyls2002,
hubertlu-tw,
iforgetmyname,
ispobock,
kkHuang-amd,
merrymercy,
ping1jing2,
rainj-me,
sogalin,
whybeyoung,
wisclmy0611 and
zijiexia
as code owners
September 29, 2026 19:50
github-actions
Bot
added
documentation
This was referenced Sep 29, 2026
This was referenced Sep 29, 2026
ch-wan
force-pushed
the
cheng/hot-fix/release-consumed-contribution
branch
from
September 29, 2026 20:00
85357f5 to
be31274
Compare
ch-wan
force-pushed
the
cheng/refactor/layer-boundary-naming
branch
2 times, most recently
from
September 29, 2026 20:29
3bf8a5f to
5ff6682
Compare
Base automatically changed from
cheng/hot-fix/release-consumed-contribution
to
main
September 29, 2026 23:38
Contributor
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Automations to automatically generate PRs for you. |
ch-wan
force-pushed
the
cheng/refactor/layer-boundary-naming
branch
from
September 29, 2026 23:47
5ff6682 to
cb4fe23
Compare
ch-wan
force-pushed
the
cheng/refactor/layer-boundary-naming
branch
from
September 29, 2026 23:56
cb4fe23 to
a6de83d
Compare
This was referenced Sep 29, 2026
StageUpdate and StageRead describe how a producer writes its output into the residual and how a consumer reads its input from it, so name them ResidualUpdate and ResidualReadout. Their implementations and the LayerResidual bundle follow (PlainAdd, NormReadout, NormQuantReadout, LayerResidualOps), and the capability flags read as predicates (is_plain_add, applied_at_exit, outlives_layer, is_plain_norm, reads_before_dp_gather). The fusion eligibility checks are predicates too: flashinfer_ar_fusion_applies and aiter_ar_fusion_applies. Renames only; no behavior change.
ch-wan
force-pushed
the
cheng/refactor/layer-boundary-naming
branch
from
September 30, 2026 00:33
a6de83d to
661ceee
Compare
- Contracts: StageDecl, EdgeDecl, StageInput and StageOutput become StageContract, EdgeContract, InputContract and OutputContract, so the resolved contract no longer differs from StageDeclaration by one abbreviation. Output permissions read as permissions (always_partial, may_defer_to_next, may_reduce_scatter[v]), and ProducerReduction values state the contract (ALWAYS_PARTIAL, EXIT_SCOPED, TAIL_AFTER_SUM). - "handoff" now means only a layer handing its output on: the MoE finalize left for the next consumer is DeferredFinalize (finalize_norm on the final norm), the required output rows are ExitRows, and the attention input hook on an entry path is attn_input_adapter. - Bound paths: StageEntry/StageSteps/StageEdges become EntryPath, StagePath and VariantEdges; make_boundary/make_output_boundary become bind_entry/bind_exit. Entry steps are named by what they run in order (e.g. _reduce_update_read, _dp_gather_sum_read) instead of _mlp_input_*, since attention inputs use them too. Communication primitives are named axis then action (attn_tp_gather, dp_gather_sum, moe_cp_take_back, ...) instead of redistribute_*. - Exit and stream: OutputBoundary/FfnCompletion become ExitPolicy and ExitDecision; CarriedSum becomes DeferredSum next to DeclaredSum; ResidualStream.leave/finish/arrive become record/export/from_handoff, so export no longer shares a name with StageBoundary.finish. - StageBoundary.postprocess becomes finish_complete_output: its input is an output compute already reduced, and it still applies the output transform and a producer-side residual update. Renames only; no behavior change.
- The layer-stack helpers say what they do: residual_batch.norm becomes final_norm, written becomes set_written, and current becomes stream_of. - Attention prepare takes capture= (a callback on the producer's rows) and capture_gathered= (a collector on the compute-input rows) instead of capture_output= and captured_last_layer_outputs=. The first no longer shares its name with StageBoundary.capture_output(), which returns a snapshot. Models keep their own layer-forward argument names. - declare_ffn(next_sparse=) becomes next_layer_sparse, since it is the next decoder layer's FFN rather than the next stage, and declare_attn(gathers_tp_input=) becomes gathers_attn_tp_input. - Configuration queries read as queries: enable_moe_dense_fully_dp() becomes is_dense_ffn_fully_dp(), and the MoE-CP pair becomes moe_gathers_over_moe_cp() / batch_gathers_over_moe_cp(batch). - TokenAxis.ATTN_TP_SCATTER becomes ATTN_TP: membership in a layout already means the rows are split across that axis. Renames only; no behavior change.
ch-wan
force-pushed
the
cheng/refactor/layer-boundary-naming
branch
from
September 30, 2026 01:01
661ceee to
9a817ba
Compare
This was referenced Sep 30, 2026
3 of 5 tasks
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR is part of a stack (oldest at bottom):
Motivation
Several names in
layers/layer_boundarydid not say what the object does, and one concept often had several words:StageUpdate/StageReadare the operations that write a producer's output into the residual and read a consumer's input from it;LayerResidualis a bundle of those operations, not a tensor.StageDeclandStageDeclarationdiffered by one abbreviation but are different layers._mlp_input_*, and communication primitivesredistribute_*, after the old communicator.This PR renames them. It does not change behavior.
Modifications
Three commits, each renaming definitions together with every use, including the developer guide:
ResidualUpdate/ResidualReadout,PlainAdd(PLAIN_ADD),NormReadout/NormQuantReadout,LayerResidualOps(PLAIN_RESIDUAL_OPS); capability flags as predicates (is_plain_add,applied_at_exit,outlives_layer,is_plain_norm,reads_before_dp_gather);flashinfer_ar_fusion_applies/aiter_ar_fusion_applies.StageContract,EdgeContract,InputContract,OutputContractwithalways_partial/may_defer_to_next/may_reduce_scatter[v];ProducerReduction.ALWAYS_PARTIAL/EXIT_SCOPED/TAIL_AFTER_SUM;DeferredFinalize(finalize_norm=on the final norm),ExitRows,attn_input_adapter;EntryPath,StagePath,bind_entry/bind_exit; entry steps named by what they run (_reduce_update_read,_dp_gather_sum_read, ...); primitives named axis then action (attn_tp_gather,dp_gather_sum,moe_cp_take_back, ...);ExitPolicy,ExitDecision,DeferredSum;ResidualStream.record/export/from_handoff;StageBoundary.finish_complete_output(waspostprocess).residual_batch.final_norm/set_written/stream_of;prepare(capture=, capture_gathered=)(models keep their own layer-forward argument names);declare_ffn(next_layer_sparse=),declare_attn(gathers_attn_tp_input=);is_dense_ffn_fully_dp(),moe_gathers_over_moe_cp()/batch_gathers_over_moe_cp(batch);TokenAxis.ATTN_TP.Not renamed: the forward flags read by compute modules (
fuse_mlp_allreduce,mlp_reduce_scatter,defer_moe_finalize),AttnTpContext/AttentionInputs, server arguments,StageBoundary.capture_output(), and generic contract fields (layout,group,read).Accuracy Tests
Renames only. Each commit was checked against its parent by walking both ASTs file by file: the structure is identical and every differing name or string is one of the renames. The only other differences are one import statement that isort merged and two docstrings describing the renamed capture arguments. Every changed module imports on each commit, and 97 affected unit test files (2,081 tests) pass on each commit.
Speed Tests and Profiling
Not applicable: no behavior change.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #36653139245
Latest PR Test (Extra): ❌ Run #36653139182
Latest PR Test (AMD ROCm 10): ❌ Run #36653139349