# Vision Encoder Parallelism

<!-- meta: description: Vision encoder parallelism design -->
<!-- meta: content_type: conceptual-deep-dive -->
<!-- meta: date_updated: 2026-06-23 -->

## Overview

The vision encoder supports independent TP and DP configurations, decoupled from the text model's parallelism. Vision TP shards encoder weights across ranks; Vision DP scatters vision blocks across ranks for independent processing, then all-gathers after the merger.

**Prefill-only.** Vision encoding runs during prefill. Decode steps do not invoke the vision encoder.

**Block-level DP.** Unlike text model DP which shards by batch, vision DP shards by the `num_blocks` dimension of the block-packed vision inputs, where each block contains at least one image. Each DP rank processes `num_blocks / dp_size` blocks independently, then results are gathered.

## Configuration

```python
VisionNeuronConfig(
    tp_size=1,    # Vision TP degree (default: 1, no weight sharding)
    dp_size=4,    # Vision DP degree (default: inferred as world_size // tp_size)
)
```

| Parameter | What it controls | Default |
|-----------|-----------------|---------|
| `tp_size` | Weight sharding across TP ranks | 1 |
| `dp_size` | Block scatter/gather across DP ranks | Inferred: `world_size // tp_size` |

**Constraints:**

- `tp_size` must be a positive integer that divides `world_size`
- `tp_size * dp_size == world_size` (all ranks must participate — no redundant compute)
- `num_blocks` is padded to be divisible by `dp_size` for even scatter

**Resolution rules** (`resolve_tp_dp`):

- Only `tp_size` set → `dp_size` inferred as `world_size // tp_size`
- Neither set (defaults `tp_size=1, dp_size=1`) → inferred to `tp_size=1, dp_size=world_size` (full DP)
- Both set explicitly → validated that `tp_size * dp_size == world_size`

`resolve_tp_dp` is called both at worker init (for process group creation) and at model runner init (for warmup and runtime block padding). It writes back the resolved `dp_size` to the config instance so downstream callers see the correct value.

## Modes

| User config | Resolved (world_size=16) | Behavior |
|-------------|--------------------------|----------|
| (default) | tp=1, dp=16 | Full DP — each rank holds full weights, processes `num_blocks / 16` blocks |
| `tp_size=4` | tp=4, dp=4 | Combined — weights sharded across 4 TP ranks, blocks scattered across 4 DP ranks |
| `tp_size=16` | tp=16, dp=1 | Full TP — weights sharded across 16 ranks, all ranks process all blocks |
| `tp_size=4, dp_size=4` | tp=4, dp=4 | Same as specifying only `tp_size=4` (explicit dp matches inferred) |

## Process Groups

### Vision TP group

Created when `vision_tp_size != text_tp_size`. Each group contains `vision_tp_size` contiguous ranks.

```text
Example: world_size=16, vision_tp=4

  Vision TP groups (4 groups of 4):
    [0,1,2,3]  [4,5,6,7]  [8,9,10,11]  [12,13,14,15]
```

When `vision_tp_size == text_tp_size`, the text TP group is reused (no new group created).

### Vision DP group

Created when `vision_dp_size > 1`. Column groups built via `_build_dp_column_group_ranks(world_size, vision_tp_size, vision_dp_size)` — ranks at the same TP position across DP replicas.

```text
Example: world_size=16, vision_tp=4, vision_dp=4

  Vision TP groups:  [0,1,2,3]  [4,5,6,7]  [8,9,10,11]  [12,13,14,15]
                      │ │ │ │    │ │ │ │    │ │ │ │       │ │ │ │
  Vision DP columns: [0,4,8,12] [1,5,9,13] [2,6,10,14]  [3,7,11,15]
```

```text
Example: world_size=8, vision_tp=1, vision_dp=4 (full DP, no weight sharding)

  Vision TP groups:  [0] [1] [2] [3] [4] [5] [6] [7]   (each rank = own TP group)
  Vision DP columns: [0,1,2,3]  [4,5,6,7]                (4 ranks per DP group)
```

### Constraint: tp_size * dp_size == world_size

Vision TP and DP must fully utilize all ranks — no redundant compute. At runtime, `resolve_tp_dp` enforces `tp_size * dp_size == world_size`. If the user specifies only `tp_size`, `dp_size` is inferred as `world_size // tp_size`. If neither is specified, defaults to full DP (`tp_size=1, dp_size=world_size`).

## Forward Flow

```text
Input: pixel_values [num_blocks, block_size, patch_dim]
       (+ pos_emb_idx, pos_emb_weight, cos, sin, bound_min, bound_max, unpack_indices)

Step 0 — DP scatter (if dp_size > 1):
  blocks_per_rank = num_blocks // dp_size
  Each rank slices its portion: [start:end] along dim 0
  pixel_values: [num_blocks, ...] → [blocks_per_rank, ...]

Step 1 — Patch embedding + position embedding:
  pixel_values → hidden_states [blocks_per_rank, block_size, hidden]

Step 2 — Transformer blocks + deepstack extraction:
  For each layer: attention + MLP (TP collectives within vision TP group)
  At deepstack layers: extract merged features

Step 3 — Final merger:
  hidden_states → main_merged [blocks_per_rank, merged_bs, D]

Step 4 — DP gather (if dp_size > 1):
  all_gather(main_merged, dim=0) across vision DP group
  main_merged: [blocks_per_rank, ...] → [num_blocks, merged_bs, D]
  Same for deepstack features.

Step 5 — Unpack to per-image order:
  index_select with unpack_indices → [merged_seq_len, D]
  (Operates on full gathered tensor — same on all ranks)

Step 6 — Construct fat tensor (main + deepstack):
  Concatenate main and deepstack features → final output
```

## Block Padding

`num_blocks` is padded to be divisible by `dp_size` at three points:

1. **Config `__post_init__`**: `num_total_vision_attention_blocks` padded for compiled graph shapes
2. **`select_vision_bucket`**: runtime bucket selection pads `num_blocks`
3. **Vision warmup** (`neuron_worker.py`): warmup shapes padded

Padding blocks contain zeros and produce zero outputs. The `unpack_indices` tensor ensures only valid merged tokens are selected in the final output.

## Weight Loading

Vision DP does not affect weight loading — each DP rank holds a **full copy** of the vision encoder weights (dp_size > 1 reduces activation memory, not weight memory).

Vision TP shards weights across the vision TP group using the same column/row parallel patterns as the text model, but scoped to the vision TP group.

## When to Use Block-Level DP

- **Large images / many images per request**: High `num_blocks` makes the per-rank activation memory the bottleneck. DP reduces activation memory linearly with `dp_size`.
- **Small vision encoder**: When the encoder fits in a single device's weight memory, TP provides no benefit. Full DP (`tp_size=1`) avoids TP communication overhead entirely.
- **Combined TP+DP**: For very large encoders that don't fit on one device AND have high block counts, use both.

## Block-Level DP vs. Model-Level DP

The current implementation is **block-level DP**: blocks from a single `embed_multimodal` call are scattered across ranks within one compiled graph. All ranks participate in every encoder invocation and the result is all-gathered before returning.

**Model-level DP** is a different approach: multiple independent encoder replicas, each handling separate requests (or subsets of images) concurrently with no communication between replicas.

| Aspect | Block-level DP (current) | Model-level DP (future) |
|--------|--------------------------|-------------------------|
| Scope | Within one `embed_multimodal` call | Across concurrent calls |
| What's distributed | Blocks from one batched image set | Entire requests to independent replicas |
| Communication | All-gather after compute | None between replicas |
| Compiled bucket | One bucket per call (all images share it) | Each replica selects its own bucket |
| Latency | Reduces single-call latency | Does not reduce single-call latency |
| Throughput | Increases via cross-request batching | Increases via parallel independent calls |

### Cross-request batching

The scheduler already batches images from multiple requests into a single `embed_multimodal` call (via `group_and_batch_mm_kwargs`). This means block-level DP naturally benefits from concurrent request load — more pending images → more blocks → better DP utilization.

### Where model-level DP adds unique value

Block-level DP has a fundamental constraint: all images in one call must use the **same compiled bucket**. When a small image (256 tokens → bucket 1024) and a large image (4096 tokens → bucket 4096) are batched together, the small image pays the latency cost of the larger bucket.

Model-level DP solves this by routing images to replicas with different optimal buckets:

```text
Block-level DP (current):
  Small + Large images → both use bucket 4096 → wasted compute on small image

Model-level DP (future):
  Replica 0: Small image → bucket 1024 (fast)
  Replica 1: Large image → bucket 4096 (parallel)
```