Fit 0.5b training on 32GB: SDPA MLA, block checkpoint, chunked CE

Whole-mixer checkpoint plus T×T MLA scores OOM'd a 31GB GPU on backward.
Checkpoint each AttnRes block, run absorbed MLA through SDPA, and compute
CE in vocab chunks so [B,T,V] logits are never materialized.

--max-tokens is now the training budget; default --steps 2000 no longer
caps a 1B-token run at 250 optimizer steps.

kda/layers/attn_res.py

@@ -20,6 +20,7 @@ import torch
 import torch.nn.functional as F
 from einops import rearrange
 from torch import Tensor, nn
+from torch.utils.checkpoint import checkpoint as activation_checkpoint
 
 
 ATTNRES_MODES = ("off", "full", "block")
@@ -247,6 +248,7 @@ class BlockAttnResStack(nn.Module):
             if is_final_aggregate
             else None
         )
+        self.gradient_checkpointing = False
 
     def forward_naive(self, x: Tensor) -> Tensor:
         blocks = [x]  # b_0=embedding/input representation
@@ -294,9 +296,20 @@ class BlockAttnResStack(nn.Module):
         blocks = [x]
         depth = len(self.layers)
         start = 0
+        use_ckpt = (
+            self.gradient_checkpointing and self.training and torch.is_grad_enabled()
+        )
         while start < depth:
             end = min(start + self.block_size, depth)
-            blocks.append(self._run_block_two_phase(blocks, start, end))
+            if use_ckpt:
+                def _run(*srcs, _start=start, _end=end):
+                    return self._run_block_two_phase(list(srcs), _start, _end)
+
+                blocks.append(
+                    activation_checkpoint(_run, *blocks, use_reentrant=False)
+                )
+            else:
+                blocks.append(self._run_block_two_phase(blocks, start, end))
             start = end
 
         return (

kda/layers/mla.py

@@ -78,20 +78,14 @@ class GatedMLA(nn.Module):
         w_uk = w[: H * self.qk_nope_head_dim].view(H, self.qk_nope_head_dim, r)
         w_uv = w[H * self.qk_nope_head_dim :].view(H, self.v_head_dim, r)
 
-        # 吸收 W_UK 进 query: score = (q @ W_UK^T) @ c^T
+        # 吸收 W_UK 进 query: score = (q @ W_UK^T) @ c^T, scale=1 matches the unscaled einsum.
         q_absorb = torch.einsum("bthd,hdj->bthj", q, w_uk)      # [B, T, H, r]
-        scores = torch.einsum("bthj,bsj->bhts", q_absorb, c)    # [B, H, T, T]
-
-        mask = torch.triu(
-            torch.ones(T, T, dtype=torch.bool, device=x.device), diagonal=1
-        )
-        scores = scores.masked_fill(mask, float("-inf"))
-        attn = F.softmax(scores, dim=-1)                        # [B, H, T, T]
-
-        # 先在 latent 加权, 再乘 W_UV^T 还原 v —— 永不解压
-        latent_out = torch.einsum("bhts,bsj->bhtj", attn, c)    # [B, H, T, r]
-        o_heads = torch.einsum("bhtj,hvj->bhtv", latent_out, w_uv)  # [B, H, T, d_v]
-
+        q_h = q_absorb.transpose(1, 2)                          # [B, H, T, r]
+        kv = c.unsqueeze(1).expand(B, H, T, r)
+        latent_out = F.scaled_dot_product_attention(
+            q_h, kv, kv, is_causal=True, scale=1.0
+        )                                                       # [B, H, T, r]
+        o_heads = torch.einsum("bhtr,hvr->bhtv", latent_out, w_uv)
         o_heads = o_heads.transpose(1, 2).reshape(B, T, H * self.v_head_dim)
         gate = torch.sigmoid(self.gate(x))                      # [B, T, H*d_v]
         return self.o_proj(gate * o_heads)                      # [B, T, d]