"""End-to-end integration tests for the composed HRDiT pipeline (plan phase P9). - T9.1 ``test_e2e_spa_plus_hap_2k`` — SPA + HAP composed over 8 simulated denoising steps with step gating (``spa_steps=3``): the leading steps run the full averaged-variant passes through the HAP kernel; the later steps run a single kernel pass each; the ORIGINAL dense attention is never called while HAP is live. - T9.2 ``test_e2e_hap_standalone`` — SPA off (``bundle_size=1``), HAP on: every step is a single kernel pass whose output equals the dense-mask reference, and differs from the no-HAP baseline. - T5.1 ``TestE2EAnimaRegression`` — Anima-style cosmos-convention calls (``skip_reshape=True`` + ``transformer_options`` kwargs) through the REAL shared unet wrapper: matching plan engages the kernel (output == dense-mask reference, orig never called); mismatched plan falls back gracefully with a one-time warning (the exact scenario that crashed pre-fix). - T5.2 ``test_e2e_cross_attention_mix`` — a forward mixing square self-attn and non-square cross-attn calls: self-attn HAP-masked, cross-attn plain, layer counter advances for BOTH (alignment preserved). These drive the REAL shared unet wrapper (``_hrdit_install_hook``) with a decreasing-sigma schedule so the step-count gate behaves exactly as in a real ComfyUI forward. All math is on the mock SDPA backend (no CUDA required). Markers: @pytest.mark.mock_integration Accept (user-run): pytest tests/test_hap_integration.py -k e2e_spa pytest tests/test_hap_integration.py -k e2e_hap_standalone pytest tests/test_hap_integration.py -k anima pytest tests/test_hap_integration.py -k cross_attention """ import pytest import torch import torch.nn.functional as F from src import hap from src.spa import _hrdit_install_hook from src.spa_context import ( SPAContext, set_hap_context, set_hrdit_layer_idx, set_hrdit_proportional, set_spa_context, set_spa_layer_filter, set_spa_step_gate, ) NUM_LAYERS = 4 HEADS = 2 DIM = 16 SEQ_LEN = 64 def _rand_qkv(seed=0): g = torch.Generator().manual_seed(seed) q = torch.randn(1, HEADS, SEQ_LEN, DIM, generator=g) k = torch.randn(1, HEADS, SEQ_LEN, DIM, generator=g) v = torch.randn(1, HEADS, SEQ_LEN, DIM, generator=g) return q, k, v def _identity_spa_ctx(num_variants=5): """SPA context with IDENTITY variant rotations (s = (n+1)//2 = 3 -> 5 passes). Identity rotations make every variant pass mathematically plain attention, so the test isolates the PASS COUNT / kernel-routing behaviour without depending on RoPE numerics. """ L, P = SEQ_LEN, DIM // 2 eye = torch.eye(2).expand(1, 1, L, P, 2, 2).clone() return SPAContext( active=True, bundle_size=(num_variants + 1) // 2, base_pe=eye.clone(), variant_pes=[eye.clone() for _ in range(num_variants)], variant_deltas=[eye.clone() for _ in range(num_variants)], pre_roped=True, fmt="flux", text_len=0, ) def _plan(num_layers=NUM_LAYERS, num_heads=HEADS, alpha=64.0, beta=0.0): return hap.ScopePlan( alphas=[[alpha] * num_heads for _ in range(num_layers)], betas=[[beta] * num_heads for _ in range(num_layers)], ) def _hap_ctx(num_layers=NUM_LAYERS, text_len=0, backend="dense"): return hap.HapContext(active=True, plan=_plan(num_layers), text_len=text_len, backend=backend) class _MockModel: def __init__(self): self._unet_wrapper = None self._spa_installed = None self._spa_orig_optimized_attention = None self._hrdit_consumers = None self._hap_ctx = None # Step-gating state (normally set by apply_spa_to_model). self._spa_steps = 0 self._spa_start_sigma = 1.0 self._spa_step_counter = 0 self._spa_last_sigma = None self._spa_layer_filter = None self._hrdit_proportional_attention = False def set_model_unet_function_wrapper(self, fn): self._unet_wrapper = fn @pytest.fixture def mock_attn(): import comfy.ldm.modules.attention as attn_mod return attn_mod @pytest.fixture(autouse=True) def _clean_state(): hap.HapRuntime.reset() set_hrdit_layer_idx(0) set_spa_layer_filter(None) set_hrdit_proportional(False) yield set_hap_context(None) set_spa_context(None) set_spa_step_gate(True) set_spa_layer_filter(None) set_hrdit_layer_idx(0) set_hrdit_proportional(False) hap.HapRuntime.reset() def _install_spy_orig(mock_attn): """Wrap the pristine SDPA so we can assert it is NEVER called while HAP is live (the wrapper's ``orig`` fallback must not fire).""" pristine = mock_attn.optimized_attention orig_calls = [] def spy_orig(q, k, v, heads, *a, **kw): orig_calls.append(1) return pristine(q, k, v, heads, *a, **kw) mock_attn.optimized_attention = spy_orig return orig_calls # --------------------------------------------------------------------------- # T9.1 — SPA + HAP composed over a gated 8-step schedule # --------------------------------------------------------------------------- @pytest.mark.mock_integration class TestE2ESpaPlusHap: def test_e2e_spa_plus_hap_2k(self, mock_attn): orig_calls = _install_spy_orig(mock_attn) m = _MockModel() m._spa_steps = 3 # SPA on the 3 LEADING steps only m._spa_start_sigma = 1.0 # no sigma-threshold gating m._hap_ctx = _hap_ctx() # HAP live on EVERY step _hrdit_install_hook(m, "flux", consumer="spa") _hrdit_install_hook(m, "flux", consumer="hap") # shared wrapper assert m._hrdit_consumers == {"spa", "hap"} q, k, v = _rand_qkv() kernel_calls = [] real_attn = hap.HapRuntime.attn def spy(self, qq, kk, vv, layer, **kw): kernel_calls.append(layer) return real_attn(self, qq, kk, vv, layer, **kw) def model_fn(x, t, **c): # Simulate the embedder forward: register the SPA variants, then run # one attention call per layer through the patched symbol. set_spa_context(_identity_spa_ctx(num_variants=5)) for _ in range(NUM_LAYERS): mock_attn.optimized_attention(q, k, v, HEADS) return x # Decreasing sigma -> no jump-up -> the step counter runs 0..7. sigmas = [1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3] per_step_kernel_counts = [] hap.HapRuntime.attn = spy try: for sigma in sigmas: kernel_calls.clear() m._unet_wrapper(model_fn, { "input": torch.zeros(1), "timestep": torch.tensor(sigma), "c": {}, }) per_step_kernel_counts.append(len(kernel_calls)) finally: hap.HapRuntime.attn = real_attn # Leading 3 steps: SPA active -> 5 variant passes x 4 layers = 20 each. # Later 5 steps: SPA gated off -> 1 pass x 4 layers = 4 each. assert per_step_kernel_counts == [20, 20, 20, 4, 4, 4, 4, 4] # The original dense attention never fires while HAP is live. assert orig_calls == [] def test_e2e_spa_plus_hap_kernel_layer_sequence(self, mock_attn): """Within one SPA-active step the kernel sees each layer index repeated per variant pass (5x), in block order.""" _install_spy_orig(mock_attn) m = _MockModel() m._spa_steps = 1 m._hap_ctx = _hap_ctx() _hrdit_install_hook(m, "flux", consumer="spa") _hrdit_install_hook(m, "flux", consumer="hap") q, k, v = _rand_qkv() kernel_calls = [] real_attn = hap.HapRuntime.attn def spy(self, qq, kk, vv, layer, **kw): kernel_calls.append(layer) return real_attn(self, qq, kk, vv, layer, **kw) def model_fn(x, t, **c): set_spa_context(_identity_spa_ctx(num_variants=5)) for _ in range(NUM_LAYERS): mock_attn.optimized_attention(q, k, v, HEADS) return x hap.HapRuntime.attn = spy try: m._unet_wrapper(model_fn, { "input": torch.zeros(1), "timestep": torch.tensor(1.0), "c": {}, }) finally: hap.HapRuntime.attn = real_attn assert kernel_calls == [0] * 5 + [1] * 5 + [2] * 5 + [3] * 5 # --------------------------------------------------------------------------- # T9.2 — HAP standalone (SPA off) # --------------------------------------------------------------------------- @pytest.mark.mock_integration class TestE2EHapStandalone: def test_e2e_hap_standalone(self, mock_attn): """SPA off (no SPA context), HAP on: single kernel pass per layer whose output equals the dense-mask reference and differs from the baseline.""" orig_calls = _install_spy_orig(mock_attn) m = _MockModel() m._hap_ctx = _hap_ctx(text_len=0) _hrdit_install_hook(m, "flux", consumer="hap") assert m._hrdit_consumers == {"hap"} q, k, v = _rand_qkv(seed=7) kernel_calls = [] captured = {} real_attn = hap.HapRuntime.attn def spy(self, qq, kk, vv, layer, **kw): kernel_calls.append(layer) return real_attn(self, qq, kk, vv, layer, **kw) def model_fn(x, t, **c): # No SPA context registered -> spa_active is False on every layer. # skip_output_reshape=True -> wrapper returns head format (B, H, T, D) # to match the hap_attn_dense reference below (math comparison). for layer in range(NUM_LAYERS): out = mock_attn.optimized_attention(q, k, v, HEADS, skip_output_reshape=True) if layer == 0: captured["out"] = out # capture while HAP ctx is live return x hap.HapRuntime.attn = spy try: # Two steps; HAP is not step-gated, so both run the kernel. for sigma in (1.0, 0.5): kernel_calls.clear() m._unet_wrapper(model_fn, { "input": torch.zeros(1), "timestep": torch.tensor(sigma), "c": {}, }) # Single pass per layer, in block order. assert kernel_calls == [0, 1, 2, 3] finally: hap.HapRuntime.attn = real_attn # Output (captured inside the forward, while HAP was live) equals the # dense-mask reference (alpha=64 -> half=0 band). out = captured["out"] mask = hap.build_band_mask(SEQ_LEN, 0, [0] * HEADS, 0) ref = hap.hap_attn_dense(q, k, v, mask) assert torch.allclose(out, ref, atol=1e-6) # ...and differs from the unmasked baseline (pruning is observable). baseline = F.scaled_dot_product_attention(q, k, v, scale=1.0) assert not torch.allclose(out, baseline, atol=1e-6) # The wrapper's orig fallback never fired while HAP was live. assert orig_calls == [] # --------------------------------------------------------------------------- # T5.1 / T5.2 — Anima regression e2e (plan 2026-08-16 P5) # --------------------------------------------------------------------------- @pytest.mark.mock_integration class TestE2EAnimaRegression: """Anima-style cosmos-convention calls through the REAL shared unet wrapper. Anima (``comfy.ldm.cosmos.predict2``) calls ``optimized_attention`` UNMASKED with ``skip_reshape=True`` and ``transformer_options`` as keyword args. The pre-fix wrapper fed ``skip_reshape`` into the real ``mask`` slot -> ``mask.ndim`` -> ``AttributeError``. These tests prove the fixed wrapper serves the cosmos convention end-to-end: - T5.1 matching plan (heads == q heads): the HAP kernel engages, the output equals the dense-mask reference, and the original attention never fires. - T5.1 mismatched plan (FLUX 24-head plan on a 2-head model): graceful plain-attention fallback, no exception, one-time warning naming both counts. - T5.2 mixed self/cross-attention: self-attn is HAP-masked, cross-attn is plain, and the layer counter advances for BOTH (alignment is sacred). """ def test_e2e_anima_matching_plan(self, mock_attn): """T5.1 matching: cosmos convention + SPA active + HAP active (heads match) -> kernel engages, output == dense-mask reference, orig never called.""" orig_calls = _install_spy_orig(mock_attn) m = _MockModel() m._spa_steps = 1 m._spa_start_sigma = 1.0 m._hap_ctx = _hap_ctx() # HEADS=2 == q heads -> kernel engages _hrdit_install_hook(m, "flux", consumer="spa") _hrdit_install_hook(m, "flux", consumer="hap") assert m._hrdit_consumers == {"spa", "hap"} q, k, v = _rand_qkv(seed=11) captured = {} def model_fn(x, t, **c): set_spa_context(_identity_spa_ctx(num_variants=5)) # Cosmos convention: skip_reshape + transformer_options kwargs, NO mask. # skip_output_reshape=True -> wrapper returns head format (B, H, T, D) # to match the hap_attn_dense reference below (math comparison). out = mock_attn.optimized_attention( q, k, v, HEADS, skip_reshape=True, skip_output_reshape=True, transformer_options={} ) captured["out"] = out return x m._unet_wrapper(model_fn, { "input": torch.zeros(1), "timestep": torch.tensor(1.0), "c": {}, }) # Identity SPA rotations -> every variant pass is the same HAP-masked # attention, so the averaged output equals the dense-mask reference. out = captured["out"] mask = hap.build_band_mask(SEQ_LEN, 0, [0] * HEADS, 0) ref = hap.hap_attn_dense(q, k, v, mask) assert torch.allclose(out, ref, atol=1e-6) # The original dense attention never fires while HAP is live. assert orig_calls == [] def test_e2e_anima_mismatched_plan_fallback(self, mock_attn, caplog): """T5.1 mismatched: FLUX plan (24 heads) on a 2-head model -> graceful plain-attention fallback, no exception, one-time warning naming both counts.""" import logging orig_calls = _install_spy_orig(mock_attn) m = _MockModel() m._spa_steps = 1 # FLUX-shaped plan (24 heads) vs q with HEADS=2 -> head-count mismatch. mismatched_plan = hap.ScopePlan( alphas=[[64.0] * 24 for _ in range(NUM_LAYERS)], betas=[[0.0] * 24 for _ in range(NUM_LAYERS)], ) m._hap_ctx = hap.HapContext( active=True, plan=mismatched_plan, text_len=0, backend="dense" ) _hrdit_install_hook(m, "flux", consumer="hap") q, k, v = _rand_qkv(seed=12) captured = {} def model_fn(x, t, **c): out = mock_attn.optimized_attention( q, k, v, HEADS, skip_reshape=True, transformer_options={} ) captured["out"] = out return x with caplog.at_level(logging.WARNING, logger="src.hap"): m._unet_wrapper(model_fn, { "input": torch.zeros(1), "timestep": torch.tensor(1.0), "c": {}, }) # No exception (the pre-fix crash) and a graceful plain-attention fallback. out = captured["out"] baseline = F.scaled_dot_product_attention(q, k, v, scale=1.0) assert torch.allclose(out, baseline, atol=1e-6) # The wrapper fell back to the original attention exactly once. assert len(orig_calls) == 1 # One-time warning naming both the plan's and the model's head counts. mismatch = [r for r in caplog.records if "heads" in r.message] assert len(mismatch) == 1 assert "24" in mismatch[0].message assert "2" in mismatch[0].message def test_e2e_cross_attention_mix(self, mock_attn): """T5.2: a forward mixing square self-attn and non-square cross-attn calls. Self-attn is HAP-masked, cross-attn is plain, and the RAW layer counter advances for BOTH (alignment is sacred). UPDATED (2026-08-19, plan-layer ordinal fix): the HAP PLAN-LAYER ORDINAL advances ONLY for plan-covered calls (square + unmasked + head-match), mirroring calibration — calibration SKIPS non-square cross-attention, so the runtime must not consume a plan slot for it either. The self-attn calls therefore receive plan layers [0, 1] (not [0, 2]), and the cross-attn calls keep their raw indices [1, 3] but decline to plain attention via the non-square guard. The pre-fix expectation ``[(0,..),(1,..),(2,..),(3,..)]`` encoded the buggy raw-counter indexing that shifted Krea2's main blocks by its 4 aux calls. """ orig_calls = _install_spy_orig(mock_attn) m = _MockModel() m._hap_ctx = _hap_ctx() # HEADS=2 == q heads _hrdit_install_hook(m, "flux", consumer="hap") q_self, k_self, v_self = _rand_qkv(seed=21) # square (SEQ_LEN x SEQ_LEN) g = torch.Generator().manual_seed(22) q_cross = torch.randn(1, HEADS, SEQ_LEN, DIM, generator=g) k_cross = torch.randn(1, HEADS, SEQ_LEN // 2, DIM, generator=g) # shorter kv v_cross = torch.randn(1, HEADS, SEQ_LEN // 2, DIM, generator=g) kernel_calls = [] captured = {} real_attn = hap.HapRuntime.attn def spy(self, qq, kk, vv, layer, **kw): kernel_calls.append((layer, int(qq.shape[-2]), int(kk.shape[-2]))) return real_attn(self, qq, kk, vv, layer, **kw) def model_fn(x, t, **c): # Block 0: self-attn (square) then cross-attn (non-square). # Self-attn: skip_output_reshape=True -> head format (B, H, T, D) to # match the hap_attn_dense reference below (math comparison). # Cross-attn declines HAP (non-square) -> orig fallback (pristine SDPA, # head format) — unaffected by the output-reshape fix. captured["self0"] = mock_attn.optimized_attention( q_self, k_self, v_self, HEADS, skip_reshape=True, skip_output_reshape=True ) captured["cross0"] = mock_attn.optimized_attention( q_cross, k_cross, v_cross, HEADS, skip_reshape=True ) # Block 1: self-attn then cross-attn. mock_attn.optimized_attention(q_self, k_self, v_self, HEADS, skip_reshape=True, skip_output_reshape=True) mock_attn.optimized_attention(q_cross, k_cross, v_cross, HEADS, skip_reshape=True) return x hap.HapRuntime.attn = spy try: m._unet_wrapper(model_fn, { "input": torch.zeros(1), "timestep": torch.tensor(1.0), "c": {}, }) finally: hap.HapRuntime.attn = real_attn # The RAW layer counter advances for BOTH self-attn and cross-attn # (alignment is sacred), but the HAP PLAN-LAYER ORDINAL advances only for # the covered (square) self-attn calls. So the self-attn calls receive # plan layers [0, 1] and the cross-attn calls keep their raw indices # [1, 3] (declined to plain attention via the non-square guard). assert kernel_calls == [ (0, SEQ_LEN, SEQ_LEN), # block 0 self-attn (square) -> plan 0 (1, SEQ_LEN, SEQ_LEN // 2), # block 0 cross-attn (non-square) -> raw 1, declines (1, SEQ_LEN, SEQ_LEN), # block 1 self-attn (square) -> plan 1 (3, SEQ_LEN, SEQ_LEN // 2), # block 1 cross-attn (non-square) -> raw 3, declines ] # Cross-attn (non-square) declined to plain attention: 2 orig fallbacks. assert len(orig_calls) == 2 # Self-attn output is HAP-masked (equals the dense-mask reference). mask = hap.build_band_mask(SEQ_LEN, 0, [0] * HEADS, 0) ref_self = hap.hap_attn_dense(q_self, k_self, v_self, mask) assert torch.allclose(captured["self0"], ref_self, atol=1e-6) # Cross-attn output is plain attention (the pristine SDPA). ref_cross = F.scaled_dot_product_attention(q_cross, k_cross, v_cross, scale=1.0) assert torch.allclose(captured["cross0"], ref_cross, atol=1e-6) def test_e2e_anima_spa_cross_attention_no_crash(self, mock_attn): """REGRESSION (2026-08-16): the exact Anima production crash path. SPA active (identity variants) + a NON-SQUARE cosmos-convention call (image queries vs text keys) through the REAL shared unet wrapper. Pre-fix this raised ``RuntimeError: einsum(): subscript l has size 512 for operand 1 ... previously seen size 6300`` because the averaged passes applied the image-sized RoPE rotations to the text ``k``. The non-square guard must decline SPA and run plain attention. """ m = _MockModel() m._spa_steps = 1 m._spa_start_sigma = 1.0 _hrdit_install_hook(m, "flux", consumer="spa") # Identity SPA variants over the IMAGE query length (SEQ_LEN) — exactly # what PosEmbedSPAAnima registers for the T*H*W grid. q_img = torch.randn(1, HEADS, SEQ_LEN, DIM, generator=torch.Generator().manual_seed(31)) k_text = torch.randn(1, HEADS, SEQ_LEN // 2, DIM, generator=torch.Generator().manual_seed(32)) v_text = torch.randn(1, HEADS, SEQ_LEN // 2, DIM, generator=torch.Generator().manual_seed(33)) captured = {} def model_fn(x, t, **c): set_spa_context(_identity_spa_ctx(num_variants=5)) # Cosmos convention: skip_reshape kw, NO mask, q_len != k_len. captured["out"] = mock_attn.optimized_attention( q_img, k_text, v_text, HEADS, skip_reshape=True, transformer_options={}, ) return x # Pre-fix: RuntimeError in einsum. Post-fix: must complete. m._unet_wrapper(model_fn, { "input": torch.zeros(1), "timestep": torch.tensor(1.0), "c": {}, }) # The declined cross-attn call ran PLAIN attention (pristine SDPA). ref = F.scaled_dot_product_attention(q_img, k_text, v_text, scale=1.0) assert torch.allclose(captured["out"], ref, atol=1e-6)