diff --git a/examples/llm_finetune/qwen/qwen25_magi_prefix_tree_rollouts.yaml b/examples/llm_finetune/qwen/qwen25_magi_prefix_tree_rollouts.yaml
new file mode 100644
index 0000000000..161c8b9fae
--- /dev/null
+++ b/examples/llm_finetune/qwen/qwen25_magi_prefix_tree_rollouts.yaml
@@ -0,0 +1,108 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+# To run this recipe:
+#   torchrun --nproc-per-node=1 examples/llm_finetune/finetune.py --config examples/llm_finetune/qwen/qwen25_magi_prefix_tree_rollouts.yaml
+#
+# Shared-prefix RL rollout finetune with the MagiAttention (FFA) backend and a
+# block-sparse prefix-tree mask, WITHOUT context parallelism (cp_size=1). Each
+# training group is one shared prompt + N sampled completions; the collate folds
+# them into a single deduplicated flat sequence ([prompt | c_0 | c_1 | ...]) and
+# builds the prefix-tree AttnMaskSpec. Every completion attends FULL to the shared
+# prompt and CAUSAL to itself, so the prompt is encoded once instead of N times.
+# This is the verl RFC #6401 / Automodel #2385 layout (cp=1 only). Pair with the
+# Qwen3 MoE magi examples to compare against plain varlen packing.
+#
+# Notes for the magi prefix-tree backend (cp=1):
+#   * the model MUST be a registered custom model (in MODEL_ARCH_MAPPING) so
+#     backend.attn="magi" wires the custom-model attn_func, which is the only
+#     magi path that reads the out-of-band AttnMaskSpec. Qwen2 is registered.
+#     A plain HF model routes through the HF magi backend (attn_implementation=
+#     "magi"), whose fixed attention interface cannot receive the mask spec; the
+#     recipe raises NotImplementedError in that case rather than silently
+#     dropping the prefix-tree mask.
+#   * local_batch_size must be 1: each folded group is already one flat sequence
+#     with no batch dim, and the prefix-tree mask is built per group.
+#   * backend.rope_fusion must be false: magi applies RoPE positionally with the
+#     per-completion position ids the collate emits (each completion continues
+#     from the prompt length), so the fused TE rope path is bypassed.
+#   * the recipe hands the per-step AttnMaskSpec to the magi attn_func out-of-band
+#     via set_active_attn_spec; no cu_seqlens / THD packer is used here.
+
+recipe: TrainFinetuneRecipeForNextTokenPrediction
+
+step_scheduler:
+  global_batch_size: 8
+  local_batch_size: 1        # one folded rollout group per micro-step (no batch dim)
+  ckpt_every_steps: 500
+  num_epochs: 2
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 20
+
+rng:
+  _target_: nemo_automodel.components.training.rng.StatefulRNG
+  seed: 1111
+  ranked: true
+
+model:
+  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
+  pretrained_model_name_or_path: Qwen/Qwen2.5-0.5B  # registered custom model -> magi custom attn_func
+  backend:
+    _target_: nemo_automodel.components.models.common.BackendConfig
+    attn: magi              # MagiAttention Flex-Flash-Attention backend
+    linear: torch
+    rms_norm: torch
+    rope_fusion: false      # required for magi (non-fused, positional RoPE)
+    enable_hf_state_dict_adapter: true
+
+checkpoint:
+  enabled: false
+  checkpoint_dir: checkpoints/
+  model_save_format: torch_save
+  save_consolidated: false
+
+distributed:
+  strategy: fsdp2
+  tp_size: 1
+  cp_size: 1                 # prefix-tree wiring is cp=1 only
+  pp_size: 1
+  ep_size: 1
+  sequence_parallel: false
+  activation_checkpointing: false
+
+loss_fn:
+  _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy
+
+dataset:
+  _target_: nemo_automodel.components.datasets.llm.mock_prefix_tree.build_mock_rollout_dataset
+  num_groups: 64
+  completions_per_group: 4
+  prompt_len: 64
+  completion_len: 32
+
+dataloader:
+  _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  collate_fn: nemo_automodel.components.datasets.llm.prefix_tree.prefix_tree_collate_fn
+  shuffle: false
+
+optimizer:
+  _target_: torch.optim.Adam
+  betas: [0.9, 0.999]
+  eps: 1e-7
+  lr: 1.0e-4
+  weight_decay: 0
+  foreach: false
diff --git a/nemo_automodel/components/datasets/llm/mock_prefix_tree.py b/nemo_automodel/components/datasets/llm/mock_prefix_tree.py
new file mode 100644
index 0000000000..7d0de67e64
--- /dev/null
+++ b/nemo_automodel/components/datasets/llm/mock_prefix_tree.py
@@ -0,0 +1,58 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Deterministic mock shared-prefix rollout data for prefix-tree smoke runs."""
+
+import random
+
+
+def build_mock_rollout_dataset(
+    *,
+    num_groups: int = 16,
+    completions_per_group: int = 4,
+    prompt_len: int = 32,
+    completion_len: int = 16,
+    vocab_size: int = 1024,
+    seed: int = 0,
+) -> list[dict]:
+    """Build a deterministic mock shared-prefix rollout dataset for smoke runs.
+
+    Each group is one shared prompt with ``completions_per_group`` completions, in
+    the ``{"prompt_ids", "completions"}`` schema consumed by
+    ``prefix_tree_collate_fn``. Token ids are random in ``[2, vocab_size)``; this
+    is a pipeline smoke, not a quality dataset.
+
+    Args:
+        num_groups: number of rollout groups.
+        completions_per_group: completions (leaves) sharing each prompt.
+        prompt_len: shared prompt length per group.
+        completion_len: length of each completion.
+        vocab_size: upper bound (exclusive) for random token ids.
+        seed: RNG seed for reproducibility.
+
+    Returns:
+        A list of ``{"prompt_ids": list[int], "completions": list[list[int]]}``.
+    """
+    rng = random.Random(seed)
+
+    def _ids(n: int) -> list[int]:
+        return [rng.randint(2, vocab_size - 1) for _ in range(n)]
+
+    return [
+        {
+            "prompt_ids": _ids(prompt_len),
+            "completions": [_ids(completion_len) for _ in range(completions_per_group)],
+        }
+        for _ in range(num_groups)
+    ]
diff --git a/nemo_automodel/components/datasets/llm/prefix_tree.py b/nemo_automodel/components/datasets/llm/prefix_tree.py
new file mode 100644
index 0000000000..0ac29bef02
--- /dev/null
+++ b/nemo_automodel/components/datasets/llm/prefix_tree.py
@@ -0,0 +1,167 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared-prefix rollout folding for multi-turn prefix-tree attention (cp=1).
+
+Folds a group of rollouts that share one prompt prefix (one prompt -> N sampled
+completions) into a single deduplicated flat token layout plus the prefix-tree
+structure (``node_lengths`` / ``sample_paths``). The shared prompt is stored
+once; every completion attends FULL to the prompt and CAUSAL to itself.
+
+This is the verl RFC #6401 / Automodel #2385 shared-prefix RL training layout,
+restricted to the cp=1 path (no context-parallel dispatch). The collate carries
+the structure on the batch; the magi backend builds the ``AttnMaskSpec`` from it
+and activates it (the datasets layer must not import ``components.distributed``).
+Enable it with ``model.backend.attn: magi``.
+
+Branch-point note: the shared prompt's final position is stored once, so it can
+predict only one next token. The N completions diverge there, so the
+prompt -> first-completion-token transition is left unsupervised (label
+``-100`` on the last prompt token); each completion is supervised causally from
+its own first token onward. This is inherent to deduplicating the shared prefix.
+"""
+
+from dataclasses import dataclass
+
+import torch
+
+CROSS_ENTROPY_IGNORE_IDX = -100
+
+
+@dataclass
+class FoldedRollouts:
+    """Deduplicated flat layout + prefix-tree structure for one shared-prefix group.
+
+    The attention mask itself is built in the magi backend from ``node_lengths`` /
+    ``sample_paths`` (via ``AttnMaskSpec.prefix_tree``); the datasets layer must not
+    import ``components.distributed``, so it only carries the structure.
+
+    Attributes:
+        input_ids: flat ``[prompt | completion_0 | completion_1 | ...]`` tokens.
+        labels: pre-shifted next-token targets (the loss does not shift); the
+            prompt and each completion's last token are ``-100`` (ignored).
+        position_ids: prompt positions ``0..P-1``; each completion continues from
+            ``P`` (so RoPE sees ``prompt ++ completion``).
+        node_lengths: token count of each node, flat-layout order
+            ``[len(prompt), len(c_0), len(c_1), ...]``.
+        sample_paths: root -> leaf node indices per completion, e.g.
+            ``[[0, 1], [0, 2], ...]``.
+    """
+
+    input_ids: list[int]
+    labels: list[int]
+    position_ids: list[int]
+    node_lengths: list[int]
+    sample_paths: list[list[int]]
+
+
+def fold_shared_prefix_rollouts(
+    prompt_ids: list[int],
+    completions: list[list[int]],
+    *,
+    ignore_idx: int = CROSS_ENTROPY_IGNORE_IDX,
+) -> FoldedRollouts:
+    """Fold one shared-prefix rollout group into a deduplicated prefix-tree layout.
+
+    Labels follow this repo's next-token convention: they are pre-shifted (the
+    loss does not shift), so position ``p`` carries the id of the token the model
+    should predict at ``p``. Within each completion, token ``t`` predicts token
+    ``t + 1``; the completion's last token has no in-layout successor and is
+    masked. The shared prompt is masked entirely, including its last position:
+    that position is the branch point where the N completions diverge, so it
+    cannot supervise any single first-completion token (the cost of deduplicating
+    the prefix is that each completion's first token is unsupervised).
+
+    Args:
+        prompt_ids: the shared prompt tokens (node 0). May be empty.
+        completions: one token-id list per sampled completion (one leaf each).
+            Must be non-empty and every completion must be non-empty.
+        ignore_idx: label value for unsupervised positions (default ``-100``).
+
+    Returns:
+        A :class:`FoldedRollouts` with the flat tokens, labels, position ids and
+        the built :class:`AttnMaskSpec`.
+
+    Raises:
+        ValueError: if ``completions`` is empty or any completion is empty.
+    """
+    if not completions:
+        raise ValueError("completions must be non-empty (need at least one rollout).")
+    if any(len(c) == 0 for c in completions):
+        raise ValueError("every completion must be non-empty.")
+
+    prompt_len = len(prompt_ids)
+
+    # Single pass builds the deduplicated flat layout: the prompt is stored once
+    # (masked, positions 0..P-1) and each completion is appended with positions
+    # continuing from P (so RoPE sees prompt ++ completion) and pre-shifted labels
+    # (token t predicts t+1; the last token is masked).
+    input_ids = list(prompt_ids)
+    labels = [ignore_idx] * prompt_len
+    position_ids = list(range(prompt_len))
+    node_lengths = [prompt_len]
+    for completion in completions:
+        input_ids.extend(completion)
+        labels.extend(completion[1:] + [ignore_idx])
+        position_ids.extend(range(prompt_len, prompt_len + len(completion)))
+        node_lengths.append(len(completion))
+
+    # Prefix tree: node 0 is the prompt, nodes 1..N are completions; each path is
+    # prompt -> completion. A bare group (empty prompt) collapses to varlen blocks:
+    # node 0 is empty, so each completion is its own causal block.
+    if prompt_len > 0:
+        sample_paths = [[0, i + 1] for i in range(len(completions))]
+    else:
+        sample_paths = [[i + 1] for i in range(len(completions))]
+
+    return FoldedRollouts(
+        input_ids=input_ids,
+        labels=labels,
+        position_ids=position_ids,
+        node_lengths=node_lengths,
+        sample_paths=sample_paths,
+    )
+
+
+def prefix_tree_collate_fn(batch: list[dict]) -> dict:
+    """Collate one shared-prefix rollout group into a model-ready batch (cp=1).
+
+    Folds the group with :func:`fold_shared_prefix_rollouts` and emits the flat
+    tokens plus the prefix-tree structure. Only ``local_batch_size == 1`` is
+    supported: each group already packs many completions into one flat sequence,
+    and the mask is per group. The ``prefix_tree`` entry is popped by the magi
+    backend (``MagiState.prepare_llm_batch``), which builds and activates the
+    ``AttnMaskSpec`` from it.
+
+    Args:
+        batch: a length-1 list holding one rollout group dict with keys
+            ``prompt_ids`` and ``completions``.
+
+    Returns:
+        Dict with ``input_ids``, ``labels``, ``position_ids`` (each ``[1, T]``)
+        and ``prefix_tree`` (``(node_lengths, sample_paths)``).
+
+    Raises:
+        ValueError: if ``batch`` does not hold exactly one rollout group.
+    """
+    if len(batch) != 1:
+        raise ValueError(f"prefix_tree_collate_fn supports local_batch_size=1 only, got {len(batch)} groups.")
+    group = batch[0]
+    folded = fold_shared_prefix_rollouts(group["prompt_ids"], group["completions"])
+    return {
+        "input_ids": torch.tensor([folded.input_ids], dtype=torch.long),
+        "labels": torch.tensor([folded.labels], dtype=torch.long),
+        "position_ids": torch.tensor([folded.position_ids], dtype=torch.long),
+        "prefix_tree": (folded.node_lengths, folded.sample_paths),
+    }
diff --git a/nemo_automodel/components/distributed/magi_attn_utils.py b/nemo_automodel/components/distributed/magi_attn_utils.py
index 3703dfb0f0..5b4cf52ed8 100644
--- a/nemo_automodel/components/distributed/magi_attn_utils.py
+++ b/nemo_automodel/components/distributed/magi_attn_utils.py
@@ -749,6 +749,28 @@ def prepare_llm_batch(
         Returns ``(train_ctx, batch)``. magi does its own CP, so ``train_ctx`` is
         always ``nullcontext`` (no torch-native DTensor CP context).
         """
+        # cp=1 prefix-tree mask: the datasets layer cannot import this module (component
+        # independence), so the collate attaches the tree structure and the spec is built
+        # and activated here, out-of-band, for the magi attn_func. Setting it every step
+        # (a spec or None) is self-clearing, so a stale spec never leaks into the next
+        # batch; plain batches omit "prefix_tree".
+        prefix_tree = batch.pop("prefix_tree", None)
+        if prefix_tree is not None and self.hf_dispatch:
+            # The prefix-tree mask is handed to the attn_func out-of-band (the HF
+            # attention interface has a fixed signature and cannot receive a custom
+            # mask spec argument). Only the custom-model attn_func reads it; the HF
+            # magi backend uses the plain causal varlen key from magi_prepare_batch
+            # and would silently drop the prefix-tree mask. Fail loudly instead.
+            raise NotImplementedError(
+                "The prefix-tree attention mask is only supported on the custom-model magi "
+                "backend (model.backend.attn='magi'), not the HF magi backend "
+                "(model.attn_implementation='magi'). HF's attention interface cannot receive "
+                "the out-of-band AttnMaskSpec, so the mask would be silently dropped. Load a "
+                "model registered in MODEL_ARCH_MAPPING and set model.backend.attn='magi'."
+            )
+        node_lengths, sample_paths = prefix_tree if prefix_tree is not None else (None, None)
+        spec = AttnMaskSpec.prefix_tree(node_lengths, sample_paths)[0] if prefix_tree is not None else None
+        set_active_attn_spec(spec)
         if self.hf_dispatch:
             # HF path: dispatch the (single causal) sequence across the CP group.
             batch, _ = magi_prepare_batch(model, batch, self.cp_group)
diff --git a/tests/functional_tests/attention/_prefix_tree_reference.py b/tests/functional_tests/attention/_prefix_tree_reference.py
new file mode 100644
index 0000000000..24198c2361
--- /dev/null
+++ b/tests/functional_tests/attention/_prefix_tree_reference.py
@@ -0,0 +1,48 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Independent oracle mask for the prefix-tree parity checks.
+
+Shared by ``prefix_tree_flex_parity.py`` and ``prefix_tree_magi_parity.py`` so
+the two checks validate against the same reference and cannot silently diverge.
+"""
+
+import torch
+
+
+def build_reference_mask(prompt_len: int, completion_lens: list[int], total: int) -> torch.Tensor:
+    """Dense boolean attend-mask, built straight from the rollout structure.
+
+    ``allowed[q, k]`` is True iff query ``q`` may attend to key ``k`` under the
+    "prompt ++ completion as an independent causal sequence" semantics:
+      * prompt query: causal within the prompt only;
+      * completion query: FULL to the prompt, CAUSAL within its own completion.
+
+    This is built from the rollout structure, NOT from the spec under test, so a
+    wrong spec or a wrong realization of it is caught.
+    """
+    node_of = [-1] * prompt_len
+    for c, n in enumerate(completion_lens):
+        node_of.extend([c] * n)
+    assert len(node_of) == total
+
+    allowed = torch.zeros(total, total, dtype=torch.bool)
+    for q in range(total):
+        qc = node_of[q]
+        for k in range(q + 1):  # causal upper bound: never attend to the future
+            kc = node_of[k]
+            if kc == -1 or kc == qc:
+                # prompt key (full for completions, causal for prompt) or same completion.
+                allowed[q, k] = True
+    return allowed
diff --git a/tests/functional_tests/attention/prefix_tree_flex_parity.py b/tests/functional_tests/attention/prefix_tree_flex_parity.py
new file mode 100644
index 0000000000..e5f43f0cfe
--- /dev/null
+++ b/tests/functional_tests/attention/prefix_tree_flex_parity.py
@@ -0,0 +1,117 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Backend-agnostic correctness check for the cp=1 prefix-tree mask.
+
+Realizes the folded :class:`AttnMaskSpec` as a ``flex_attention`` mask (pure
+torch, runs on Ampere/A100 and Hopper alike, no magi/FA4 build) and compares its
+output, position by position, against an independently constructed dense-mask
+SDPA reference. The reference is built straight from the rollout structure (each
+completion attends FULL to the prompt and CAUSAL to itself), NOT from the spec,
+so a wrong spec or a wrong realization of it is caught.
+
+This validates the part we added (fold rollouts -> AttnMaskSpec, realize the
+mask); the magi FFA kernel that consumes the same spec on Hopper is checked
+separately by ``prefix_tree_magi_parity.py``. Run with::
+
+    python tests/functional_tests/attention/prefix_tree_flex_parity.py
+
+Exit code 0 means parity holds.
+"""
+
+from __future__ import annotations
+
+import sys
+
+import torch
+import torch.nn.functional as F
+from _prefix_tree_reference import build_reference_mask
+from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+
+from nemo_automodel.components.datasets.llm.prefix_tree import fold_shared_prefix_rollouts
+from nemo_automodel.components.distributed.magi_attn_utils import AttnMaskSpec
+
+# fp32 both sides: a tight tolerance that flags any mask error rather than noise.
+MAX_DIFF_TOL = 1e-3
+COS_SIM_TOL = 0.9999
+
+
+def _spec_mask_mod(spec):
+    """Turn an :class:`AttnMaskSpec` (AttnSlice rectangles) into a flex mask_mod."""
+    rects = list(zip(spec.q_ranges, spec.k_ranges, spec.mask_types))
+
+    def mask_mod(b, h, q_idx, kv_idx):
+        keep = q_idx < 0  # all-False seed of the right broadcast shape/dtype
+        for (qs, qe), (ks, ke), mt in rects:
+            rect = (q_idx >= qs) & (q_idx < qe) & (kv_idx >= ks) & (kv_idx < ke)
+            if mt == "causal":
+                rect = rect & (kv_idx <= q_idx)
+            keep = keep | rect
+        return keep
+
+    return mask_mod
+
+
+def main() -> int:
+    if not torch.cuda.is_available():
+        print("SKIP: no CUDA device.")
+        return 0
+    torch.cuda.set_device(0)
+    device = torch.device("cuda")
+
+    # One shared-prefix rollout group: prompt + 3 completions of differing lengths.
+    torch.manual_seed(0)
+    prompt_ids = list(range(1, 49))  # P = 48
+    completion_lens = [16, 24, 8]
+    completions = [list(range(100 + 1000 * i, 100 + 1000 * i + n)) for i, n in enumerate(completion_lens)]
+    folded = fold_shared_prefix_rollouts(prompt_ids, completions)
+    spec, sample_token_ranges = AttnMaskSpec.prefix_tree(folded.node_lengths, folded.sample_paths)
+    prompt_len = len(prompt_ids)
+    total = len(folded.input_ids)
+
+    num_heads, head_dim = 8, 64
+    scale = head_dim**-0.5
+    q = torch.randn(1, num_heads, total, head_dim, device=device, dtype=torch.float32)
+    k = torch.randn(1, num_heads, total, head_dim, device=device, dtype=torch.float32)
+    v = torch.randn(1, num_heads, total, head_dim, device=device, dtype=torch.float32)
+
+    # --- under test: flex_attention with the spec realized as a block mask ---
+    block_mask = create_block_mask(_spec_mask_mod(spec), B=1, H=1, Q_LEN=total, KV_LEN=total, device=device)
+    out_flex = flex_attention(q, k, v, block_mask=block_mask, scale=scale)  # [1, H, T, D]
+
+    # --- reference: exact SDPA with the independently built dense mask ---
+    allowed = build_reference_mask(prompt_len, completion_lens, total).to(device)
+    out_ref = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed, scale=scale)
+
+    a = out_flex.float()
+    b = out_ref.float()
+    max_diff = (a - b).abs().max().item()
+    mean_diff = (a - b).abs().mean().item()
+    cos_sim = F.cosine_similarity(a.flatten(), b.flatten(), dim=0).item()
+    print(f"overall: max_diff={max_diff:.6e}, mean_diff={mean_diff:.6e}, cosine_sim={cos_sim:.8f}")
+
+    # Per-completion slice diffs (output laid out as [1, H, T, D] -> slice on T).
+    ok = True
+    for i, (start, end) in enumerate([rng[-1] for rng in sample_token_ranges]):
+        cd = (a[:, :, start:end] - b[:, :, start:end]).abs().max().item()
+        print(f"  completion[{i}] tokens [{start}:{end}] max_diff={cd:.6e}")
+        ok = ok and cd < MAX_DIFF_TOL
+
+    passed = ok and max_diff < MAX_DIFF_TOL and cos_sim > COS_SIM_TOL
+    print("PARITY PASSED" if passed else "PARITY FAILED")
+    return 0 if passed else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/tests/functional_tests/attention/prefix_tree_magi_parity.py b/tests/functional_tests/attention/prefix_tree_magi_parity.py
new file mode 100644
index 0000000000..1572ee2d57
--- /dev/null
+++ b/tests/functional_tests/attention/prefix_tree_magi_parity.py
@@ -0,0 +1,125 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Correctness check for the cp=1 prefix-tree mask on the magi attention backend.
+
+Runs the MagiAttention FFA kernel with a shared-prefix prefix-tree spec and
+compares its output, position by position, against an independently constructed
+dense-mask SDPA reference. The reference mask is built directly from the rollout
+structure (each completion attends FULL to the prompt and CAUSAL to itself), NOT
+from the spec, so a wrong spec or a wrong kernel application is caught.
+
+Requires one GPU and a source-built ``magi_attention`` (not on PyPI). Run with::
+
+    torchrun --nproc-per-node=1 tests/functional_tests/attention/prefix_tree_magi_parity.py
+
+Exit code 0 means parity holds within the bf16 tolerance.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from _prefix_tree_reference import build_reference_mask
+
+from nemo_automodel.components.datasets.llm.prefix_tree import fold_shared_prefix_rollouts
+from nemo_automodel.components.distributed.magi_attn_utils import (
+    AttnMaskSpec,
+    is_magi_available,
+    make_magi_attn_func,
+    set_active_attn_spec,
+    set_active_cp_group,
+)
+
+# bf16 vs fp32-reference tolerance (the FFA kernel runs in bf16).
+MAX_DIFF_TOL = 2e-2
+COS_SIM_TOL = 0.999
+
+
+def main() -> int:
+    if not torch.cuda.is_available():
+        print("SKIP: no CUDA device.")
+        return 0
+    if not is_magi_available():
+        print("SKIP: magi_attention not importable (source CUDA build required).")
+        return 0
+
+    if not dist.is_initialized():
+        os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+        os.environ.setdefault("MASTER_PORT", "29555")
+        os.environ.setdefault("RANK", "0")
+        os.environ.setdefault("WORLD_SIZE", "1")
+        dist.init_process_group(backend="nccl")
+    torch.cuda.set_device(0)
+    device = torch.device("cuda")
+
+    # cp=1: a size-1 group makes the magi attn_func take its cp=1 flex-key path
+    # (it short-circuits to SDPA only when the active cp group is None).
+    set_active_cp_group(dist.group.WORLD)
+
+    # One shared-prefix rollout group: prompt + 3 completions of differing lengths.
+    torch.manual_seed(0)
+    prompt_ids = list(range(1, 49))  # P = 48
+    completion_lens = [16, 24, 8]
+    completions = [list(range(100 + 1000 * i, 100 + 1000 * i + n)) for i, n in enumerate(completion_lens)]
+    folded = fold_shared_prefix_rollouts(prompt_ids, completions)
+    spec, sample_token_ranges = AttnMaskSpec.prefix_tree(folded.node_lengths, folded.sample_paths)
+    prompt_len = len(prompt_ids)
+    total = len(folded.input_ids)
+
+    num_heads, head_dim = 8, 64
+    scale = head_dim**-0.5
+    q = torch.randn(total, num_heads, head_dim, device=device, dtype=torch.bfloat16)
+    k = torch.randn(total, num_heads, head_dim, device=device, dtype=torch.bfloat16)
+    v = torch.randn(total, num_heads, head_dim, device=device, dtype=torch.bfloat16)
+
+    # --- magi (under test): FFA with the prefix-tree spec, THD [T, H, D] ---
+    attn_func = make_magi_attn_func(softmax_scale=scale)
+    set_active_attn_spec(spec)
+    out_magi = attn_func(q, k, v)  # [T, H, D]
+    set_active_attn_spec(None)
+
+    # --- reference: exact SDPA with the independently built dense mask, fp32 ---
+    allowed = build_reference_mask(prompt_len, completion_lens, total).to(device)
+    q4, k4, v4 = (t.float().transpose(0, 1).unsqueeze(0) for t in (q, k, v))  # [1, H, T, D]
+    out_ref = F.scaled_dot_product_attention(q4, k4, v4, attn_mask=allowed, scale=scale)
+    out_ref = out_ref.squeeze(0).transpose(0, 1)  # [T, H, D]
+
+    a = out_magi.float()
+    b = out_ref.float()
+    max_diff = (a - b).abs().max().item()
+    mean_diff = (a - b).abs().mean().item()
+    cos_sim = F.cosine_similarity(a.flatten(), b.flatten(), dim=0).item()
+    print(f"overall: max_diff={max_diff:.6e}, mean_diff={mean_diff:.6e}, cosine_sim={cos_sim:.8f}")
+
+    # Per-completion slice diffs, to localize any divergence to a specific branch.
+    ok = True
+    for i, (start, end) in enumerate([rng[-1] for rng in sample_token_ranges]):
+        cd = (a[start:end] - b[start:end]).abs().max().item()
+        print(f"  completion[{i}] tokens [{start}:{end}] max_diff={cd:.6e}")
+        ok = ok and cd < MAX_DIFF_TOL
+
+    passed = ok and max_diff < MAX_DIFF_TOL and cos_sim > COS_SIM_TOL
+    print("PARITY PASSED" if passed else "PARITY FAILED")
+    if dist.is_initialized():
+        dist.destroy_process_group()
+    return 0 if passed else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/tests/unit_tests/datasets/llm/test_prefix_tree.py b/tests/unit_tests/datasets/llm/test_prefix_tree.py
new file mode 100644
index 0000000000..37e2b4e2e7
--- /dev/null
+++ b/tests/unit_tests/datasets/llm/test_prefix_tree.py
@@ -0,0 +1,119 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo_automodel.components.datasets.llm.mock_prefix_tree import build_mock_rollout_dataset
+from nemo_automodel.components.datasets.llm.prefix_tree import (
+    CROSS_ENTROPY_IGNORE_IDX,
+    FoldedRollouts,
+    fold_shared_prefix_rollouts,
+    prefix_tree_collate_fn,
+)
+
+
+def test_fold_basic_layout_and_mask():
+    folded = fold_shared_prefix_rollouts([1, 2], [[10, 11], [20, 21, 22]])
+
+    # Flat dedup layout: prompt once, then each completion.
+    assert folded.input_ids == [1, 2, 10, 11, 20, 21, 22]
+    # Pre-shifted next-token labels: prompt masked, each completion predicts its
+    # own next token, and each completion's last token is masked (no successor).
+    _ = CROSS_ENTROPY_IGNORE_IDX
+    assert folded.labels == [_, _, 11, _, 21, 22, _]
+    # Each completion's positions continue from the prompt length (P=2).
+    assert folded.position_ids == [0, 1, 2, 3, 2, 3, 4]
+    # Tree structure: node 0 is the prompt, nodes 1..N the completions; each path
+    # is prompt -> completion. The magi backend turns this into the AttnMaskSpec.
+    assert folded.node_lengths == [2, 2, 3]
+    assert folded.sample_paths == [[0, 1], [0, 2]]
+
+
+def test_fold_single_completion():
+    folded = fold_shared_prefix_rollouts([5], [[7, 8]])
+    assert folded.input_ids == [5, 7, 8]
+    # Prompt masked; token 7 predicts 8; last token 8 masked.
+    assert folded.labels == [CROSS_ENTROPY_IGNORE_IDX, 8, CROSS_ENTROPY_IGNORE_IDX]
+    assert folded.position_ids == [0, 1, 2]
+    assert folded.node_lengths == [1, 2]
+    assert folded.sample_paths == [[0, 1]]
+
+
+def test_fold_empty_prompt_collapses_to_varlen_blocks():
+    folded = fold_shared_prefix_rollouts([], [[10, 11], [20]])
+    assert folded.input_ids == [10, 11, 20]
+    # No prompt: each completion still shifts within itself, last token masked.
+    assert folded.labels == [11, CROSS_ENTROPY_IGNORE_IDX, CROSS_ENTROPY_IGNORE_IDX]
+    # Each completion restarts at position 0.
+    assert folded.position_ids == [0, 1, 0]
+    # Empty prompt node 0; each completion is its own block (no shared ancestor).
+    assert folded.node_lengths == [0, 2, 1]
+    assert folded.sample_paths == [[1], [2]]
+
+
+def test_fold_custom_ignore_idx():
+    folded = fold_shared_prefix_rollouts([1], [[10, 11]], ignore_idx=-1)
+    # Custom ignore value used for the masked prompt and masked last token.
+    assert folded.labels == [-1, 11, -1]
+
+
+def test_fold_returns_folded_rollouts_instance():
+    folded = fold_shared_prefix_rollouts([1], [[2]])
+    assert isinstance(folded, FoldedRollouts)
+
+
+@pytest.mark.parametrize(
+    "prompt, completions, match",
+    [
+        ([1], [], "non-empty"),
+        ([1], [[10], []], "every completion must be non-empty"),
+    ],
+)
+def test_fold_validation_errors(prompt, completions, match):
+    with pytest.raises(ValueError, match=match):
+        fold_shared_prefix_rollouts(prompt, completions)
+
+
+def test_collate_builds_batched_tensors():
+    group = {"prompt_ids": [1, 2], "completions": [[10, 11], [20]]}
+    batch = prefix_tree_collate_fn([group])
+
+    assert set(batch) == {"input_ids", "labels", "position_ids", "prefix_tree"}
+    for key in ("input_ids", "labels", "position_ids"):
+        assert batch[key].shape == (1, 5)
+        assert batch[key].dtype == torch.long
+    assert torch.equal(batch["input_ids"], torch.tensor([[1, 2, 10, 11, 20]]))
+    assert torch.equal(batch["labels"], torch.tensor([[-100, -100, 11, -100, -100]]))
+    # Tree structure (node_lengths, sample_paths) for the magi backend to build the spec.
+    assert batch["prefix_tree"] == ([2, 2, 1], [[0, 1], [0, 2]])
+
+
+def test_collate_rejects_multi_group_batch():
+    group = {"prompt_ids": [1], "completions": [[2]]}
+    with pytest.raises(ValueError, match="local_batch_size=1 only"):
+        prefix_tree_collate_fn([group, group])
+
+
+def test_build_mock_rollout_dataset_shape_and_determinism():
+    ds = build_mock_rollout_dataset(num_groups=3, completions_per_group=4, prompt_len=8, completion_len=5)
+    assert len(ds) == 3
+    group = ds[0]
+    assert len(group["prompt_ids"]) == 8
+    assert len(group["completions"]) == 4
+    assert all(len(c) == 5 for c in group["completions"])
+    # Same seed -> identical data; it must be foldable by the collate.
+    again = build_mock_rollout_dataset(num_groups=3, completions_per_group=4, prompt_len=8, completion_len=5)
+    assert ds == again
+    assert prefix_tree_collate_fn([group])["input_ids"].shape == (1, 8 + 4 * 5)
diff --git a/tests/unit_tests/distributed/test_magi_attn_utils.py b/tests/unit_tests/distributed/test_magi_attn_utils.py
index 99c619d686..c6842d0d0a 100644
--- a/tests/unit_tests/distributed/test_magi_attn_utils.py
+++ b/tests/unit_tests/distributed/test_magi_attn_utils.py
@@ -141,6 +141,28 @@ def test_defaults_disabled(self):
     def test_hf_dispatch(self, enabled, custom, expected):
         assert MagiState(enabled=enabled, custom=custom).hf_dispatch is expected
 
+    def test_prepare_llm_batch_hf_prefix_tree_raises(self):
+        # HF magi backend cannot receive the out-of-band prefix-tree mask spec, so
+        # the recipe must fail loudly rather than silently drop it.
+        st = MagiState(enabled=True, custom=False, cp_group=None, cp_size=1)
+        batch = {"input_ids": torch.zeros(1, 4, dtype=torch.long), "prefix_tree": ([1, 1, 1], [[0, 1], [0, 2]])}
+        with pytest.raises(NotImplementedError, match="prefix-tree attention mask is only supported"):
+            st.prepare_llm_batch(model=None, batch=batch, device_mesh=None, is_thd=False, pad_id=0, num_chunks=1)
+
+    def test_prepare_llm_batch_custom_prefix_tree_ok(self):
+        # Custom backend (cp=1, no THD packing) accepts the prefix-tree mask: it
+        # activates the spec out-of-band and returns the batch with the key popped.
+        st = MagiState(enabled=True, custom=True, cp_group=None, cp_size=1)
+        batch = {"input_ids": torch.zeros(1, 4, dtype=torch.long), "prefix_tree": ([1, 1, 1], [[0, 1], [0, 2]])}
+        try:
+            _, out = st.prepare_llm_batch(
+                model=None, batch=batch, device_mesh=None, is_thd=False, pad_id=0, num_chunks=1
+            )
+            assert "prefix_tree" not in out
+            assert mu.get_active_attn_spec() is not None
+        finally:
+            mu.set_active_attn_spec(None)
+
 
 # --------------------------------------------------------------------------- #
 # setup_magi