> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.speculative.target_cp

Target-side context parallelism for speculative-decoding draft training.

The draft trainers (EAGLE-3, DFlash, DSpark) shard the FROZEN target's forward
along the sequence and gather its captured hidden states / logits back to the
full sequence before handing them to the draft. That flow is specific to
speculative decoding -- it needs no gradients through the target and no
model-owned CP sharder -- so it lives here rather than in the shared
`components/distributed/cp_utils` surface.

## Module Contents

### Functions

| Name                                                                                                                    | Description                                                           |
| ----------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------- |
| [`attach_cp_kv_gather_hooks`](#nemo_automodel-components-speculative-target_cp-attach_cp_kv_gather_hooks)               | Context-parallel self-attention for a FROZEN (forward-only) target.   |
| [`gather_cp_seq`](#nemo_automodel-components-speculative-target_cp-gather_cp_seq)                                       | Gather context-parallel sharded `tensors` back to the full sequence.  |
| [`make_target_cp_ctx`](#nemo_automodel-components-speculative-target_cp-make_target_cp_ctx)                             | Build a context-parallel context for a frozen target forward.         |
| [`run_target_cp_forward_and_gather`](#nemo_automodel-components-speculative-target_cp-run_target_cp_forward_and_gather) | Run a frozen target under context parallelism and gather its outputs. |

### API

```python
nemo_automodel.components.speculative.target_cp.attach_cp_kv_gather_hooks(
    model: torch.nn.Module,
    cp_mesh
) -> None
```

Context-parallel self-attention for a FROZEN (forward-only) target.

Torch's `context_parallel` ring dispatch does not fire for a plain HuggingFace
forward -- q/k/v reach `F.scaled_dot_product_attention` as ordinary local
tensors, so each rank silently attends only to its own sequence shard and every
position past the first shard is wrong. Because the target is frozen (no
backward through attention), the correct fix is simple: all-gather K/V across the
cp group and attend the local Q against the full K/V with a global causal mask.
The O(S^2) attention matrix stays sharded `[S/cp, S]` per rank (the memory
win); only the O(S) K/V is replicated.

Assumes contiguous (non-load-balanced) sharding -- rank `r` holds global
positions `[r*S_local, (r+1)*S_local)` -- which is what :func:`make_target_cp_ctx`
produces (it disables load balancing). Q/K/V are `[B, nH, S_local, D]` at the
SDPA call, so the sequence dim is 2.

```python
nemo_automodel.components.speculative.target_cp.gather_cp_seq(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh,
    tensors: typing.List[torch.Tensor],
    seq_dim: int,
    orig_len: int
)
```

Gather context-parallel sharded `tensors` back to the full sequence.

Inverse of the sharding done by :func:`make_target_cp_ctx`. Uses torch's
`context_parallel_unshard` with `load_balancer=None` (matching the
load-balancing-disabled sharding) and slices the right-pad back off.

**Parameters:**

The context-parallel device (sub)mesh used to shard.

Local-shard tensors (e.g. captured aux hidden states, logits),
each sharded to `T/cp` along `seq_dim`.

The sequence dimension to gather along.

The pre-pad sequence length to slice back to.

**Returns:**

A list of full-sequence tensors of length `orig_len` along `seq_dim`.

```python
nemo_automodel.components.speculative.target_cp.make_target_cp_ctx(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh,
    input_ids,
    position_ids = None
)
```

Build a context-parallel context for a frozen target forward.

Shards `input_ids` (and `position_ids`) along the sequence dim across
`cp_mesh` so the target's self-attention runs as ring attention. Unlike
:func:`make_cp_batch_and_ctx`, this does not require `labels` and is meant
for the EAGLE-3 target wrapper, which gathers the aux/logits back to the full
sequence (see :func:`gather_cp_seq`) before handing them to the draft.

Load balancing is disabled (`_cp_options.enable_load_balance = False`) so
each rank holds a contiguous sequence chunk and the gather is a plain ordered
concat (no round-robin un-permute). The sharding is thrown away right after
the forward, so load balancing buys nothing here, and the ordered shard makes
the gather deterministic. This is a process-global torch flag; the EAGLE-3
recipe is the only context-parallel user in its process.

The sequence is right-padded to a multiple of `cp_size`; the returned
`orig_len` lets the caller slice the gathered outputs back down.

**Parameters:**

The context-parallel device (sub)mesh.

`[B, T]` token ids.

Optional `[B, T]` (or `[1, T]`) position ids; an arange
is injected when omitted.

**Returns:**

`(cp_ctx, sharded_input_ids, sharded_position_ids, orig_len)`. Enter

```python
nemo_automodel.components.speculative.target_cp.run_target_cp_forward_and_gather(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh,
    model: torch.nn.Module,
    input_ids: torch.Tensor,
    forward_kwargs: dict,
    collect: typing.Callable[[object], typing.List[torch.Tensor]],
    position_ids: typing.Optional[torch.Tensor] = None,
    filter_kwargs: bool = False
) -> tuple
```

Run a frozen target under context parallelism and gather its outputs.

Shards `input_ids` (and `position_ids`) along the sequence via
:func:`make_target_cp_ctx`, runs `model` as ring attention with
`attention_mask=None` (the `self_attn` hooks force `is_causal`), and --
still inside the CP context, before any capture hooks are removed -- gathers
the tensors returned by `collect(outputs)` back to the full sequence with
:func:`gather_cp_seq` (`seq_dim=1`, un-padded to `orig_len`).

Centralizes the gather-inside-context invariant and the `seq_dim`/`orig_len`
contract shared by the eagle3/dflash/dspark target wrappers, so a fix lands in
one place instead of drifting across three copies.

**Parameters:**

The `cp` device submesh.

The frozen target module.

Full (unsharded) `[B, T]` token ids.

Extra kwargs for `model.forward`. Must not include
`input_ids` / `attention_mask` / `position_ids` -- those are
injected here (CP forces `attention_mask=None`).

`callable(outputs) -> list[Tensor]` selecting the tensors to
gather; invoked inside the CP context after the forward, so it also
sees any tensors captured by forward hooks.

Optional `[B, T]` / `[1, T]` positions; an arange is
injected by :func:`make_target_cp_ctx` when omitted.

Drop kwargs the model's forward does not accept (via
:func:`filter_forward_kwargs`) -- needed for the VLM/MoE targets.

**Returns:** `tuple`

`(outputs, gathered)` -- the raw model outputs and the list of