No <eos> token in the training? Clarification: Cross-Sample Information Flow in PrefixLM Attention (Pass 2)

[MugeTong]

Hello,

Thank you very much for sharing so exciting and awsome work!

Description

I've been studying the implementation of the two-pass PrefixLM attention mechanism in flash_attention_prefixlm_v2.py, and I have a question about the design choice regarding cross-sample information flow during Pass 2 (causal attention).

Observation

In the forward pass implementation, there are two separate attention computations:

Pass 1 (Bidirectional):

• Operates only on instruction/prefix portions
• Each sample's instruction is masked to attend only to itself
• seqused_q=prefix_lens, seqused_k=prefix_lens

Pass 2 (Causal):

• Operates on response/causal portions
• Uses complete cu_seqlens for key/value, allowing cross-sample access
• cu_seqlens_q=cu_seqlens_shifted, cu_seqlens_k=cu_seqlens

The Question

Based on the code structure, in a batch containing multiple (instruction, response) pairs concatenated without separators:

[inst1, resp1, inst2, resp2, inst3, resp3, ...]

During Pass 2 (causal attention):

• resp1 can causally attend to [inst1, resp1_prefix] ✓
• resp2 can causally attend to [inst1, resp1, inst2, resp2_prefix] ← cross-sample
• resp3 can causally attend to [inst1, resp1, inst2, resp2, inst3, resp3_prefix] ← cross-sample

This is because:

1. cu_seqlens_k remains the complete sequence boundaries (all samples)
2. causal=True enforces causal masking within the response portion
3. But there's no explicit sample-level boundary masking to prevent cross-sample attention

So does this means there is cross-sample attention in the training?