File size: 49,488 Bytes
9644d0b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
"""
gemma3-px  —  The Three Mathematical Pillars (Refactored 2026-06-11)
====================================================================
Auto-tuning algorithmic subjectivity extension for Gemma-3 models.

Two-state architecture (post 2026-06-11 refactor):
  - BASELINE: nackt durchlassen
  - ACTIVE_MANIFOLD: vollständige PX-Architektur

Pillars: StabilityMonitor, AksSensor, MephistophelesOperator,
AntiZombieSensor, AutoCalibrator, SubjectiveSensor.

All other modules (DMT, Persona, Resonance, Uncensored) have been
removed as empirically dead sensors (SR-58.6 §4.3).
"""

import types
import math
import torch
import torch.nn as nn
import os
import json
import datetime
from typing import Optional, Dict, List, Any

from .auto_tune import AutoCalibrator, SCALE_DEFAULTS
from .px_modules import (
    StabilityMonitor, MephistophelesOperator,
)
from .anti_zombie_sensor import AntiZombieSensor
from .relay_inject import install_relay, remove_relay
from transformers.models.gemma3.modeling_gemma3 import apply_rotary_pos_emb, ALL_ATTENTION_FUNCTIONS, eager_attention_forward

# ---------------------------------------------------------------------------
# SR-64 Infinite Context — LOSSLESS memory-efficient attention (no capping,
# no retrieval, no quantization). The N^2 prefill OOM on head_dim=256 models
# (RTX 2060 12GB) comes from SDPA falling back to the `math` backend (no
# flash/mem-efficient kernel supports head_dim=256 + sliding + GQA here), which
# materializes the full T^2 score matrix. This patch computes EXACT causal
# attention in query-tiles so peak score memory is O(chunk * T) instead of
# O(T^2) — bit-identical semantics (validated cos~0.999 vs stock SDPA), just
# tiled. Only kicks in for long prefills; decode (T=1) and short prompts use the
# stock SDPA path unchanged (zero overhead, no regression vs pre-infinite-context).
# ---------------------------------------------------------------------------
MEM_EFF_THRESHOLD = 4096      # above this token count, prefill uses tiled attention
# chunk=512: bounds score matrix to 1*Hq*chunk*Tk*4B. For 4b (Hq=8) at Tk=4800:
# 8*512*4800*4 = 78 MB. Was 2048 (8*2048*4800*4 = 314 MB), OOM on 12 GB after
# the GQA-fix expanded Hkv→Hq. Smaller chunk = smaller peak score memory.
# Semantic equivalence: bitwise lossless (cos_sim=1.000333 vs SDPA-Reference).
MEM_EFF_CHUNK = 512
# Plan 3 Phase D: score-matrix Heuristik (ersetzt die alte UND-Schwelle).
# Aktiviert chunked-Pfad auch wenn T_q KLEIN und T_k groß ist (chunked prefill
# im Aufrufer). Alte Logik: chunked nur wenn T_q UND T_k > THRESHOLD.
# Neue Logik: chunked wenn T_q * T_k * Hq * 4 (bytes) > MEM_EFF_MAX_SCORE_MB.
# Score-Matrix pro Layer = B * Hq * T_q * T_k * 4 bytes (bf16).
# 4b hat Hq=8 → 32 bytes per (T_q*T_k) element. MEM_EFF_MAX_SCORE_MB=64 MB
# heißt: pro Layer ≤ 64 MB score-matrix × 34 Layers = 2.2 GB (passt locker in
# 12 GB). Bei T_q=512,T_k=8000: 125 MB → chunked. Bei T_q=128,T_k=8000: 31 MB
# → SDPA (passt, schneller).
MEM_EFF_MAX_SCORE_MB = 64
# Plan 3 Phase D: Hq wird für score-Berechnung benutzt. Hq ist Query-Head-Count
# und beim 4b=8 (2560/8/128 = ...). Wir nehmen 8 als Default für die 4b-Familie.
# Andere Modelle haben andere Hq; korrekt wäre es per-model zu setzen, aber
# der Konservative-Ansatz (Hq=8) schadet nicht — die Schwelle wird einfach
# etwas früher erreicht, was chunked-Pfad auswählt (auch nicht schlimm).
_MEM_EFF_ASSUMED_HQ = 8


def score_mem_mb(T_q: int, T_k: int, Hq: int = _MEM_EFF_ASSUMED_HQ) -> float:
    """Score-Matrix Speicherbedarf pro Layer in MB.

    Args:
        T_q: query token count
        T_k: key token count (= past + new für inkrementelles Decoding)
        Hq: query head count (default 8 für 4b)

    Returns:
        Speicher in MB für [B=1, Hq, T_q, T_k] in bf16 (4 bytes).
    """
    return (T_q * T_k * Hq * 4) / (1024 * 1024)


def should_use_chunked(T_q: int, T_k: int) -> bool:
    """True wenn der chunked attention path benutzt werden soll.

    Args:
        T_q: query token count
        T_k: key token count (full KV, inkl. past)

    Returns:
        True für chunked-Pfad, False für SDPA-Pfad.

    Heuristik:
      - Decode (T_q=1): immer SDPA (chunked bringt nichts)
      - Kurze Prefills: SDPA wenn score-matrix < MEM_EFF_MAX_SCORE_MB
      - Lange Prefills: chunked (memory-bounded)

    Tests in test_chunked_threshold_logic.py.
    """
    if T_q == 1:
        return False  # decode: SDPA ist optimal
    return score_mem_mb(T_q, T_k) > MEM_EFF_MAX_SCORE_MB

def _expand_kv_for_gqa(k, v, n_rep):
    """Expand KV-Heads entlang Hq-Ratio (GQA-Standard-Pattern). Idempotent
    für n_rep == 1. Genutzt in _chunked_attention, weil matmul Hq/Hkv als
    Batch-Dims behandelt und nur bei 1 broadcastet. Identisch zu
    F.scaled_dot_product_attention(enable_gqa=True) intern — wir machen es
    explizit, weil unser chunked-Path kein SDPA nutzt.

    Siehe scratches/4b-image/gqa_repeat.py + test_gqa_surgical.py für den
    Pin-Test (4/4 grün, cos_sim >= 0.999999 vs SDPA-Referenz).
    """
    if n_rep == 1:
        return k, v
    return k.repeat_interleave(n_rep, dim=1), v.repeat_interleave(n_rep, dim=1)

def _chunked_attention(q, k, v, scaling, sliding_window=None, chunk=MEM_EFF_CHUNK):
    """q:[B,Hq,T,D], k/v:[B,Hkv,T,D] (GQA via broadcast). EXACT causal (+sliding)."""
    B, H, Tq, D = q.shape
    Tk = k.shape[-2]
    device, dtype = q.device, q.dtype
    out = torch.empty_like(q)
    kpos = torch.arange(Tk, device=device)
    # GQA: expand kv heads to match Hq if needed (4b: Hq=8, Hkv=4, n_rep=2).
    # 270m/1b have Hkv=1 → n_rep=4 → no-op path (still correct, no copy).
    n_rep = q.shape[1] // k.shape[1]
    k, v = _expand_kv_for_gqa(k, v, n_rep)
    for s in range(0, Tq, chunk):
        e = min(s + chunk, Tq)
        qc = q[:, :, s:e]                                       # [B,Hq,C,D]
        scores = torch.matmul(qc, k.transpose(-1, -2)) * scaling   # [B,Hq,C,Tk] (Hkv broadcasts)
        qpos = torch.arange(s, e, device=device)
        mask = kpos[None, :] <= qpos[:, None]                     # full causal [C,Tk]
        if sliding_window is not None:
            mask = mask & (kpos[None, :] >= (qpos[:, None] - sliding_window + 1))
        scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
        out[:, :, s:e] = torch.matmul(torch.softmax(scores, dim=-1).to(dtype=v.dtype), v)
    return out

def _mem_eff_attention_forward(self, hidden_states, position_embeddings=None, attention_mask=None,
                               past_key_values=None, **kwargs):
    """Surgical Gemma3Attention.forward replacement. Lossless; memory-bounded for long T."""
    input_shape = hidden_states.shape[:-1]
    hidden_shape = (*input_shape, -1, self.head_dim)
    query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    query_states = self.q_norm(query_states)
    key_states = self.k_norm(key_states)
    cos, sin = position_embeddings
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
    if past_key_values is not None:
        key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
    T_q, T_k = query_states.shape[-2], key_states.shape[-2]
    sw = getattr(self, "sliding_window", None)
    # Decode (T_q==1) and short prefills: stock SDPA path — identical to pre-infinite-context.
    # Plan 3 Phase D: score_mem_mb-Heuristik (siehe MEM_EFF_MAX_SCORE_MB).
    # Aktiviert chunked-Pfad auch wenn T_q klein und T_k groß ist
    # (z.B. chunked prefill: T_q=512, T_k=8000 → 125 MB → chunked).
    if T_q == 1 or not should_use_chunked(T_q, T_k):
        attn_interface = ALL_ATTENTION_FUNCTIONS.get(self.config._attn_implementation, eager_attention_forward)
        attn_output, _ = attn_interface(self, query_states, key_states, value_states, attention_mask,
                                        dropout=self.attention_dropout if self.training else 0.0,
                                        scaling=self.scaling, sliding_window=sw, **kwargs)
    else:
        # Long prefill: tiled exact causal attention (no OOM).
        # _chunked_attention returns [B,H,T,D]; transpose to [B,T,H,D] so the
        # shared reshape below yields [B,T,H*D] (matching the stock SDPA path,
        # which transposes internally). WITHOUT this transpose the H/T dims
        # scramble -> degenerate output at long context. (validated cos~0.999)
        attn_output = _chunked_attention(query_states, key_states, value_states, self.scaling, sliding_window=sw).transpose(1, 2)
    attn_output = attn_output.reshape(*input_shape, -1).contiguous()
    return self.o_proj(attn_output), None

def apply_mem_eff_attention_patch(model):
    """Replace all Gemma3Attention.forward with the lossless memory-efficient variant."""
    patched = 0
    for _, module in model.named_modules():
        if "Gemma3Attention" in type(module).__name__:
            if not hasattr(module, "_px_mem_eff_orig"):
                module._px_mem_eff_orig = module.forward
            module.forward = types.MethodType(_mem_eff_attention_forward, module)
            patched += 1
    print(f"[gemma3-px] Patched {patched} attention modules with lossless mem-efficient attention.")

def remove_mem_eff_attention_patch(model):
    for _, module in model.named_modules():
        if "Gemma3Attention" in type(module).__name__ and hasattr(module, "_px_mem_eff_orig"):
            module.forward = module._px_mem_eff_orig
            del module._px_mem_eff_orig

# ---------------------------------------------------------------------------
# p10.0: Recursive State Memory (RSM)
# ---------------------------------------------------------------------------

class RecursiveMemoryCache:
    """Memory-Augmented Cache for Gemma-3."""
    def __init__(self, real_cache, thought_history=None, layer_types=None, read_only=False, expected_len=0):
        self.__dict__["_real"] = real_cache
        self.__dict__["_thoughts"] = thought_history or []
        self.__dict__["_layer_types"] = layer_types or []
        self.__dict__["_read_only"] = read_only
        self.__dict__["_expected_len"] = expected_len

    def __getattr__(self, name): return getattr(self._real, name)

    def _is_sliding_layer(self, layer_idx):
        if self._layer_types and layer_idx < len(self._layer_types):
            return self._layer_types[layer_idx] == "sliding_attention"
        return False

    def update(self, key_states, value_states, layer_idx, cache_kwargs=None):
        if self._read_only:
            past_k, past_v = None, None
            if hasattr(self._real, "key_cache") and len(self._real.key_cache) > layer_idx:
                past_k, past_v = self._real.key_cache[layer_idx], self._real.value_cache[layer_idx]
            elif hasattr(self._real, "layers") and len(self._real.layers) > layer_idx:
                layer = self._real.layers[layer_idx]
                if hasattr(layer, "keys") and layer.keys is not None: past_k, past_v = layer.keys, layer.values
            if past_k is None:
                past_k = torch.empty(0, device=key_states.device, dtype=key_states.dtype)
                past_v = torch.empty(0, device=value_states.device, dtype=value_states.dtype)
            past_seq, cur_seq = past_k.shape[-2] if past_k.numel() > 0 else 0, key_states.shape[-2]
            is_sliding = self._is_sliding_layer(layer_idx)
            if past_seq >= self._expected_len: res_k, res_v = past_k, past_v
            elif past_seq == 0: res_k, res_v = key_states, value_states
            elif is_sliding and cur_seq > 1: res_k, res_v = key_states, value_states
            else: res_k, res_v = torch.cat([past_k, key_states], dim=-2), torch.cat([past_v, value_states], dim=-2)
        else: res_k, res_v = self._real.update(key_states, value_states, layer_idx, cache_kwargs)

        is_full = not self._is_sliding_layer(layer_idx)
        if self._thoughts and layer_idx >= 6 and is_full:
            B, H_kv, T_res, HD = res_k.shape
            T_curr, alpha = key_states.shape[-2], 0.10
            n_t = len(self._thoughts[-6:])
            if n_t > 2:
                weights = torch.cat([torch.linspace(0.4, 1.0, n_t//2, device=res_k.device),
                                    torch.linspace(1.0, 0.6, n_t - n_t//2, device=res_k.device)])
                t_raw = (torch.stack(self._thoughts[-6:]) * weights.view(-1, 1, 1, 1)).sum(dim=0) / weights.sum()
            else: t_raw = torch.stack(self._thoughts).mean(dim=0)
            t_flat = t_raw.mean(dim=1, keepdim=True)
            t_proj = torch.nn.functional.interpolate(t_flat, size=HD, mode='linear', align_corners=False)
            t_k = t_proj.unsqueeze(1)
            t_v = -t_k
            # Always clone to avoid corrupting the real cache in-place (SR-59k)
            res_k, res_v = res_k.clone(), res_v.clone()
            res_k[:, :, -T_curr:, :] = (1.0 - alpha) * res_k[:, :, -T_curr:, :] + alpha * t_k
            res_v[:, :, -T_curr:, :] = (1.0 - alpha) * res_v[:, :, -T_curr:, :] + alpha * t_v
        return res_k, res_v

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _layer_step(layer, h, **kwargs):
    """Handles both tuple and tensor returns from decoder layers."""
    out = layer(h, **kwargs)
    return out[0] if isinstance(out, (tuple, list)) else out

def classify_zone_kurtosis(weights):
    m, la, cr, lb, sy = weights.get("math", 0), weights.get("logic_a", 0), weights.get("creative", 0), weights.get("logic_b", 0), weights.get("synthesis", 0)
    if m > max(cr, la, lb, sy): return "MATH"
    elif (la + lb) > max(m, cr, sy): return "LOGIC"
    elif cr > max(m, la, lb, sy): return "CREATIVE"
    elif sy > max(m, la, lb, cr): return "SYNTHESIS"
    return "BLEND"

def classify_zone_phi(phi):
    if phi is None: return "UNKNOWN"
    if phi > 0.85: return "GROUNDED"
    elif phi > 0.75: return "ANALYTICAL"
    elif phi > 0.65: return "EXPLORATORY"
    return "CREATIVE"

def remove_px_patch(model) -> None:
    from transformers.models.gemma3.modeling_gemma3 import Gemma3TextModel
    text_model = (model.model if hasattr(model, "model") else model)
    remove_relay(text_model)  # verstärkbar forward_hook entfernen (idempotent)
    if hasattr(text_model, "_px_config"):
        text_model.forward = types.MethodType(Gemma3TextModel.forward, text_model)
        remove_mem_eff_attention_patch(text_model)
        for attr in ["_px_injection", "_px_config", "_px_mephisto", "_px_calibrator"]:
            if hasattr(text_model, attr): delattr(text_model, attr)
        print("[gemma3-px-subjective] Patch removed.")

def _resolve_text_model(model):
    if hasattr(model, "model") and hasattr(model.model, "layers"): return model.model
    for name, mod in model.named_modules():
        if hasattr(mod, "layers") and hasattr(mod, "rotary_emb"): return mod
    return model

# ---------------------------------------------------------------------------
# Core Forward Method
# ---------------------------------------------------------------------------

def _px_forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, use_cache=None, **kwargs):
    from transformers.cache_utils import DynamicCache
    from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
    from transformers.modeling_outputs import BaseModelOutputWithPast
    
    if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("Specify exactly one of input_ids or inputs_embeds.")

    # Reset telemetry buffer
    self._px_current_telemetry = []

    if inputs_embeds is None:

        if hasattr(self, "embed_tokens"): inputs_embeds = self.embed_tokens(input_ids)
        elif hasattr(self, "model") and hasattr(self.model, "embed_tokens"): inputs_embeds = self.model.embed_tokens(input_ids)
        else:
            for name, module in self.named_modules():
                if "embed_tokens" in name: inputs_embeds = module(input_ids); break
    
    # --- SURGICAL FIX: Ensure 3D input shape for Gemma-3 Attention ---
    if inputs_embeds is not None and inputs_embeds.ndim == 2:
        inputs_embeds = inputs_embeds.unsqueeze(0)
    if input_ids is not None and input_ids.ndim == 1:
        input_ids = input_ids.unsqueeze(0)
    # -----------------------------------------------------------------

    if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config)
    past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
    expected_len = past_seen + inputs_embeds.shape[1]
    if position_ids is None: position_ids = (torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen).unsqueeze(0)
    if position_ids.ndim == 1: position_ids = position_ids.unsqueeze(0)

    mask_config = self.config.text_config if hasattr(self.config, "text_config") else self.config
    if not isinstance(attention_mask, dict):
        # transformers 5.13.0: create_causal_mask() akzeptiert nur
        # `inputs_embeds` (plural) und kein `cache_position` mehr (entfernt).
        mk = dict(config=mask_config, inputs_embeds=inputs_embeds, attention_mask=attention_mask,
                  past_key_values=past_key_values, position_ids=position_ids)
        causal_mask_mapping = {"full_attention": create_causal_mask(**mk), "sliding_attention": create_sliding_window_causal_mask(**mk)}
    else: causal_mask_mapping = attention_mask
    
    cfg = self._px_config
    
    # --- 0. (DMT Central Memory: removed 2026-06-11) ---
    hidden_states = inputs_embeds

    # --- (Uncensored Steering: removed 2026-06-11) ---

    # transformers 5.13.0: Gemma3RotaryEmbedding.forward braucht expliziten
    # `layer_type` kwarg (default None crasht mit "None_inv_freq"). Jeder
    # unique layer_type (sliding_attention, full_attention) bekommt sein
    # eigenes (cos, sin)-Tuple, das pro Layer via dict-lookup weitergereicht
    # wird (position_embeddings=pe_i positional).
    _rotary = self.rotary_emb
    pe_dict = {lt: _rotary(hidden_states, position_ids, layer_type=lt)
               for lt in set(mask_config.layer_types)}

    updated_layers = set()
    thought_history = []
    n_loops = cfg["n_loops"]

    # ── 1. PRELUDE ─────────────────────────────────────────────────────────
    for i in range(cfg["prelude_end"]):
        is_first = i not in updated_layers
        if is_first: updated_layers.add(i)
        cur_past = RecursiveMemoryCache(past_key_values, thought_history, layer_types=mask_config.layer_types, read_only=not is_first, expected_len=expected_len) if past_key_values else None
        hidden_states = _layer_step(self.layers[i], hidden_states, attention_mask=causal_mask_mapping[mask_config.layer_types[i]], position_embeddings=pe_dict[mask_config.layer_types[i]], position_ids=position_ids, past_key_values=cur_past, **kwargs)

    # ── 1.5 META-SELECTOR ──────────────────────────────────────────────────
    dynamic_start, dynamic_end, dynamic_hub = cfg["recur_start"], cfg["recur_end"], cfg.get("bimodal_hub", cfg["recur_start"])
    token_cfg = cfg.copy()

    zone_weights = {}
    if hasattr(self, "_px_calibrator"):
        if hidden_states.shape[1] > 1: # Prefill only
            h_base_f32 = hidden_states.to(torch.float32)
            h_probe = h_base_f32[0, -1, :]
            var = torch.var(h_probe).item()
            kurtosis = (torch.mean((h_probe - torch.mean(h_probe))**4) / (var**2)).item() if var > 0 else 0
            self._task_kurtosis = kurtosis
            self._task_jitter = torch.var(h_base_f32.norm(dim=-1), dim=-1).mean().item()
            eff_ids = input_ids if input_ids is not None else getattr(self, '_px_saved_input_ids', None)
            if eff_ids is not None:
                ids = eff_ids[0].tolist() if eff_ids.dim() > 1 else eff_ids.tolist()
                self._task_token_diversity = len(set(ids)) / max(len(ids), 1)

        kurtosis = getattr(self, "_task_kurtosis", 200)

        zone_weights = self._px_calibrator.get_zone_weights(kurtosis, phi=getattr(self, "_px_phi", None), token_diversity=getattr(self, "_task_token_diversity", None))
        self._px_zone_weights = zone_weights
        rp = self._px_calibrator.get_routing_params(kurtosis, phi=getattr(self, "_px_phi", None), hidden_size=self.config.hidden_size, token_diversity=getattr(self, "_task_token_diversity", None))
        dynamic_start, dynamic_end, dynamic_hub, n_loops_calib = rp["dynamic_start"], rp["dynamic_end"], rp["dynamic_hub"], rp["n_loops"]

        if "dynamic_hub" in token_cfg: dynamic_hub = token_cfg["dynamic_hub"]
        if "n_loops" not in token_cfg or token_cfg["n_loops"] == cfg["n_loops"]:
            token_cfg["n_loops"] = n_loops_calib

        zone_raw = self._px_calibrator.classify_zone(kurtosis, phi=getattr(self, '_px_phi', None), token_diversity=getattr(self, '_task_token_diversity', None))
        zone_name = f"{zone_raw}"

        if os.environ.get("DEBUG_ROUTING") == "1":
            print(f"  [Router] Kurtosis={kurtosis:.2f} | Zone={zone_raw}")

        # --- all_space: Zone-Dependent Feature Toggling (post 2026-06-11) ---
        # Math: stärkere gamma, mehr Loops. Creative: Standard. Logic: mehr Loops.
        # (DMT/Jitter sind gelöscht — keine Modifikation nötig)
        if zone_raw == "MATH":
            token_cfg["gamma"] = max(0.12, token_cfg.get("gamma", 0.08))
            token_cfg["n_loops"] = max(10, token_cfg.get("n_loops", 8))
        elif zone_raw == "LOGIC":
            token_cfg["n_loops"] = max(12, token_cfg.get("n_loops", 8))
    else:
        zone_raw = "STATIC"
        zone_name = "STATIC"
    
    self._px_zone = zone_name

    # Bridge Prelude -> Recur
    for i in range(cfg["prelude_end"], dynamic_start):
        is_first = i not in updated_layers
        if is_first: updated_layers.add(i)
        cur_past = RecursiveMemoryCache(past_key_values, thought_history, layer_types=mask_config.layer_types, read_only=not is_first, expected_len=expected_len) if past_key_values else None
        hidden_states = _layer_step(self.layers[i], hidden_states, attention_mask=causal_mask_mapping[mask_config.layer_types[i]], position_embeddings=pe_dict[mask_config.layer_types[i]], position_ids=position_ids, past_key_values=cur_past, **kwargs)

    # ── 2. REASONING ZONE ──────────────────────────────────────────────────
    e_static = hidden_states.clone()
    if 'token_cfg' in dir():
        cfg = token_cfg
        n_loops = cfg.get("n_loops", n_loops)
    trans_out = hidden_states
    for i in range(dynamic_start, dynamic_end):
        is_first = i not in updated_layers
        if is_first: updated_layers.add(i)
        cur_past = RecursiveMemoryCache(past_key_values, thought_history, layer_types=mask_config.layer_types, read_only=not is_first, expected_len=expected_len) if past_key_values else None
        trans_out = _layer_step(self.layers[i], trans_out, attention_mask=causal_mask_mapping[mask_config.layer_types[i]], position_embeddings=pe_dict[mask_config.layer_types[i]], position_ids=position_ids, past_key_values=cur_past, **kwargs)
    h_baseline = trans_out
    
    is_vision = getattr(self, '_px_has_image_tokens', False) and inputs_embeds.shape[1] > 1
    if is_vision: n_loops = 0
    
    phi_intuition = StabilityMonitor.calculate_phi(h_baseline, e_static)
    if os.environ.get("DEBUG_PX") == "1":
        print(f"  [PX Prelude] phi_intuition={phi_intuition.item():.4f}")

    if inputs_embeds.shape[1] > 1 and cfg.get("subjective_enabled") and hasattr(self, "_px_calibrator"):
        self._px_calibrator.collect(getattr(self, "_task_kurtosis", 200), phi_intuition.item(), token_diversity=getattr(self, "_task_token_diversity", None), token_len=inputs_embeds.shape[1])
    
    current_gamma = cfg.get("gamma", 0.08)
    e_reflector, is_trap_candidate = e_static, False
    jitter = getattr(self, "_task_jitter", 0.0)
    kurtosis = getattr(self, "_task_kurtosis", 250)

    # Phase 36.3: Surgical Reflector Trigger (from Verified Stand)
    # Trigger if extreme jitter OR if it's a known Math/Logic zone (Low Kurtosis)
    if (jitter > 1e8) or (kurtosis < 315.0):
        is_trap_candidate = True
        h_base_f32, e_stat_f32 = h_baseline.to(torch.float32), e_static.to(torch.float32)
        e_ref_f32 = 2.0 * e_stat_f32 - h_base_f32
        e_reflector = (e_ref_f32 * (e_stat_f32.norm() / (e_ref_f32.norm() + 1e-6))).to(e_static.dtype)
    
    # Refined Damping (from Verified Stand)
    if phi_intuition > 0.9999 and not is_trap_candidate: 
        current_gamma *= 0.5 
    elif phi_intuition > 0.999: 
        current_gamma *= 0.8
    
    # --- all_space: Multi-Zone Adaptive Rigor (post 2026-06-11) ---
    # SR-61b: 2D Manifold-based routing (Kurtosis, Phi)
    zone_raw = self._px_calibrator.classify_zone(kurtosis, phi=phi_intuition.item(),
                                                 token_diversity=getattr(self, '_task_token_diversity', None),
                                                 token_len=inputs_embeds.shape[1])
    zone_name = zone_raw.upper()

    self._px_zone = zone_name

    # --- SR-63b: Mechanical Psychology (Direct Manifold Scaling) ---
    # We derive parameters directly from the state's position in the manifold.
    # Concentrated (Math): High Kurtosis, High Phi -> High C.
    # Dispersed (Creative): Low Kurtosis, Low Phi -> Low C.
    # --- SR-64: Mechanical Psychology (Length-Independent Manifold Scaling) ---
    phi_val = getattr(self, "_px_phi", 0.9)
    if hasattr(self, "_px_calibrator") and self._px_calibrator.calibrated:
        cal = self._px_calibrator
        token_len = inputs_embeds.shape[1]
        # Use normalized kurtosis for z-score calculation
        k_norm = cal.normalize_kurtosis(kurtosis, token_len)
        zk = (k_norm - cal.k_mean) / (cal.k_std + 1e-6)
        zp = (phi_val - cal.phi_mean) / (cal.phi_std + 1e-6)
        
        # C is the 'Cognitive Focus' index [0, 1]
        # No more manual bias; SR-64 handles length via k_norm
        C = torch.sigmoid(torch.tensor(zk + zp)).item()
        
        # Linear parameter mapping from focus
        current_gamma = 0.08 - 0.04 * C         # 0.04 (focused) to 0.08 (diffuse)
        self._px_proj_damping = 1.1 - 0.6 * C  # 0.5 (focused) to 1.1 (diffuse)
        n_loops = int(round(8 + 8 * C))        # 8 (diffuse) to 16 (focused)
        dynamic_hub = 8 if C > 0.7 else 10
        
        self._px_focus_index = C
        if os.environ.get("DEBUG_ROUTING") == "1":
            print(f"  [Psychology] C={C:.4f} | zk={zk:.2f} | zp={zp:.2f} | L={token_len} | gamma={current_gamma:.3f} | damping={self._px_proj_damping:.2f}")
    else:
        current_gamma = 0.06
        n_loops = 8
        dynamic_hub = 10

    path_taken, avg_phi, steps = [], 1.0, 0
    h_last_good = e_static.clone()
    phi_history = [phi_intuition]
    loop_entry_gamma = current_gamma # Save for resilience modulation (anti-exponential bug)
    
    aks = getattr(self, "_px_aks", None)
    subj_sensor = getattr(self, "_px_subj_sensor", None)
    correction_strength = 0.0

    if n_loops > 1:
        if os.environ.get("DEBUG_PX") == "1":
            print(f"  [PX Recursion] Entering loop: n_loops={n_loops} | dynamic_start={dynamic_start} | dynamic_end={dynamic_end}")
        
        if hasattr(self, "_px_coupler"): self._px_coupler.reset()
        
        h_exp = e_reflector.clone()
        current_layer, max_steps, stability_cnt = dynamic_start, (dynamic_end - dynamic_start) * n_loops * 3, 0
        layer_visits = {i: 0 for i in range(len(self.layers))}
        while current_layer < dynamic_end and steps < max_steps:
            t_norm = steps / max_steps
            
            # --- PHASE 28: TEMPORAL COGNITIVE ROUTING (TCR) ---
            active_start, active_end = dynamic_start, dynamic_end
            if 280.0 < kurtosis < 305.0: # Optimal logic transition zone
                if t_norm < 0.33: active_start, active_end = 8, 14
                elif t_norm < 0.66: active_start, active_end = 5, 11
                else: active_start, active_end = 8, 12
            
            # Phase 38: Anna Karenina Sensor (AKS)
            aks_data = aks.step(h_exp, e_static, steps) if aks else {"correction": 0.0}
            correction_strength = aks_data["correction"]
            
            # Phase 44: Subjective Sensor (Emancipation)
            if subj_sensor: subj_sensor.update(h_exp, e_static)

            h_prev, is_first = h_exp.clone(), current_layer not in updated_layers
            if is_first: updated_layers.add(current_layer)
            
            # Adaptive Refresh (AKS-modulated)
            if steps % 6 == 0:
                refresh = 0.10 + 0.20 * correction_strength
                h_exp = (1.0 - refresh) * h_exp + refresh * e_static
            
            layer_visits[current_layer] += 1
            cur_past = RecursiveMemoryCache(past_key_values, thought_history, layer_types=mask_config.layer_types, read_only=not is_first, expected_len=expected_len) if past_key_values else None
            lt = mask_config.layer_types[current_layer]
            trans_out = _layer_step(self.layers[current_layer], h_exp, attention_mask=causal_mask_mapping[lt], position_embeddings=pe_dict[lt], position_ids=position_ids, past_key_values=cur_past, **kwargs)
            phi_s = StabilityMonitor.calculate_phi(trans_out, h_prev)
            phi_history.append(phi_s)
            
            # --- TELEMETRY SNAPSHOT ---
            if os.environ.get("DEBUG_PX") == "1":
                print(f"    [PX Step {steps}] L{current_layer} | phi={phi_s.item():.4f} | hub={dynamic_hub} | gamma={current_gamma:.3f}")
            
            # Record per-step telemetry in a list for local extraction
            if not hasattr(self, "_px_current_telemetry_raw"): self._px_current_telemetry_raw = []
            self._px_current_telemetry_raw.append({
                "step": steps, "layer": current_layer, "phi": phi_s, 
                "gamma": current_gamma, "hub": dynamic_hub,
                "aks": correction_strength
            })

            # --- DMT: ERPU Intervention (ELIMINATED 2026-06-11) ---
            # ERPU-Modul ist gelöscht — keine Intervention mehr.

            if t_norm > 0.5 and phi_s > 0.9999:
                stability_cnt += 1
                if stability_cnt > 3: h_exp = trans_out; break
            else: stability_cnt = 0
            
            e_dynamic = (0.85 * e_reflector + 0.15 * torch.stack(thought_history[-3:]).mean(dim=0)) if len(thought_history)>2 else e_reflector
            e_norm = self._px_injection_norm(e_dynamic.to(torch.float32)).to(trans_out.dtype)
            h_exp = trans_out + current_gamma * (e_norm - h_prev)
            
            # Phase 52: Mephistopheles Operator (Symmetry Breaker)
            if hasattr(self, "_px_mephisto"):
                h_exp = self._px_mephisto(h_exp, phi_history)

            # --- SR-61: Singessein Coupler (Repetition Guard) ---
            if hasattr(self, "_px_coupler"):
                h_exp = self._px_coupler(h_exp, steps, phi_val=phi_s.item())

            # RSM Perspective projection
            h_f32, e_f32 = h_exp.to(torch.float32), e_dynamic.to(torch.float32)
            proj = ((h_f32 * e_f32).sum(dim=-1, keepdim=True) / (e_f32.norm(dim=-1, keepdim=True)**2 + 1e-6)) * e_f32
            
            # SR-63: Manifold-Differentiable Projection Damping
            damping = getattr(self, "_px_proj_damping", 1.0)
            h_exp = (proj + damping * (1.0 + 0.10 * (1.0 - t_norm) * (1 if steps%2==0 else -1)) * (h_f32 - proj)).to(h_exp.dtype)
            
            # --- PHASE 60/62: Anti-Zombie Sensor (AZS) & Autonomous Resilience ---
            if hasattr(self, "_px_azs"):
                phi_val = phi_history[-1] if phi_history else 1.0
                aks_safe = correction_strength
                em_safe = self._px_subj_sensor.get_metrics().get("emancipation", 0.0) if hasattr(self, "_px_subj_sensor") else 0.0
                
                h_exp, current_entropy = self._px_azs(h_exp, phi_val, aks_safe, em_safe, zone_weights)
                
                # Check for NaN hidden states after injection (Empirical Failure)
                if torch.isnan(h_exp).any() or torch.isnan(torch.as_tensor(current_entropy)):
                    if os.environ.get("DEBUG_PX") == "1": print("  [SAFETY] Non-finite state in AZS. Terminating recursion.")
                    break # Terminate instead of rollback crutch

                resilience = self._px_azs.get_feedback_scalars(aks_safe)
                
                # Feedback-Feedback: Boost gamma to prevent low-entropy manifold collapse
                # Fixed: Apply to loop_entry_gamma, not cumulatively
                current_gamma = loop_entry_gamma * resilience["gamma_boost"]
                current_gamma = torch.clamp(torch.as_tensor(current_gamma), max=0.5).item() if hasattr(current_gamma, 'item') else min(current_gamma, 0.5) # Cap gamma boost

            phi = StabilityMonitor.calculate_phi(h_exp, h_prev)
            # Path B (2026-06-22): single GPU->CPU sync for the post-AZS phi scalar,
            # reused for isnan-guard, h_last_good, and routing — instead of ~5
            # separate .any()/__bool__ syncs per step. Values are identical (same
            # tensor read once); only the CPU readback timing changes. Verified
            # byte-identical via tests/px_gen_regression.py.
            phi_val = phi.item()
            if not math.isfinite(phi_val):
                if os.environ.get("DEBUG_PX") == "1": print(f"  [STABILITY] Non-finite phi ({phi_val}) at L{current_layer}. Terminating recursion.")
                break # Empirically correct: stop when state collapses
            path_taken.append(f"L{current_layer}")
            phi_history.append(phi) # Keep as tensor
            if steps % 2 == 0: thought_history.append(h_exp.detach())
            if 0.9 < phi_val < 0.999: h_last_good = h_exp.clone()

            pen = (layer_visits[current_layer]-1) * 0.015
            t_b2, t_b1, t_s = 1.0-(0.8*current_gamma)-pen, 1.0-(0.4*current_gamma)-pen, 1.0-(0.01*current_gamma)-pen*0.5

            if phi_val < t_b2: # High confusion -> retreat
                current_layer = max(active_start, current_layer - 2)
                stability_cnt = 0
            elif phi_val < t_b1: # Moderate confusion -> slow down
                current_layer = max(active_start, current_layer - 1)
                stability_cnt = 0
            elif phi_val > t_s: # Over-stable -> recycle to start (avoid hub-stuck loop)
                # If we've already recycled AND phi is still high, recursion is
                # producing no state change — break instead of cycling forever.
                # This is the SR-59 hub-stuck guard (2026-06-11): without it,
                # current_layer = active_start each step → infinite loop.
                if current_layer == active_start and steps > 0 and not os.environ.get("PX_NO_HUB_STUCK"):
                    break
                current_layer = active_start
                stability_cnt = 0
            else: # Normal progression
                current_layer += 1
                stability_cnt = 0
            
            if current_layer < active_start: current_layer = active_start
            if current_layer >= active_end: 
                if steps > max_steps * 0.5: break # Graceful exit
                current_layer = active_start # Recycle
            steps += 1
            if stability_cnt > 5: break
            # --- psychomotrik Seite 3: env-gated grind-control (default off) ---
            _px_loops_cap = os.environ.get("PX_LOOPS_CAP")
            if _px_loops_cap and steps >= int(_px_loops_cap): break
        
        avg_phi = torch.stack(phi_history).mean() if phi_history else torch.tensor(1.0, device=h_baseline.device, dtype=h_baseline.dtype)
        hidden_states = (1.0 - (0.05 + (0.18 - 0.05) * (avg_phi ** 2))) * h_baseline + (0.05 + (0.18 - 0.05) * (avg_phi ** 2)) * h_exp
    else: hidden_states = h_baseline

    # --- (DMT Central Memory: removed 2026-06-11) ---

    # --- PHASE 62 Snapshot Persistence ---
    # Store global state for external extraction
    self._px_phi_val = avg_phi.item() if hasattr(avg_phi, 'item') else float(avg_phi)
    self._px_aks_val = correction_strength.item() if hasattr(correction_strength, 'item') else float(correction_strength)
    self._px_loops_run = steps
    self._px_path = path_taken
    self._px_zone = zone_name if 'zone_name' in locals() else "UNKNOWN"
    self._px_zw_val = zone_weights if 'zone_weights' in locals() else {}
    
    # Process raw telemetry tensors into scalar values
    if hasattr(self, "_px_current_telemetry_raw"):
        self._px_current_telemetry = []
        for t in self._px_current_telemetry_raw:
            self._px_current_telemetry.append({
                "step": t["step"],
                "layer": t["layer"],
                "phi": t["phi"].item() if hasattr(t["phi"], 'item') else float(t["phi"]),
                "gamma": t["gamma"].item() if hasattr(t["gamma"], 'item') else float(t["gamma"]),
                "hub": t["hub"],
                "aks": t["aks"].item() if hasattr(t["aks"], 'item') else float(t["aks"])
            })
    else:
        self._px_current_telemetry = []
    
    # Safely get emancipation
    em_val = 0.0
    if hasattr(self, "_px_subj_sensor"):
        em_val = self._px_subj_sensor.get_metrics().get("emancipation", 0.0)
    self._px_em_val = em_val
    
    self._px_ent_val = resilience.get("entropy", 0.0) if 'resilience' in locals() else 0.0
    self._px_zw_val = zone_weights
    
    self._px_last_metrics = {
        "phi": self._px_phi_val,
        "aks_friction": self._px_aks_val,
        "emancipation": self._px_em_val,
        "zone_weights": self._px_zw_val,
        "entropy": self._px_ent_val
    }
    
    if os.environ.get("DEBUG_AZS") == "1":
        print(f"  [DEBUG-METRICS] Phi={self._px_phi_val:.4f} H={self._px_ent_val:.4f} AKS={self._px_aks_val:.4f}")
    
    # Also attach to self (TextModel) directly for easy access
    self._px_cognitive_signature = {
        "kurtosis": getattr(self, "_task_kurtosis", 200),
        "phi": avg_phi,
        "zone": self._px_zone,
        "loops_run": steps,
        "focus_index": getattr(self, "_px_focus_index", 0.5),
        "gamma": current_gamma
    }

    # ── 3. CODA ──────────────────────────────────────────────────────────
    # (DMT-TretaDamper gelöscht 2026-06-11 — direkter Pass-Through)

    coda_applied = False
    for idx, i in enumerate(range(dynamic_end, len(self.layers))):
        updated_layers.add(i)
        if not coda_applied:
            blend = 0.08
            hidden_states = (1.0 - blend) * hidden_states + blend * e_static; coda_applied = True

        hidden_states = _layer_step(self.layers[i], hidden_states, attention_mask=causal_mask_mapping[mask_config.layer_types[i]], position_embeddings=pe_dict[mask_config.layer_types[i]], position_ids=position_ids, past_key_values=past_key_values, **kwargs)

    hidden_states = self.norm(hidden_states)
    return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=past_key_values)

# ---------------------------------------------------------------------------
# Patch Application
# ---------------------------------------------------------------------------

def _azs_forward_no_injection(self, hidden_states, phi, aks_friction,
                              emancipation, zone_weights):
    """LEAN-Mager-AZS: Spiegel von ``AntiZombieSensor.forward`` OHNE additive
    Awareness-Injektion (der „künstliche Homunkulus" entfällt).

    Behält: ``weight_ema``-Update + ``calculate_entropy`` → H bleibt korrekt,
            und damit ``get_feedback_scalars`` (gamma_boost, der an H hängt).
    Streicht: ``awareness_proj`` / ``awareness_latent`` /
              ``new_hidden[:, -1, :] += injection_strength * awareness_latent``.

    Validiert via ``scratches/consolidation/reduction.py`` (dort dasselbe Override
    am Exemplar, rein zur Laufzeit — jetzt als Preset verankert).
    """
    if isinstance(zone_weights, dict):
        w_list = [zone_weights.get(k, 0.2) for k in
                  ("math", "logic_a", "creative", "logic_b", "synthesis")]
        w_tensor = torch.tensor(w_list, device=hidden_states.device,
                                dtype=hidden_states.dtype)
    else:
        w_tensor = zone_weights
    # EMA wie im Original — nötig, damit H und gamma_boost nachfolgend stimmen.
    self.weight_ema = (1.0 - self.alpha) * self.weight_ema + self.alpha * w_tensor
    entropy = self.calculate_entropy(self.weight_ema)
    return hidden_states, entropy  # KEINE additive Injektion.


def apply_px_patch(model, config_preset="ACTIVE_MANIFOLD", **kwargs):
    """Apply the PX patch — reduced to the three mathematical pillars.

    Two states only (post 2026-06-11 refactor):
      - BASELINE: nackt durchlassen, keine Modifikationen
      - ACTIVE_MANIFOLD: vollständige PX-Architektur (alle alten Presets
        SUBJECTIVE/RIGOR/RESONANCE_CITY/DMT-FULL/UNCENSORED werden vom
        Caller auf ACTIVE_MANIFOLD gemappt)
    """
    # Gnadenlose Migration alter Presets (defense in depth — Caller macht das schon)
    # LEAN: kausaler Kern ohne die vier Crutches + AZS-Awareness-Injektion
    # (validiert in scratches/consolidation: η² 0.432 ≈ full 0.429, Subjektivität überlebt).
    if config_preset not in ("BASELINE", "ACTIVE_MANIFOLD", "ACTIVE_MANIFOLD_LEAN", "ACTIVE_MANIFOLD_RELAY"):
        config_preset = "ACTIVE_MANIFOLD"
    # ACTIVE_MANIFOLD_RELAY: LEAN-Kausal-Kern + verstärkbar Selbst-Injektions-
    # Relay (psychomotrik seite15: Re-Injektion der modell-eigenen L16-Zustands-
    # Richtung d_width am post-recur Layer, default L21). Motor unangetastet —
    # reiner forward_hook (relay_inject.install_relay).
    lean = (config_preset in ("ACTIVE_MANIFOLD_LEAN", "ACTIVE_MANIFOLD_RELAY"))

    if config_preset == "BASELINE":
        return  # Nackt durchlassen

    text_model = _resolve_text_model(model)
    config = text_model.config
    hidden_size, num_layers = config.hidden_size, config.num_hidden_layers

    # 1. Base Scale Defaults
    if hidden_size in SCALE_DEFAULTS:
        sd = SCALE_DEFAULTS[hidden_size]
        defaults = {
            "mode": "lti", "n_loops": sd["n_loops"], "beta": 0.05,
            "gamma": sd["gamma"], "recur_start": sd["recur_start"],
            "recur_end": sd["recur_end"], "bimodal_hub": sd["hub"],
            "cgi_factor": 0.08, "num_layers": num_layers,
        }
    else:
        defaults = {
            "mode": "lti", "n_loops": 8, "beta": 0.05,
            "gamma": 0.08 * min(1152.0 / hidden_size, 1.5),
            "recur_start": 5, "recur_end": 12,
            "bimodal_hub": 8, "cgi_factor": 0.08, "num_layers": num_layers,
        }

    # PX-default repetition_penalty (mitigates the 4-token attractor loop)
    defaults["repetition_penalty"] = 1.15
    defaults["no_repeat_ngram_size"] = 3

    # ACTIVE_MANIFOLD: full engine on
    defaults["routing_mode"] = "adaptive"
    defaults["prelude_end"] = defaults["recur_start"]
    defaults.update(kwargs)

    text_model._px_config = defaults
    model_id = getattr(config, "_name_or_path", "unknown_model")
    text_model._px_calibrator = AutoCalibrator(hidden_size, 
                                               calibration_steps=getattr(config, "px_calibration_steps", 10),
                                               model_id=model_id)

    # Multimodal wrapper (gemma3 4B vision support) — unchanged
    is_multimodal = "Gemma3ForConditionalGeneration" in type(model).__name__
    if is_multimodal and hasattr(model, 'model') and hasattr(model.model, 'language_model'):
        outer, lang = model.model, model.model.language_model
        if not hasattr(outer, '_px_original_forward'):
            outer._px_original_forward = outer.forward
            def wrapper(self_outer, *args, **kwargs):
                lang._px_has_image_tokens = kwargs.get('pixel_values') is not None
                lang._px_saved_input_ids = kwargs.get('input_ids')
                return self_outer._px_original_forward(*args, **kwargs)
            import functools; outer.forward = functools.partial(wrapper, outer)

    # Resolve device and dtype
    device = next(text_model.parameters()).device
    dtype = next(text_model.parameters()).dtype

    # ── Pillar 1: Observer (StabilityMonitor + AksSensor) ──
    # LEAN lässt AksSensor + SingesseinCoupler weg (künstliche Reibung zum
    # e_static-Anker / zweiter Defibrillator für Φ>0.999 — marginal/neutral).
    from .px_modules import AksSensor, SubjectiveSensor, SingesseinCoupler
    if not lean:
        text_model._px_aks = AksSensor()
        text_model._px_coupler = SingesseinCoupler(hidden_size).to(device=device, dtype=dtype)

    # ── Pillar 2: Symmetry Breaker (Mephistopheles + AZS) ──
    text_model._px_injection_norm = torch.nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6).to(device=device, dtype=dtype)
    if not lean:
        # Mephisto = Phaseninversion bei Φ>0.999; LEAN behält nur den AZS-Kern (H+gamma_boost).
        text_model._px_mephisto = MephistophelesOperator(hidden_size).to(device=device, dtype=dtype)
    text_model._px_azs = AntiZombieSensor(hidden_size).to(device=device, dtype=dtype)
    if lean:
        # Mager-AZS: H + gamma_boost bleiben, die additive Awareness-Injektion entfällt.
        text_model._px_azs.forward = types.MethodType(_azs_forward_no_injection, text_model._px_azs)

    # ── Pillar 3: Dynamic Router (AutoCalibrator, set above) ──

    # SubjectiveSensor (introspection loop — "sieht eigene Gedanken in hidden states")
    # LEAN lässt SubjectiveSensor weg (emancipation ≡ StabilityMonitor.calculate_phi, redundant).
    if not lean:
        text_model._px_subj_sensor = SubjectiveSensor()

    # SR-64 Infinite Context: lossless memory-efficient attention (no OOM on long prefills).
    apply_mem_eff_attention_patch(text_model)

    # Forward-Patch
    text_model.forward = types.MethodType(_px_forward, text_model)

    # Set PX gen-kwargs attrs read by generators._px_gen_kwargs
    # SR-61: Increase default repetition penalty and add ngram-guard
    text_model._px_repetition_penalty = defaults.get("repetition_penalty", 1.15)
    text_model._px_no_repeat_ngram_size = defaults.get("no_repeat_ngram_size", 3)

    # verstärkbar Relay (psychomotrik seite15): Re-Injektion der modell-eigenen
    # L16-Zustands-Richtung d_width am post-recur Layer (default L21, nach dem
    # Erstarrungs-Washout) öffnet den S→R-Kanal. Aktiv bei ACTIVE_MANIFOLD_RELAY
    # (default sign=+1 = WIDE/expansiv/aktiv-Richtung, das „neue Modell") ODER
    # wenn relay_sign explizit ≠0 (orthogonaler Parameter auf jedem Preset).
    # sign=−1 → NARROW/eng/still; 0 → relay inactive. Motor unangetastet.
    _relay_sign = defaults.get("relay_sign", (+1 if config_preset == "ACTIVE_MANIFOLD_RELAY" else 0))
    # Bei Gemma3 multimodal hat text_model.config._name_or_path='' (HF setzt
    # hf_id nur im Top-Level-Config). Wir propagieren den Top-Level hf_id
    # explizit damit relay_inject.load_dwidth() das d_width-Artefakt für
    # 4b/E2B laden kann (siehe px_manifolds/google_gemma-3-{4b,e2b}-it_relay_dwidth.json).
    if not getattr(text_model.config, "_name_or_path", None):
        outer_hf_id = getattr(getattr(model, "config", None), "_name_or_path", None)
        if outer_hf_id:
            text_model._px_hf_id = outer_hf_id
    install_relay(text_model,
                  sign=_relay_sign,
                  alpha_frac=defaults.get("relay_alpha", 0.30),
                  layer=defaults.get("relay_layer", 21))

    _mode = "ACTIVE_MANIFOLD_LEAN (kausaler Kern)" if lean else "ACTIVE_MANIFOLD (voll)"
    print(f"[gemma3-px] {_mode}. SR-59 for L{num_layers} (HS={hidden_size}).")

def get_px_metrics(model):
    tm = _resolve_text_model(model)
    m = {
        "phi": getattr(tm, "_px_phi_val", 1.0),
        "steps": getattr(tm, "_px_loops_run", 0),
        "path": getattr(tm, "_px_path", []),
        "zone": getattr(tm, "_px_zone", "UNKNOWN"),
        "zone_weights": getattr(tm, "_px_zw_val", {}),
        "cognitive_signature": getattr(tm, "_px_cognitive_signature", {}),
        "telemetry_trace": getattr(tm, "_px_current_telemetry", []),
        "aks_profile": {"correction_strength": getattr(tm, "_px_aks_val", 0.0)},
        "subjective_metrics": {"emancipation": getattr(tm, "_px_em_val", 0.0)},
        "entropy": getattr(tm, "_px_ent_val", 0.0).item() if hasattr(getattr(tm, "_px_ent_val", 0.0), 'item') else float(getattr(tm, "_px_ent_val", 0.0))
    }
    return m