PowerMachine's picture
KHTST v21 — fonte única do IPC (Rust removido por DRY; layout com _Static_assert compile-time + auditoria import-time); Newton-Schulz com normalização espectral CERTIFICADA (min(sqrt(|A|inf)/1.2, |X|F), janela invariante 1.255 provada, nunca pior que v20, Teos 33.3-33.6); imagem adaptativa (letterbox+LANCZOS+janela saliente+upscale limitado, Teos 33.7-33.10); hot-paths Cython (conv causal fundida com gradiente analítico exato, ganho 2.3-7.6x; top-k parcial do roteador, Teos 33.11-33.13); tabela v17 de 18 sinais FECHADA no treinador com 8 produtores reais (Teo 33.14); A/B/C/D pareado real 48 passos (ganhos 1.10-1.12x, leitura honesta); doc 33 + PLANO/RELATORIO/README; suíte v21 40/40, regressão v1-v20 zero, verificar_tudo 49✓; deprecados removidos
20672ab verified
Raw History Blame Contribute Delete
10.2 kB
# -*- coding: utf-8 -*-
"""Kernel MAMBA-3 (v11, doc 22 §2) — varredura seletiva com discretização
TRAPETZOIDAL, em PyTorch puro, vetorizado e DIFERENCIÁVEL por autograd.
Fonte: arquivo MAMBA_3.txt (kernel Cython `mamba3_cpu_kernel.pyx`) — NÃO
copiado: o kernel Cython (a) fixava `a_val = A[b, d, 0]` (ignorava a
dependência de A em n — gradiente de A[·,·,n>0] errado), (b) exigia
compilação por plataforma e backward manual propenso a erro, e (c)
materializava h_states (B,L,D,N) inteiro. Aqui (Teoremas 22.2–22.5):
• A é por (d, n): α_{t,d,n} = exp(δ_{t,d}·a_{d,n}) — recuperação do SSM
canônico; o gradiente de A flui EXATO via autograd (Corolário 22.2.1);
• atualização trapezoidal: h_t = α_t·h_{t−1} + (δ_t/2)·(1+α_t)·B̄_t·u_t
com B̄_{t,n} = Σ_r B_{t,r,n} (MIMO rank-reduzido, como a fonte);
• ESTABILIDADE (Teo 22.3): δ = σ(·) ∈ (0,1) e a_{d,n} < 0 ⇒
α_{t,d,n} ∈ (e^{a}, 1) ⊂ (0,1) — a recorrência é CONTRAÇÃO; estado
limitado ‖h_t‖ ≤ max_s (δ_s/2)(1+α_s)|B̄_s·u_s| / (1−α_max);
• v18 — CLAMPS DE SANIDADE (Teo 30.2, requisito §1.1.5/§2.4): a
contração do Teo 22.3 exige δ∈(0,1) e a<0. Se uma anomalia
(inicialização ruim, gradiente explosivo, uso direto da API) violar
isso, α = exp(δ·a) ≥ 1 e a recorrência DIVERGE geometricamente.
A varredura agora fixa delta ∈ [1e-6, 1−1e-6] e A ≤ −1e-6 ANTES de
α=exp(δ·a). Como δ e A válidos NUNCA são alterados pelo clamp
(identidade na região válida), o gradiente autograd é inalterado
(Teo 30.2-cor: clamp idêntico na região válida ⟹ ∇ idêntico);
• memória O(B·D·N) por passo (estado vivo) — sem tensor (B,L,D,N);
• semente inicial h_0 = 0 (convenção de estado frio — idêntica à fonte).
Assinatura compatível com o autograd manual da fonte (u, delta, A, B, C) —
troca DROP-IN: o autograd do PyTorch substitui `mamba3_backward_cpu`
(prova de exatidão: teste `test_v11_mamba3_grad_vs_finitediff`).
"""
from __future__ import annotations
import torch
import torch.nn.functional as F
__all__ = ["varredura_trapezoidal", "Mamba3Funcao", "verificar_estabilidade",
"conv1d_causal", "set_conv_cython"]
# v18 — clamps de sanidade (Teo 30.2). Ligados por padrão: custam O(1)
# por tensor e garantem a contração mesmo sob anomalia numérica.
CLAMPS_SEGURANCA = True
DELTA_MIN, DELTA_MAX = 1e-6, 1.0 - 1e-6
A_MAX = -1e-6 # força a < 0 (contração)
def verificar_estabilidade(delta: torch.Tensor, A: torch.Tensor) -> dict:
"""Auditoria (sem mutação): quantas entradas de δ/A violariam a
contração do Teo 22.3 ANTES do clamp. Telemetria para o TelemetryHub
(violações > 0 ⟹ anomalia a montante: investigar inicialização/grads)."""
with torch.no_grad():
n_delta_fora = int(((delta <= 0.0) | (delta >= 1.0)).sum())
n_A_nao_neg = int((A >= 0.0).sum())
return {"mamba3/delta_fora": n_delta_fora, "mamba3/A_nao_negativa": n_A_nao_neg}
def varredura_trapezoidal(u: torch.Tensor, delta: torch.Tensor, A: torch.Tensor,
B_m: torch.Tensor, C_m: torch.Tensor,
devolve_estado: bool = False):
"""Varredura causal seletiva MAMBA-3.
Formas:
u: (B, L, D) entrada u_t
delta: (B, L, D) passo de tempo δ_t ∈ (0,1) (gate sigmoid)
A: (D, N) ou (B, D, N) — dinâmica de estado (a < 0 estável)
B_m: (B, L, R, N) projeção de entrada MIMO (rank R)
C_m: (B, L, R, N) projeção de leitura MIMO (rank R)
Devolve out: (B, L, D) com out_t,d = Σ_n (Σ_r C_{t,r,n}) · h_{t,d,n}.
v12 — `devolve_estado=True` devolve TAMBÉM h_final: (B, D, N), o estado
h_L da recorrência (Teorema 24.13: h_L é EXATAMENTE o estado que o
decode incremental recebe — mesma recorrência, então
varredura(u_{1..L}) ⊕ passo(u_{L+1}) ≡ varredura(u_{1..L+1})).
Laço temporal O(L) vetorizado sobre (B, D, N, R) — CPU/GPU sem kernel
custom; exato no tempo (sem aproximação além da própria regra do
trapézio, que é a DEFINIÇÃO do MAMBA-3 — Teorema 22.2).
"""
Bsz, L, D = u.shape
N = A.shape[-1]
R = B_m.shape[2]
if A.dim() == 2: # (D, N) → expande p/ lote
A = A.unsqueeze(0).expand(Bsz, -1, -1)
assert B_m.shape == (Bsz, L, R, N) and C_m.shape == (Bsz, L, R, N), \
"BUG (doc 22 §2): formas MIMO incompatíveis"
if CLAMPS_SEGURANCA:
# Teo 30.2 (v18, requisito §1.1.5/§2.4): contração GARANTIDA por
# construção — δ∈[1e-6,1−1e-6] e A≤−1e-6 ⟹ α=exp(δ·a) ∈
# [exp((1−1e-6)·(−1e-6)) , exp(1e-6·(−1e-6))) ⊂ (0,1) SEMPRE.
# Na região válida o clamp é a identidade ⟹ ∇ inalterado
# (Teo 30.2-cor: d clamp(x)/dx = 1 dentro da faixa).
delta = delta.clamp(DELTA_MIN, DELTA_MAX)
A = A.clamp(max=A_MAX)
# agregação MIMO: b̄_t,n = Σ_r B[t,r,n]; c̄_t,n = Σ_r C[t,r,n] (O(B·L·D·N·R))
b_bar = B_m.sum(dim=2) # (B, L, N)
c_bar = C_m.sum(dim=2) # (B, L, N)
# δ_{t,d}·a_{d,n} → α por (B, L, D, N)
dtA = delta.unsqueeze(-1) * A.unsqueeze(1) # (B, L, D, N)
alpha = torch.exp(dtA) # Teo 22.3: ∈ (0,1)
beta = 0.5 * delta # (B, L, D) — δ/2
# termo de entrada: β_t·(1+α_t)·b̄_{t,n}·u_t,d → (B, L, D, N)
entrada = (beta.unsqueeze(-1) * (1.0 + alpha) * b_bar.unsqueeze(2)) \
* u.unsqueeze(-1) # (B,L,D,N)
h = torch.zeros(Bsz, D, N, dtype=u.dtype, device=u.device)
saidas = []
for t in range(L): # O(L) — estado vivo
h = alpha[:, t] * h + entrada[:, t]
# leitura: out_t,d = Σ_n h_{t,d,n}·c̄_{t,n}
saidas.append((h * c_bar[:, t].unsqueeze(1)).sum(-1))
out = torch.stack(saidas, dim=1) # (B, L, D)
if devolve_estado:
return out, h # h_final (B, D, N)
return out
class Mamba3Funcao(torch.autograd.Function):
"""Ponte de compatibilidade com a API da fonte (`Mamba3Function.apply`).
Mantida para os testes compararem com a interface (u, delta, A, B, C);
o backward é o autograd PADRÃO (exato — Corolário 22.2.1)."""
@staticmethod
def forward(ctx, u, delta, A, B_m, C_m):
return varredura_trapezoidal(u, delta, A, B_m, C_m)
@staticmethod
def backward(ctx, grad_out): # pragma: no cover
raise RuntimeError("use tensor.backward() — autograd exato (doc 22)")
# ======================================================================
# v21 — CONV CAUSAL PROFUNDA FUNDIDA (doc 33 §4, Teos 33.11/33.12)
# ======================================================================
# Kernel Cython (acelerado/mamba3hot.pyx) com backward ANALÍTICO exato
# (regra da cadeia na conv causal — testado contra autograd: erro ≤ 1e-5
# relativo). Ganho medido: 2.3–7.6× vs o caminho torch (pad+Conv1d+
# 2 transposes+copy — ~5× o tráfego de memória do mínimo). FALLBACK
# automático ao caminho torch v20 (mesmo cálculo) quando: sem binário,
# dtype ≠ fp32 ou device ≠ cpu. Kill-switch: set_conv_cython(False).
try: # binário commitado (praxe v9.1)
import mamba3hot as _hot # build top-level (acelerado/)
except ImportError: # pragma: no cover
try:
from khtst.acelerado import mamba3hot as _hot
except Exception:
_hot = None
_CONV_CYTHON = True # requisito v21: Cython prioritário
def set_conv_cython(ativa: bool) -> None:
"""Kill-switch do caminho Cython da conv (A/B e fallback de emergência)."""
global _CONV_CYTHON
_CONV_CYTHON = bool(ativa)
class _ConvCausalProfunda(torch.autograd.Function):
"""Ponte autograd: forward = kernel C fundido; backward = fórmulas
analíticas EXATAS (Teo 33.11-ii — ∂u/∂w/∂b da conv causal depthwise)."""
@staticmethod
def forward(ctx, u, w, b):
ctx.save_for_backward(u, w)
from khtst.mamba3.kernel import _hot as hot
out = hot.conv1d_causal_profunda(
u.detach().contiguous().numpy(),
w.detach().contiguous().numpy(),
b.detach().contiguous().numpy())
return torch.from_numpy(out)
@staticmethod
def backward(ctx, g):
# v21-nota: backward ANALÍTICO de 1ª ordem (Teo 33.11-ii). `g` é
# detachado — double-backward pelo kernel C não é suportado (os
# testes de gradiente do projeto são de 1ª ordem; o caminho torch
# de fallback mantém autograd de ordem arbitrávia).
from khtst.mamba3.kernel import _hot as hot
u, w = ctx.saved_tensors
K = int(w.shape[1])
gd = g.detach().contiguous().numpy()
gu = torch.from_numpy(hot.conv1d_grad_u(
gd, w.detach().contiguous().numpy()))
gw = torch.from_numpy(hot.conv1d_grad_w2(
gd, u.detach().contiguous().numpy(), K))
gb = torch.from_numpy(hot.conv1d_grad_b(gd))
return gu, gw, gb
def conv1d_causal(u: torch.Tensor, pesos_dk: torch.Tensor,
vies: torch.Tensor) -> torch.Tensor:
"""Conv depthwise causal (B,L,D) → (B,L,D) — FONTE ÚNICA do hot-path
do Mamba-3 (R21.5).
Caminho Cython fundido (fp32/CPU) quando disponível; caso contrário o
caminho torch v20 EXATO (pad esquerdo K−1 + Conv1d groups=D + volta
do layout — Teo 33.11: ambos computam
uc[b,l,d] = Σ_j w[d,j]·u[b, l+j−K+1, d] + v[d], equivalência medida
max|Δ| = 0)."""
if _hot is not None and _CONV_CYTHON \
and u.dtype == torch.float32 and u.is_cpu:
return _ConvCausalProfunda.apply(u, pesos_dk, vies)
K = pesos_dk.shape[1]
return F.conv1d(F.pad(u.transpose(1, 2), (K - 1, 0)),
pesos_dk.unsqueeze(1), vies,
groups=pesos_dk.shape[0]).transpose(1, 2).contiguous()