Update app.py

app.py

@@ -1,6 +1,21 @@
 # ================================================================
 # MTP - app.py para Hugging Face Space (Gradio, CPU)
 # Carga el checkpoint MTP_MODEL.pt desde el repo TeszenAI/MTP-1
+#
+# OPTIMIZACIÓN DE VELOCIDAD (sin tocar arquitectura ni pesos):
+#   - KV-cache en la atención: en generación autoregresiva, cada paso
+#     antes recomputaba TODO el contexto desde cero (O(n^2) en total).
+#     Ahora se reutiliza lo ya calculado y solo se procesa el token
+#     nuevo (O(n) en total). Es el mismo cálculo matemático, solo que
+#     no se repite trabajo ya hecho.
+#   - F.scaled_dot_product_attention: kernel fusionado de PyTorch,
+#     mismo resultado que el softmax manual pero más rápido en CPU.
+#     Si la versión de PyTorch no lo trae, cae automáticamente al
+#     cálculo manual (fallback), así que no se rompe en ningún entorno.
+#   - repetition_penalty vectorizado (sin bucle Python + set() por token).
+#
+# El modelo, los pesos, el muestreo (top_k/top_p/temperature/repetition)
+# y las respuestas de la API/Gradio son EXACTAMENTE los mismos que antes.
 # ================================================================
 import os
 import math
@@ -17,10 +32,22 @@ from huggingface_hub import hf_hub_download
 # ---------------- Optimización para CPU ----------------
 # Limita hilos a los núcleos disponibles (evita overhead en Spaces pequeños)
 torch.set_num_threads(max(1, os.cpu_count() or 1))
+
+# set_num_interop_threads solo puede llamarse una vez y antes de cualquier
+# operación paralela; lo protegemos por si el entorno ya lo fijó.
+try:
+    torch.set_num_interop_threads(1)
+except RuntimeError:
+    pass
+
 torch.set_grad_enabled(False)  # solo inferencia, nunca necesitamos gradientes
 
 DEVICE = "cpu"
 
+# Disponibilidad de scaled_dot_product_attention (PyTorch >= 2.0).
+# Si no está disponible, usamos el softmax manual original como fallback.
+_HAS_SDPA = hasattr(F, "scaled_dot_product_attention")
+
 REPO_ID = "TeszenAI/MTP-1.2"
 FILENAME = "MTP_MODEL.pt"
 
@@ -35,21 +62,52 @@ class CausalSelfAttention(nn.Module):
         self.attn_dropout = nn.Dropout(dropout)
         self.resid_dropout = nn.Dropout(dropout)
         mask = torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size)
+        # Se mantiene el buffer para que el state_dict del checkpoint cargue
+        # igual que antes (la clave "attn.mask" existe en el checkpoint).
+        # Ya no se usa en el forward optimizado con SDPA; solo lo usa el
+        # fallback manual si SDPA no está disponible.
         self.register_buffer("mask", mask)
 
-    def forward(self, x):
+    def forward(self, x, past_kv=None, use_cache=False):
         B, T, C = x.shape
         qkv = self.qkv(x)
         q, k, v = qkv.split(C, dim=2)
         q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
         k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
         v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
-        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
-        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
-        att = F.softmax(att, dim=-1)
-        att = self.attn_dropout(att)
-        out = (att @ v).transpose(1, 2).contiguous().view(B, T, C)
-        return self.resid_dropout(self.proj(out))
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+
+        present_kv = (k, v) if use_cache else None
+
+        # Causal solo hace falta cuando hay varias queries nuevas sin pasado
+        # (prefill del prompt). En un paso de decodificación (T=1 con caché)
+        # el único token nuevo ya puede ver todo el pasado sin máscara.
+        is_causal = (past_kv is None) and (T > 1)
+
+        if _HAS_SDPA:
+            out = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=None,
+                dropout_p=0.0,  # en eval() el dropout original no hace nada
+                is_causal=is_causal,
+            )
+        else:
+            att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+            if is_causal:
+                Tk = k.size(-2)
+                causal_mask = torch.tril(torch.ones(T, Tk, device=x.device, dtype=torch.bool))
+                att = att.masked_fill(~causal_mask, float("-inf"))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            out = att @ v
+
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        out = self.resid_dropout(self.proj(out))
+        return out, present_kv
 
 
 class FeedForward(nn.Module):
@@ -72,10 +130,11 @@ class Block(nn.Module):
         self.ln2 = nn.LayerNorm(n_embd)
         self.ff = FeedForward(n_embd, dropout)
 
-    def forward(self, x):
-        x = x + self.attn(self.ln1(x))
+    def forward(self, x, past_kv=None, use_cache=False):
+        attn_out, present_kv = self.attn(self.ln1(x), past_kv=past_kv, use_cache=use_cache)
+        x = x + attn_out
         x = x + self.ff(self.ln2(x))
-        return x
+        return x, present_kv
 
 
 class MTP(nn.Module):
@@ -90,15 +149,22 @@ class MTP(nn.Module):
         self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
         self.lm_head.weight = self.tok_emb.weight
 
-    def forward(self, idx):
+    def forward(self, idx, past_key_values=None, use_cache=False, pos_offset=0):
         B, T = idx.shape
-        pos = torch.arange(T, device=idx.device)
+        pos = torch.arange(pos_offset, pos_offset + T, device=idx.device)
         x = self.tok_emb(idx) + self.pos_emb(pos)
         x = self.drop(x)
-        for block in self.blocks:
-            x = block(x)
+
+        new_past = [] if use_cache else None
+        for i, block in enumerate(self.blocks):
+            past_kv = past_key_values[i] if past_key_values is not None else None
+            x, present_kv = block(x, past_kv=past_kv, use_cache=use_cache)
+            if use_cache:
+                new_past.append(present_kv)
+
         x = self.ln_f(x)
-        return self.lm_head(x)
+        logits = self.lm_head(x)
+        return logits, new_past
 
 
 # ---------------- Carga del checkpoint (una sola vez, al iniciar el Space) ----------------
@@ -122,12 +188,11 @@ model = MTP(
 model.load_state_dict(checkpoint["model_state_dict"])
 model.eval()
 
-# fusiona LayerNorm/Linear estáticamente no aplica aquí, pero fija modo eval
-# y evita cualquier dropout durante inferencia.
 BLOCK_SIZE = cfg["block_size"]
 
 print(f"MTP cargado ({checkpoint['meta']['model_name']}, "
-      f"entrenado con {checkpoint['meta']['trained_examples']} ejemplos)")
+      f"entrenado con {checkpoint['meta']['trained_examples']} ejemplos)"
+      f" | SDPA={'sí' if _HAS_SDPA else 'no (fallback manual)'}")
 
 
 def encode_text(s):
@@ -138,17 +203,48 @@ def decode_ids(ids):
     return "".join(itos.get(i, "") for i in ids if i not in (PAD_ID, BOS_ID, EOS_ID))
 
 
-# ---------------- Generación ----------------
+# ---------------- Generación (con KV-cache) ----------------
 @torch.inference_mode()
 def generate(idx, max_new_tokens, temperature, top_k, top_p, repetition_penalty):
+    past_key_values = None
+    cache_len = 0  # cuántos tokens del extremo derecho de `idx` ya están en la caché
+
     for _ in range(max_new_tokens):
-        idx_cond = idx[:, -BLOCK_SIZE:]
-        logits = model(idx_cond)
-        logits = logits[:, -1, :] / max(temperature, 1e-5)
+        total_len = idx.shape[1]
+
+        if total_len <= BLOCK_SIZE:
+            if past_key_values is None:
+                # Primer paso: una sola pasada ("prefill") sobre todo el prompt.
+                logits, past_key_values = model(idx, use_cache=True)
+                cache_len = total_len
+            else:
+                # Pasos siguientes: solo se procesa el último token generado,
+                # reutilizando la caché de todo lo anterior.
+                last_token = idx[:, -1:]
+                logits, past_key_values = model(
+                    last_token,
+                    past_key_values=past_key_values,
+                    use_cache=True,
+                    pos_offset=cache_len,
+                )
+                cache_len += 1
+            logits = logits[:, -1, :]
+        else:
+            # Se superó block_size: mismo comportamiento que el modelo original
+            # (ventana deslizante recalculada por completo). Solo ocurre en
+            # respuestas muy largas; la caché se reinicia para esa ventana.
+            idx_cond = idx[:, -BLOCK_SIZE:]
+            logits, past_key_values = model(idx_cond, use_cache=True)
+            cache_len = BLOCK_SIZE
+            logits = logits[:, -1, :]
+
+        logits = logits / max(temperature, 1e-5)
 
         if repetition_penalty and repetition_penalty != 1.0:
-            for token_id in set(idx[0].tolist()):
-                logits[0, token_id] /= repetition_penalty
+            # Vectorizado: antes era `for token_id in set(idx[0].tolist())`,
+            # un bucle Python nuevo por cada token generado.
+            unique_ids = torch.unique(idx[0])
+            logits[0, unique_ids] /= repetition_penalty
 
         if top_k is not None and top_k > 0:
             v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
@@ -171,6 +267,7 @@ def generate(idx, max_new_tokens, temperature, top_k, top_p, repetition_penalty)
         idx = torch.cat([idx, next_id], dim=1)
         if next_id.item() == EOS_ID:
             break
+
     return idx
 
 
@@ -178,7 +275,7 @@ def run_inference(text, max_new_tokens=None, temperature=None, top_k=None, top_p
     """Núcleo de generación, reutilizado por la UI de Gradio y por la API /generate.
     No reduce calidad por estar en CPU: usa exactamente el mismo muestreo
     (top_k + top_p + repetition_penalty) que en la Celda 2 de entrenamiento,
-    solo que tarda más en devolver el resultado."""
+    solo que ahora con KV-cache es notablemente más rápido en respuestas largas."""
     max_new_tokens = int(max_new_tokens) if max_new_tokens else gen_defaults["max_new_tokens"]
     temperature = float(temperature) if temperature is not None else gen_defaults["temperature"]
     top_k = int(top_k) if top_k is not None else gen_defaults["top_k"]