feat(tree-native): add tree-attention metadata + position_ids surface for tree training

[racky-scitix]

Summary

Adds the minimal Megatron-LM surface that the slime-trainer tree-training path
needs to flow tree-attention topology and per-token positional ids from the
data iterator down to TE's TreeFlashAttention backend. All changes are
opt-in: existing callers see no behavior change because every new field
defaults to None and the new AttnBackend.tree enum is only selected by
slime's tree path.

Scope (5 files, +91 / -2)

• megatron/core/transformer/tree_metadata.py (new) — TreeMetadata
  dataclass carrying cu_node_lens, node_parent, tree_position_ids,
  padded_size, num_nodes, and an opaque precomputed dict reused across
  layers. Lightweight, no CUDA assumptions at definition time.
• megatron/core/packed_seq_params.py — adds
  PackedSeqParams.tree_metadata: Optional[TreeMetadata] = None. Slime's
  TreeDataIterator populates it when --enable-tree-training is on; all
  existing callers keep None.
• megatron/core/transformer/enums.py — adds AttnBackend.tree = 6 for
  data-driven dispatch (TEDotProductAttention detects tree_metadata and
  routes to TE's TreeFlashAttention; cp_size == 1 required).
• megatron/core/models/common/embeddings/rotary_pos_embedding.py —
  RotaryEmbedding.forward accepts an optional position_ids tensor and,
  when provided, returns emb[position_ids] so each packed token gets the
  RoPE frequency for its logical (tree) position. Includes a defensive
  bounds check against rotary_seq_len. Default behavior unchanged.
• megatron/core/models/gpt/gpt_model.py — threads the existing
  position_ids argument into self.rotary_pos_emb(...) so the new RoPE
  path receives it.

Design notes

• Data-driven dispatch. Earlier iterations introduced spec-level
  TreeFlashAttention / TETreeDotProductAttention classes; those have
  been removed. The branch now keeps the spec untouched and dispatches
  purely on whether PackedSeqParams.tree_metadata is populated.
• Generic position_ids surface. RoPE stays generic — it does not
  know about tree topology. Whatever produces tree_position_ids (slime's
  TreeDataIterator today) is responsible for getting them onto
  PackedSeqParams.position_ids and into GPTModel.forward(position_ids=...).
• No public API removed. No upstream-Megatron public contract changes;
  every new field is optional with a None default.

Commit history

8 commits on top of sirl-dev:

5b385ae58 feat(tree-native): add Megatron PackedSeqParams.tree_metadata + TreeMetadata dataclass
5ba1a419b feat(tree-native): add Megatron tree-attention extensions
252c31867 refactor(tree-native): slim TETreeDotProductAttention to TE delegation
b2dbf0024 refactor(tree-native): data-driven tree dispatch, eliminate spec-level tree classes
8b600e448 refactor(tree-native): remove tree-specific code from TEDotProductAttention
fd65546a0 refactor(tree-native): generic position_ids on PackedSeqParams for RoPE
7a6d6e2ae refactor(tree-native): position_ids flows through model.forward, not PackedSeqParams
1813b6948 refactor(tree-native): drop stray blank line in transformer_engine.py

The history is intentionally kept in narrative order rather than squashed —
the later refactor commits show how the surface was simplified down to the
small, additive net diff above. Squashing on merge is fine.

Test plan

• git diff origin/sirl-dev sirl-dev-tree-native --shortstat =
  5 files, +91 / -2; no upstream-Megatron public APIs removed.
• Net transformer_engine.py diff vs sirl-dev is empty (the lone
  blank-line touch from earlier was reverted by the last commit).
• Pinned and consumed by slime-trainer's tree-native lane
  (vendors/megatron-lm gitlink at 1813b6948); precision +
  benchmark validation lives in slime-trainer
  docs/engineering/tree-attention-{benchmark,precision}.md.
• Reviewer to confirm RotaryEmbedding.forward position_ids defaulting
  to None is the right shape for upstream consumers.

🤖 Generated with Claude Code