# Lesson 001: CPU Flash Attention Silently Drops attn_mask

**Date:** 2026-01-28
**Severity:** Critical (produces visual artifacts)
**Files:** `src/neurobrix/core/runtime/graph/compiled_ops.py`

## Symptom

PixArt-Alpha produces mosaic/blocky artifacts. PixArt-Sigma (traced on GPU) works fine.

## Root Cause

When the model is traced on CPU, PyTorch selects `_scaled_dot_product_flash_attention_for_cpu` as the attention backend. This op does NOT natively accept `attn_mask`. However, PixArt-Alpha's cross-attention passes the encoder attention mask as a bias tensor via kwargs:

```json
{
  "op_type": "aten::_scaled_dot_product_flash_attention_for_cpu",
  "input_tensor_ids": ["q", "k", "v"],
  "kwargs": {"attn_mask": {"type": "tensor", "tensor_id": "..."}}
}
```

The compiled handler extracted `scale` from kwargs but ignored `attn_mask`:

```python
# BROKEN
def flash_cpu_attention(q, k, v, dropout_p=0.0, is_causal=False, **kwargs):
    scale = kwargs.get("scale", None)
    output = F.scaled_dot_product_attention(q, k, v, scale=scale, ...)
    # attn_mask silently dropped!
```

Without the mask, all 120 token positions (including ~115 padding tokens) are treated as valid. The padding embeddings pollute cross-attention output across all 28 transformer blocks.

## Fix

```python
def flash_cpu_attention(q, k, v, dropout_p=0.0, is_causal=False, **kwargs):
    scale = kwargs.get("scale", None)
    attn_mask = kwargs.get("attn_mask", None)
    output = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, scale=scale, ...)
```

## Why Sigma Was Unaffected

PixArt-Sigma was traced on GPU, producing `_scaled_dot_product_efficient_attention` ops. The handler for that variant already accepted `attn_bias` as a positional argument (input[3]), so the mask was correctly forwarded.

## Debugging Timeline

| Step | Finding | Was it the cause? |
|------|---------|-------------------|
| pos_embed buffer loading | Fixed None return, buffer identical to diffusers | No |
| Weight comparison | All 612 weights match diffusers (diff=0.0) | No |
| VAE comparison | Identical output (diff=0.0) | No |
| Scheduler comparison | Identical output (diff=0.0) | No |
| Text encoder comparison | Identical output (diff=0.0) | No |
| Transformer with random inputs | mean_diff=0.25, cosine=0.87 | Misleading (different backends) |
| Topology/sequence length | Correct (120 everywhere) | No |
| Trace mask chain in graph | Mask reaches cross-attn as kwargs.attn_mask | **Yes** |
| compiled_ops handler | attn_mask extracted but not forwarded | **Root cause** |

## Lesson

When redirecting device-specific attention variants to `F.scaled_dot_product_attention`, forward ALL kwargs — not just the ones you expect. The graph may pass additional tensors (like `attn_mask`) that the original op variant doesn't natively support but that were captured during tracing because a higher-level module passed them.

**General rule:** Any `**kwargs` sink in a dispatch handler should be audited for silently dropped tensors.
