# [model, trainer, engine] feat: prefix-tree MAGI attention for verl SFT and RL (draft)

[arvyanh]

What does this PR do?

The RFC: #6401

Adds prefix-tree shared-prefix deduplication for verl SFT and GRPO training using MAGI attention. Token-by-token trie detection discovers shared prefixes across rollout samples without requiring rollout-side metadata. The flat token layout + attention rectangle spec is dispatched to MAGI calc_attn, which computes correct prefix-tree attention patterns while internally deduplicating shared KV tokens.

Key features

• Dynamic trie detection: build_tree_dynamic builds a compressed trie from input tokens — no prior knowledge of turn boundaries needed.
• Multi-level tree support: handles multi-turn rollout sharing with zero-length leaf nodes for samples that terminate at intermediate levels.
• MAGI integration: patches Megatron-LM's TEDotProductAttention through the SelfAttention → TransformerLayer → GPTModel chain to inject MAGI/flex attention.
• prefix-length aware dynamic bsz grouping and mbs reordering
• Configurable old-log-prob backend: Currently the old_log_prob's log_prob seems to diverge

Design

This PR implements the design described in RFC #6401. All n GRPO rollout samples are packed into a single flat [prefix | leaf_0 | ... | leaf_{n-1}] sequence, run through one transformer forward pass, with cross-leaf attention blocked. The result is mathematically equivalent to n independent forwards.

Attention rectangle spec

The attention pattern is encoded as rectangle specs:

          k: prefix    k: leaf0    k: leaf1
q: prefix   causal       ✗           ✗
q: leaf0     full      causal         ✗
q: leaf1     full        ✗         causal

MAGI interprets these rectangles natively.

Prefix detection (dynamic trie)

Prefix detection runs entirely at train time — no rollout-side metadata required. build_tree_dynamic performs token-by-token compressed trie insertion over the micro-batch's input sequences. The trie is converted to a TreeNode tree, then fed to build_prefix_tree_attention_spec, which emits (q_ranges, k_ranges, mask_types) for every node.

but in case we could retrieve the tree at runtime, we could provide the trie data to the traininer

The tree generalizes to arbitrary-depth multi-level trees. For example, in multi-turn agent RL where responses branch into sub-groups sharing a turn-2 prefix:

S0: turn0 + turn1_A + turn2_B1
S1: turn0 + turn1_A + turn2_B2
S2: turn0 + turn1_C + turn2_D1
S3: turn0 + turn1_C + turn2_D2

This produces a depth-3 tree with turn0 as root, turn1_A/turn1_C as intermediate nodes, and four leaves. Zero-length leaf nodes are inserted when a sample terminates at an intermediate node.

Megatron integration

The patch chain injects magi_attention_key through upstream Megatron-LM:

GPTModel.forward → TransformerBlock → TransformerLayer
    → SelfAttention → TEDotProductAttention (magi_attn_forward: dispatch → calc_attn → undispatch)

File summary

Area	Path
Trie + load balancing	verl/utils/prefix_tree/dynamic.py
Layout + attention spec	verl/utils/prefix_tree/utils.py
MAGI key + batch	verl/utils/prefix_tree/magi.py
Trainer helpers	verl/utils/prefix_tree/trainer.py
Megatron patch	verl/models/mcore/prefix_tree_merge.py

Result

[TODO]

Test

CPU unit tests (59 tests, all pass)

python3 prefix_script/test/test_prefix_tree_full.py  # temporary location

GPU training

Validated on Megatron-LM training with GRPO on CoQA — prefix-tree MAGI attention runs correctly with CP=2, SP=1, PP=1.

Checklist Before Starting

• Search for similar PRs: upstream verl prefix-tree PRs
• Format the PR title as [{modules}] {type}: {description}
• Read the Contribute Guide
• Apply pre-commit checks

[gemini-code-assist]

Code Review

This pull request introduces prefix-tree shared-prefix deduplication for actor and SFT training, adding dynamic-trie prefix-tree building, MAGI/flex attention integration, and trainer load-balancing utilities. Feedback on the changes highlights a critical logic modification in the non-prefix-tree fallback path within model_forward.py where the condition for passing position_ids was altered, as well as the high maintainability risks associated with the extensive monkey-patching of Megatron-LM core components in prefix_tree_merge.py.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The condition for passing position_ids to the model has been changed from mtp_enable_train to not vision_model. This is a significant logic change in the non-prefix-tree fallback path and appears unrelated to the main feature of this PR. This could have unintended consequences on model correctness or performance for non-vision models that are not using MTP, as they will now receive position_ids where they previously did not. Furthermore, the original comment # position_ids is only needed for MTP is now outdated and misleading. This change needs to be justified, and if it is an intentional bug fix, it should be explained and ideally handled in a separate PR.

Suggested change

	position_ids=position_ids_rmpad if not vision_model else None,
	position_ids=position_ids_rmpad if mtp_enable_train else None, # position_ids is only needed for MTP

[gemini-code-assist]

The extensive monkey-patching of multiple Megatron-LM core components (GPTModel, TransformerBlock, TransformerLayer, SelfAttention, TEDotProductAttention, RotaryEmbedding) introduces a significant maintainability risk. This approach is fragile and highly susceptible to breaking with future updates to the upstream Megatron-LM library. While the implementation is carefully written with sentinels to prevent re-patching and try...finally blocks to restore methods, a breaking change in the upstream API could lead to silent failures or hard-to-debug issues.

Consider exploring less intrusive integration methods, such as subclassing or dependency injection where possible.