Cu Lengths Attention Flow

Purpose / 用途

Use this skill when reasoning about attention boundaries in the LLaVA-OneVision2 forward pass — specifically how cu_lengths and patch_positions control attention at different stages of the model.

在分析 LLaVA-OneVision2 前向传播中的 attention 边界时使用这个 skill——具体来说，cu_lengths 和 patch_positions 如何在模型的不同阶段控制 attention。

This skill is specifically for:

• Understanding the difference between ViT-level and LLM-level attention control
• Debugging packed vs non-packed attention behavior
• Reasoning about cross-sample isolation in packed sequences
• Understanding why patch_positions grouping does NOT carry into the LLM

这个 skill 专门用于：

• 理解 ViT 层和 LLM 层 attention 控制的区别
• 调试 packed 和 non-packed 的 attention 行为
• 分析 packed 序列中跨样本隔离机制
• 理解为什么 patch_positions 的分组不会延续到 LLM 中

Key Files / 关键文件

Core Concept: Two Independent Attention Control Mechanisms / 核心概念：两套独立的 Attention 控制机制

Overview Diagram / 概览图

┌─────────────────────────────────────────────────┐
│  ViT Encoder                                    │
│                                                 │
│  Control: patch_positions (temporal dimension)  │
│  Effect:  Local/shared attention                │
│           e.g. 4 images share one attention     │
│           window via same temporal index        │
│                                                 │
│  Output:  visual embeddings                     │
└──────────────────┬──────────────────────────────┘
                   │  (embeddings replace image
                   │   placeholder tokens)
                   ▼
┌─────────────────────────────────────────────────┐
│  LLM Decoder                                    │
│                                                 │
│  Control: cu_lengths (cumulative sub-seq lens)  │
│  Effect:  Determines attention domain           │
│                                                 │
│  NON-PACKED: cu_lengths == [[0]]                │
│    → full causal attention                      │
│    → ALL tokens see ALL previous tokens         │
│    → patch_positions grouping is GONE           │
│                                                 │
│  PACKED: cu_lengths = [0, a, a+b, ...]          │
│    → block-diagonal causal attention            │
│    → sub-sequences isolated from each other     │
│    → within each sub-seq: full causal           │
└─────────────────────────────────────────────────┘

Mechanism Details / 机制详解

1. cu_lengths Generation / cu_lengths 的产生

Source A: Offline Packed Data (PackedCaptioningSample)

# task_encoder.py → pack_selected_samples()
cu_lengths = [0]
for sample in samples:
    current_length += sample.total_len
    cu_lengths.append(current_length)
# Result: [0, 512, 1024, 1389] — 3 sub-samples packed together

Each sub-sample was independently encoded (tokenized + image processed) then concatenated into one long sequence. cu_lengths records the boundaries.

每个子样本独立编码（tokenize + 图像处理），然后拼接成一个长序列。cu_lengths 记录边界。

Source B: Non-Packed Single Samples

# task_encoder.py → batch()
if self.is_packing_enabled or int(os.environ.get("OFFLINE_PACKED_DATA", 0)) == 1:
    cu_lengths = torch.stack([s.cu_lengths for s in samples])
else:
    cu_lengths = torch.tensor([[0]], dtype=torch.int32)  # dummy value

Non-packed samples get cu_lengths = [[0]] with shape [1, 1].

非 packed 样本得到 cu_lengths = [[0]]，shape 为 [1, 1]。

2. Forward Function Branching / 前向函数分支

In pretrain_llava_onevision2.py:

if cu_lengths.shape == torch.Size([1, 1]):
    # ===== NON-PACKED PATH =====
    # Uses attn_mask for padding_causal attention
    # Every token attends to all previous tokens (full causal)
    # packed_seq_params = None
    for i in range(attn_mask.shape[0]):
        loss_mask[i, (attn_mask[i] == False).sum() - 1] = 0
else:
    # ===== PACKED PATH =====
    # micro-batch-size must be 1 for packing
    assert cu_lengths.shape[0] == 1
    attn_mask = None  # not needed — cu_seqlens defines boundaries
    packed_seq_params = PackedSeqParams(
        qkv_format="thd",
        cu_seqlens_q=cu_lengths[0],      # → Flash Attention kernel
        cu_seqlens_kv=cu_lengths[0],
        max_seqlen_q=max_lengths[0].item(),
        max_seqlen_kv=max_lengths[0].item(),
    )

3. LLM Attention Behavior / LLM 层的 Attention 行为

Non-packed (cu_lengths == [[0]]) → Full Causal Attention：

• The entire sequence is ONE attention domain
• Token at position i can attend to positions 0..i
• ALL visual tokens (from any image) + ALL text tokens are mutually visible
• Compute cost: O(n²) where n = total sequence length
• packed_seq_params = None → standard causal mask

整个序列是一个 attention 域，所有 visual token（来自任何图片）和所有 text token 互通可见。计算量 O(n²)。

Packed (cu_lengths = [0, a, a+b, ...]) → Block-Diagonal Causal Attention：

• Each sub-sequence [cu_lengths[i], cu_lengths[i+1]) is an independent attention domain
• Within each sub-sequence: full causal attention
• Between sub-sequences: ZERO attention (completely isolated)
• Implemented via Flash Attention's cu_seqlens parameter
• Compute cost: O(a² + b² + c² + ...) << O((a+b+c+...)²)

每个子序列 [cu_lengths[i], cu_lengths[i+1]) 是独立的 attention 域，子序列之间完全隔离。通过 Flash Attention 的 cu_seqlens 参数实现。

4. Sequence Parallelism Padding / 序列并行填充

When args.sequence_parallel is enabled, the sequence must be divisible by TP size (and TP×CP×2 if CP > 1). For packed sequences, padding tokens are appended as a dummy extra sub-sequence:

当启用 sequence_parallel 时，序列长度必须被 TP size 整除。对 packed 序列，padding token 作为一个额外的 dummy 子序列追加：

if packed_seq_params is not None:
    new_end_q = packed_seq_params.cu_seqlens_q[-1:] + pad_size
    packed_seq_params = PackedSeqParams(
        cu_seqlens_q=torch.cat([packed_seq_params.cu_seqlens_q, new_end_q]),
        cu_seqlens_kv=torch.cat([packed_seq_params.cu_seqlens_kv, new_end_kv]),
        max_seqlen_q=max(packed_seq_params.max_seqlen_q, pad_size),
        max_seqlen_kv=max(packed_seq_params.max_seqlen_kv, pad_size),
    )

Critical Insight: patch_positions Scope / 关键洞察：patch_positions 的作用域

ViT Layer: patch_positions Controls Attention Grouping

In the ViT encoder, patch_positions has a temporal dimension (t, h, w). Images sharing the same temporal index share one attention window. For example, 4 images treated as "video frames" share attention via their temporal coordinates.

在 ViT encoder 中，patch_positions 有时间维度 (t, h, w)。共享相同时间索引的图片共享一个 attention window。例如，4 张图片被当作"视频帧"通过时间坐标共享 attention。

LLM Layer: patch_positions Has NO Effect on Attention

patch_positions is NOT used to control LLM attention. It is only passed through the data pipeline for potential use in position embeddings or other purposes, but the LLM's attention boundaries are controlled EXCLUSIVELY by cu_lengths.

patch_positions 不控制 LLM 的 attention。 它只是在数据 pipeline 中传递，可能用于 position embedding 等目的，但 LLM 的 attention 边界完全由 cu_lengths 控制。

This means:

这意味着：