Upload 11 files

MonoidForCausalLM.py

@@ -14,11 +14,13 @@ Architecture / 架构概要:
 
     Monoid attention compresses the entire causal history into a
     fixed-size state matrix S_t ∈ ℝ^{d×d} per head:
-      S_t = α_t · S_{t-1} + k_t ⊗ v_t       (explicit causal recurrence)
-      o_t = q_t · S_t                          (state readout)
+      S_t = diag(α_t) · S_{t-1} + k_t ⊗ v_t   (vector decay recurrence)
+      o_t = q_t · S_t                            (state readout)
+    where α_t ∈ ℝ^d is a per-dimension vector decay gate.
     幺半群注意力将完整因果历史压缩到每个头一个固定大小的状态矩阵 S_t:
-      S_t = α_t · S_{t-1} + k_t ⊗ v_t         (显式因果递推)
-      o_t = q_t · S_t                          (状态读出)
+      S_t = diag(α_t) · S_{t-1} + k_t ⊗ v_t     (向量衰减递推)
+      o_t = q_t · S_t                              (状态读出)
+    其中 α_t ∈ ℝ^d 是逐维度的向量衰减门。
 
     This is a monoid because the binary operator:
       (log_α, S) ⊕ (log_β, X) = (log_α + log_β, exp(log_β)·S + X)
@@ -69,12 +71,12 @@ except ImportError:
     # Slower than the fused CUDA kernel but numerically identical.
 
     def parallel_scan(log_alpha: Tensor, kv: Tensor) -> Tensor:
-        """Sequential prefix scan fallback: S_t = exp(log_α_t)·S_{t-1} + kv_t."""
+        """Sequential prefix scan fallback: S_t[i,:] = exp(log_α_t[i])·S_{t-1}[i,:] + kv_t[i,:]."""
         B, H, T, d1, d2 = kv.shape
         states = torch.zeros(B, H, T, d1, d2, device=kv.device, dtype=kv.dtype)
         S = torch.zeros(B, H, d1, d2, device=kv.device, dtype=kv.dtype)
         for t in range(T):
-            decay = torch.exp(log_alpha[:, :, t])                    # [B, H, 1]
+            decay = torch.exp(log_alpha[:, :, t])                    # [B, H, d]
             while decay.dim() < S.dim():
                 decay = decay.unsqueeze(-1)
             S = S * decay + kv[:, :, t]
@@ -86,7 +88,7 @@ except ImportError:
         B, H, T, d1, d2 = kv.shape
         states = torch.zeros(B, H, T, d1, d2, device=kv.device, dtype=kv.dtype)
         S = torch.zeros(B, H, d1, d2, device=kv.device, dtype=kv.dtype)
-        log_acc = torch.zeros(B, H, 1, device=log_alpha.device, dtype=log_alpha.dtype)
+        log_acc = torch.zeros(B, H, d1, device=log_alpha.device, dtype=log_alpha.dtype)
         for t in range(T):
             decay = torch.exp(log_alpha[:, :, t])
             while decay.dim() < S.dim():
@@ -217,11 +219,12 @@ def monoid_op(
     b: tuple[Tensor, Tensor],
 ) -> tuple[Tensor, Tensor]:
     """
-    The monoid binary operator ⊕ on (log-space decay, state matrix) pairs.
-    幺半群二元算子 ⊕，作用于 (对数衰减, 状态矩阵) 对。
+    The monoid binary operator ⊕ on (log-space vector decay, state matrix) pairs.
+    幺半群二元算子 ⊕，作用于 (对数向量衰减, 状态矩阵) 对。
 
     Definition / 定义:
-      (log_α, S) ⊕ (log_β, X) = (log_α + log_β, exp(log_β)·S + X)
+      (log_α, S) ⊕ (log_β, X) = (log_α + log_β, diag(exp(log_β))·S + X)
+      where log_α, log_β ∈ ℝ^d are per-dimension log decay vectors.
 
     Why this is a monoid / 为什么这是幺半群:
       • Associativity / 结合律:
@@ -326,15 +329,19 @@ class MonoidAttention(nn.Module):
 
         # --- Decay gate (novel component, randomly initialized) ---
         # --- 衰减门 (全新组件, 随机初始化) ---
-        # Projects hidden_size → num_heads, yielding one scalar α per head.
-        # After sigmoid: α_t ∈ (0,1) controls per-head forgetting rate.
-        # This is the key to *explicit causal modeling*: the model learns
-        # a content-dependent decay, not a fixed positional bias.
-        # 将 hidden_size 投影到 num_heads, 每个头产生一个标量 α。
-        # 经过 sigmoid 后: α_t ∈ (0,1) 控制每个头的遗忘速率。
-        # 这是 *显式因果建模* 的关键: 模型学习的是内容相关的衰减,
-        # 而非固定的位置偏置。
-        self.decay_proj = nn.Linear(config.hidden_size, self.num_heads, bias=True)
+        # Projects hidden_size → num_heads * head_dim, yielding a VECTOR per head.
+        # Activation: log_α = -softplus(Wx + b), giving α ∈ (0, 1].
+        # Vector decay: S_t = diag(α_t) · S_{t-1} + k_t ⊗ v_t
+        # Different feature dimensions can have independent lifetimes:
+        #   - fast-decaying dims for local syntax
+        #   - slow-decaying dims for global entity/fact memory
+        # 将 hidden_size 投影到 num_heads * head_dim, 每个头产生一个向量。
+        # 激活: log_α = -softplus(Wx + b), 使 α ∈ (0, 1]。
+        # 向量衰减: S_t = diag(α_t) · S_{t-1} + k_t ⊗ v_t
+        # 不同特征维度拥有独立的生命周期:
+        #   - 快速衰减的维度负责局部语法结构
+        #   - 慢速衰减的维度负责全局实体和事实记忆
+        self.decay_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=True)
 
         # --- QK-Norm (novel component, randomly initialized) ---
         # --- QK 归一化 (全新组件, 随机初始化) ---
@@ -358,12 +365,17 @@ class MonoidAttention(nn.Module):
     def forward(
         self,
         hidden_states: Tensor,
+        attention_mask: Tensor | None = None,
         monoid_cache: MonoidCache | None = None,
         use_cache: bool = False,
     ) -> tuple[Tensor, tuple[Tensor, Tensor] | None]:
         """
         Args:
             hidden_states: [B, T, hidden_size]
+            attention_mask: [B, T] with 1=real token, 0=pad.
+                            For PAD positions: α=1 (preserve state), kv=0 (no contribution).
+                            掩码: 1=真实token, 0=填充。
+                            填充位置: α=1 (保持状态不变), kv=0 (无贡献)。
             monoid_cache: O(1) state cache for inference
                           推理用 O(1) 状态缓存
             use_cache: whether to use/update the cache
@@ -399,13 +411,40 @@ class MonoidAttention(nn.Module):
         # PSD 保证信息单调积累。
         k = torch.nn.functional.silu(k)
 
-        # --- Compute per-head decay gate α_t ---
-        # --- 计算每头衰减门 α_t ---
-        # sigmoid ensures α ∈ (0,1), then log-space for numerical stability.
-        # sigmoid 保证 α ∈ (0,1), 然后转到对数空间保证数值稳定性。
-        alpha = torch.sigmoid(self.decay_proj(hidden_states))       # [B,T,H]
-        alpha = alpha.transpose(1, 2).unsqueeze(-1)                 # [B,H,T,1]
-        log_alpha = torch.log(alpha.clamp(min=1e-6))
+        # --- Compute per-dimension vector decay gate α_t ---
+        # --- 计算每维度向量衰减门 α_t ---
+        # Negative Softplus: log_α = -softplus(Wx + b)
+        # Value range: log_α ∈ (-∞, 0), i.e. α ∈ (0, 1].
+        # When Wx → -∞: softplus → 0, α → 1 (perfect memory, no forgetting)
+        # When Wx → +∞: softplus → Wx, α → 0 (complete forgetting)
+        # This avoids α > 1 explosion (unlike SiLU) while still allowing
+        # α = 1 for lossless memory (unlike Sigmoid which caps at <1).
+        # Each dimension of the d-vector decays independently:
+        #   S_t[i,j] = α_t[i] · S_{t-1}[i,j] + k_t[i] · v_t[j]
+        #
+        # 负 Softplus: log_α = -softplus(Wx + b)
+        # 值域: log_α ∈ (-∞, 0), 即 α ∈ (0, 1]。
+        # 当 Wx → -∞: softplus → 0, α → 1 (完美记忆, 不遗忘)
+        # 当 Wx → +∞: softplus → Wx, α → 0 (完全遗忘)
+        # 避免了 SiLU 的 α > 1 爆炸, 同时允许 α = 1 无损记忆 (Sigmoid 无法做到)。
+        # d-向量的每个维度独立衰减:
+        #   S_t[i,j] = α_t[i] · S_{t-1}[i,j] + k_t[i] · v_t[j]
+        raw = self.decay_proj(hidden_states)                                # [B,T,H*d]
+        log_alpha = -torch.nn.functional.softplus(raw)                      # [B,T,H*d]
+        log_alpha = log_alpha.view(B, T, H, d).transpose(1, 2)             # [B,H,T,d]
+
+        # --- Apply attention_mask: PAD tokens must be invisible to the recurrence ---
+        # --- 应用注意力掩码: PAD token 必须对递推不可见 ---
+        # For PAD positions (mask=0): set log_α=0 (α=1, preserve state) and kv=0 (no contribution).
+        # This makes S_t = 1·S_{t-1} + 0 = S_{t-1}, i.e. PAD is a no-op on the state.
+        # 对于 PAD 位置 (mask=0): 设 log_α=0 (α=1, 保持状态) 且 kv=0 (无贡献)。
+        # 这使得 S_t = 1·S_{t-1} + 0 = S_{t-1}, 即 PAD 对状态是空操作。
+        if attention_mask is not None:
+            # attention_mask: [B, T] → [B, 1, T, 1] for broadcasting with [B, H, T, d]
+            mask = attention_mask[:, None, :, None].to(log_alpha.dtype)     # [B,1,T,1]
+            log_alpha = log_alpha * mask    # PAD → log_α=0 → α=1
+            k = k * mask                    # PAD → k=0
+            v = v * mask                    # PAD → v=0 → kv=0
 
         # ══════════════════════════════════════════════════════════
         # Inference path (RNN mode): O(1) per token per layer
@@ -427,7 +466,7 @@ class MonoidAttention(nn.Module):
             # Outer product: k_t ⊗ v_t ∈ ℝ^{H×d×d}
             # 外积: k_t ⊗ v_t ∈ ℝ^{H×d×d}
             kv_t = torch.einsum('bhd, bhe -> bhde', k[:, :, 0], v[:, :, 0])
-            log_t = log_alpha[:, :, 0]  # [B,H,1]
+            log_t = log_alpha[:, :, 0]  # [B,H,d]
 
             prev = monoid_cache.get_state(self.layer_idx) if monoid_cache else None
             if prev is None:
@@ -467,16 +506,16 @@ class MonoidAttention(nn.Module):
             kv = torch.einsum('bhtd, bhte -> bhtde', k, v)      # [B,H,T,d,d]
             states, (log_acc, S_T) = parallel_scan_with_state(log_alpha, kv)
 
-            # Add h0 contribution: S_t += (∏_{i=0}^{t} α_i) · h0
-            # 叠加 h0 贡献: S_t += (∏_{i=0}^{t} α_i) · h0
-            cum_log_alpha = torch.cumsum(log_alpha, dim=2)       # [B,H,T,1]
-            h0_decay = torch.exp(cum_log_alpha).unsqueeze(-1)    # [B,H,T,1,1]
+            # Add h0 contribution: S_t += diag(∏_{i=0}^{t} α_i) · h0
+            # 叠加 h0 贡献: S_t += diag(∏_{i=0}^{t} α_i) · h0
+            cum_log_alpha = torch.cumsum(log_alpha, dim=2)       # [B,H,T,d]
+            h0_decay = torch.exp(cum_log_alpha).unsqueeze(-1)    # [B,H,T,d,1]
             states = states + h0_decay * self.h0.unsqueeze(2)    # broadcast h0 [1,H,1,d,d]
 
             # Final state includes h0 contribution
             # 最终状态包含 h0 贡献
-            total_h0_decay = torch.exp(log_acc).unsqueeze(-1)    # [B,H,1,1]
-            S_final = S_T + total_h0_decay * self.h0.squeeze(0)  # [B,H,d,d] (squeeze batch dim of h0)
+            total_h0_decay = torch.exp(log_acc).unsqueeze(-1)    # [B,H,d,1]
+            S_final = S_T + total_h0_decay * self.h0.squeeze(0)  # [B,H,d,d]
             # h0 is [1,H,d,d], squeeze(0) removed for clarity but expand also works
             final_state = (log_acc, S_final)
 
@@ -507,16 +546,16 @@ class MonoidAttention(nn.Module):
         # 向量化外积: 一次性计算所有 t 的 k_t ⊗ v_t
         kv = torch.einsum('bhtd, bhte -> bhtde', k, v)           # [B,H,T,d,d]
 
-        # Parallel prefix scan: S_t = α_t·S_{t-1} + kv_t (from S=0)
-        # 并行前缀扫描: S_t = α_t·S_{t-1} + kv_t (从 S=0 开始)
-        # Keep log_alpha as [B,H,T,1] — CUDA kernel backward expects this shape.
-        # 保持 log_alpha 为 [B,H,T,1] — CUDA kernel 反向传播需要此形状。
+        # Parallel prefix scan: S_t = diag(α_t)·S_{t-1} + kv_t (from S=0)
+        # 并行前缀扫描: S_t = diag(α_t)·S_{t-1} + kv_t (从 S=0 开始)
+        # log_alpha is [B,H,T,d] — vector decay per dimension.
+        # log_alpha 为 [B,H,T,d] — 每维度向量衰减。
         states = parallel_scan(log_alpha, kv)                     # [B,H,T,d,d]
 
-        # Add h0 contribution: S_t += (∏_{i=0}^{t} α_i) · h0
-        # 叠加 h0 贡献: S_t += (∏_{i=0}^{t} α_i) · h0
-        cum_log_alpha = torch.cumsum(log_alpha, dim=2)            # [B,H,T,1]
-        h0_decay = torch.exp(cum_log_alpha).unsqueeze(-1)         # [B,H,T,1,1]
+        # Add h0 contribution: S_t += diag(∏_{i=0}^{t} α_i) · h0
+        # 叠加 h0 贡献: S_t += diag(∏_{i=0}^{t} α_i) · h0
+        cum_log_alpha = torch.cumsum(log_alpha, dim=2)            # [B,H,T,d]
+        h0_decay = torch.exp(cum_log_alpha).unsqueeze(-1)         # [B,H,T,d,1]
         states = states + h0_decay * self.h0.unsqueeze(2)         # broadcast h0 [1,H,1,d,d]
 
         # Vectorized readout: o_t = q_t · S_t for all t at once
@@ -557,6 +596,7 @@ class MonoidDecoderLayer(nn.Module):
     def forward(
         self,
         hidden_states: Tensor,
+        attention_mask: Tensor | None = None,
         monoid_cache: MonoidCache | None = None,
         use_cache: bool = False,
     ) -> Tensor:
@@ -564,7 +604,7 @@ class MonoidDecoderLayer(nn.Module):
         # --- 注意力块 + 残差连接 ---
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
-        hidden_states, _ = self.self_attn(hidden_states, monoid_cache=monoid_cache, use_cache=use_cache)
+        hidden_states, _ = self.self_attn(hidden_states, attention_mask=attention_mask, monoid_cache=monoid_cache, use_cache=use_cache)
         hidden_states = residual + hidden_states
 
         # --- FFN block with residual ---
@@ -600,7 +640,7 @@ class MonoidPreTrainedModel(PreTrainedModel):
                 module.weight.data[module.padding_idx].zero_()
 
         if isinstance(module, MonoidAttention):
-            nn.init.constant_(module.decay_proj.bias, 4.0)
+            nn.init.constant_(module.decay_proj.bias, 1.0)
 
 class MonoidModel(MonoidPreTrainedModel):
     """
@@ -625,6 +665,7 @@ class MonoidModel(MonoidPreTrainedModel):
     def forward(
         self,
         input_ids: Tensor | None = None,
+        attention_mask: Tensor | None = None,
         inputs_embeds: Tensor | None = None,
         monoid_cache: MonoidCache | None = None,
         use_cache: bool = False,
@@ -638,11 +679,12 @@ class MonoidModel(MonoidPreTrainedModel):
                 hidden_states = self._gradient_checkpointing_func(
                     layer.__call__,
                     hidden_states,
+                    attention_mask,
                     monoid_cache,
                     use_cache,
                 )
             else:
-                hidden_states = layer(hidden_states, monoid_cache=monoid_cache, use_cache=use_cache)
+                hidden_states = layer(hidden_states, attention_mask=attention_mask, monoid_cache=monoid_cache, use_cache=use_cache)
 
         hidden_states = self.norm(hidden_states)
 
@@ -723,9 +765,13 @@ class MonoidForCausalLM(MonoidPreTrainedModel, GenerationMixin):
             # Cache exists → only feed the latest token (O(1) inference)
             # 缓存已存在 → 只需输入最新的 token (O(1) 推理)
             input_ids = input_ids[:, -1:]
+            # Decode step: single real token, no PAD → mask not needed
+            # 解码步: 单个真实token, 无PAD → 不需要掩码
+            attention_mask = None
 
         model_inputs = {
             "input_ids": input_ids,
+            "attention_mask": attention_mask,
             "monoid_cache": past_key_values,
             "use_cache": True,
         }
@@ -734,8 +780,8 @@ class MonoidForCausalLM(MonoidPreTrainedModel, GenerationMixin):
     def forward(
         self,
         input_ids: Tensor | None = None,
-        attention_mask: Tensor | None = None,    # kept for API compat; monoid ignores this
-                                                  # 保留 API 兼容性; 幺半群不使用
+        attention_mask: Tensor | None = None,    # [B,T] 1=real, 0=pad — used to mask PAD from recurrence
+                                                  # [B,T] 1=真实token, 0=填充 — 用于屏蔽PAD对递推的影响
         position_ids: Tensor | None = None,      # kept for API compat; monoid ignores this
                                                   # 保留 API 兼容性; 幺半群不使用
         past_key_values: MonoidCache | None = None,
@@ -762,6 +808,7 @@ class MonoidForCausalLM(MonoidPreTrainedModel, GenerationMixin):
 
         outputs = self.model(
             input_ids=input_ids,
+            attention_mask=attention_mask,
             inputs_embeds=inputs_embeds,
             monoid_cache=cache,
             use_cache=bool(use_cache),

README.md

@@ -9,6 +9,7 @@ tags:
   - linear-attention
   - state-space
   - O(1)-inference
+  - vector-decay
   - reasoning
 pipeline_tag: text-generation
 model-index:
@@ -20,124 +21,113 @@ model-index:
 
 A 1.3B parameter language model that replaces softmax attention with **causal monoid state compression**, achieving **O(1) time per token** and **O(1) memory** at inference — regardless of sequence length.
 
-Fine-tuned for enhanced reasoning with structured chain-of-thought data.
-
 ## Monoid Attention — Internal Structure
 
 ```
                           MonoidAttention (per layer, per head)
- ┌─────────────────────────────────────────────────────────────────────┐
- │                                                                     │
- │   x_t ∈ R^{2048}                                                   │
- │    │                                                                │
- │    ├──> q_proj ──> RMSNorm ──> q_t ∈ R^{d}     (query)            │
- │    │                                                                │
- │    ├──> k_proj ──> RMSNorm ──> SiLU ──> k_t ∈ R^{d}  (key, >= 0) │
- │    │                                                                │
- │    ├──> v_proj ──> v_t ∈ R^{d}                  (value)            │
- │    │                                                                │
- │    └──> decay_proj ──> sigmoid ──> alpha_t ∈ (0,1)  (decay gate)   │
- │                                                                     │
- │         k_t (x) v_t                                                 │
- │            │            ┌──────────────────────────────┐            │
- │            │            │  State Matrix S_t ∈ R^{d x d} │            │
- │            v            │                              │            │
- │    S_t = alpha_t * S_{t-1} + k_t (x) v_t              │            │
- │            │            │  "Compressed causal history" │            │
- │            │            └──────────────────────────────┘            │
- │            v                                                        │
- │    o_t = q_t . S_t ──> o_proj ──> output                           │
- │                                                                     │
- └─────────────────────────────────────────────────────────────────────┘
-```
-
-## Monoid State Diagonal — O(1) Compression Contour
-
-The state matrix `S_t` accumulates causal history along its diagonal. Each head maintains an independent `d x d` state that compresses ALL past tokens into a fixed footprint:
-
-```
-   State Matrix S_t ∈ R^{64 x 64}  (one per head, 32 heads per layer)
-
-   k-dim -->
-   0   8  16  24  32  40  48  56  63
-   ┌───┬───┬───┬───┬───┬───┬───┬───┐  0
-   │***│** │*  │   │   │   │   │   │     v-dim
-   │***│** │*  │.  │   │   │   │   │      |
-   ├───┼───┼───┼───┼───┼───┼───┼───┤  8   |
-   │** │***│** │*  │.  │   │   │   │      v
-   │*  │***│** │*  │.  │   │   │   │
-   ├───┼───┼───┼───┼───┼───┼───┼───┤ 16
-   │*  │** │***│** │*  │.  │   │   │
-   │.  │*  │***│** │*  │.  │   │   │
-   ├───┼───┼───┼───┼───┼───┼───┼───┤ 24
-   │   │.  │** │***│** │*  │.  │   │
-   │   │   │*  │***│** │*  │.  │   │
-   ├───┼───┼───┼───┼───┼───┼───┼───┤ 32
-   │   │   │.  │** │***│** │*  │.  │
-   │   │   │   │*  │***│** │*  │.  │
-   ├───┼───┼───┼───┼───┼───┼───┼───┤ 40
-   │   │   │   │.  │** │***│** │*  │
-   │   │   │   │   │*  │***│** │*  │
-   ├───┼───┼───┼───┼───┼───┼───┼───┤ 48
-   │   │   │   │   │.  │** │***│** │
-   │   │   │   │   │   │*  │***│** │
-   ├───┼───┼───┼───┼───┼───┼───┼───┤ 56
-   │   │   │   │   │   │.  │** │***│
-   │   │   │   │   │   │   │*  │***│
-   └───┴───┴───┴───┴───┴─��─┴───┴───┘ 63
-
-   Legend:  *** = high activation (recent tokens, alpha^0 ~ alpha^2)
-            **  = medium (alpha^3 ~ alpha^5)
-            *   = fading (alpha^6 ~ alpha^10)
-            .   = near-zero (alpha^11+, effectively forgotten)
-                = zero (never reached or fully decayed)
-
-   The diagonal band emerges because S_t = SUM_{i<=t} alpha^{t-i} * k_i (x) v_i.
-   Recent outer products dominate near the diagonal; older ones decay
-   exponentially via alpha, creating this characteristic contour.
+ ┌─────────────────────────────────────────────────────────────────────────┐
+ │                                                                         │
+ │   x_t ∈ R^{2048}                                                       │
+ │    │                                                                    │
+ │    ├──> q_proj ──> RMSNorm ──> q_t ∈ R^d          (query, scaled 1/√d) │
+ │    │                                                                    │
+ │    ├──> k_proj ──> RMSNorm ──> SiLU ──> k_t ∈ R^d (key, non-negative) │
+ │    │                                                                    │
+ │    ├──> v_proj ──> v_t ∈ R^d                       (value)             │
+ │    │                                                                    │
+ │    └──> decay_proj ──> -Softplus ──> log α_t ∈ R^d (vector decay gate) │
+ │                                                                         │
+ │         k_t ⊗ v_t                                                       │
+ │            │             ┌─────────────────────────────────┐            │
+ │            │             │  State Matrix S_t ∈ R^{d x d}   │            │
+ │            v             │  "Compressed causal history"    │            │
+ │    S_t = diag(α_t) · S_{t-1} + k_t ⊗ v_t                 │            │
+ │            │             │  α_t ∈ (0,1]^d per dimension    │            │
+ │            │             └─────────────────────────────────┘            │
+ │            v                                                            │
+ │    o_t = q_t · S_t ──> o_proj ──> output                               │
+ │                                                                         │
+ └─────────────────────────────────────────────────────────────────────────┘
 ```
 
-
 ## Key Properties
 
 | Property | Transformer (Llama) | Spartacus (Monoid) |
 |---|---|---|
-| Inference time per token | O(T) -- scans full KV-cache | **O(1)** -- single state update |
-| Inference memory per layer | O(T) -- stores all past K,V | **O(1)** -- fixed d x d state matrix |
-| Sequence length extrapolation | Degrades beyond training length | **Unlimited** -- state size is constant |
+| Inference time per token | O(T) — scans full KV-cache | **O(1)** — single state update |
+| Inference memory per layer | O(T) — stores all past K,V | **O(1)** — fixed d×d state matrix |
+| Sequence length extrapolation | Degrades beyond training length | **Unlimited** — state size is constant |
 | Causality | Imposed via attention mask | **Built into the recurrence** |
-| Training complexity | O(T^2) | **O(T)** via parallel prefix scan |
+| Training complexity | O(T²) | **O(T)** via parallel prefix scan |
 
 ## The Monoid Recurrence
 
 Standard attention computes:
 
 ```
-o_t = sum_{i<=t} softmax(q_t . k_i) v_i    -- requires O(T) KV-cache
+o_t = Σ_{i≤t} softmax(q_t · k_i) v_i    — requires O(T) KV-cache
 ```
 
 Monoid attention compresses the entire causal history into a **fixed-size state matrix** S_t per head:
 
 ```
-S_t = alpha_t * S_{t-1} + k_t (x) v_t      -- explicit causal recurrence
-o_t = q_t . S_t                              -- state readout
+S_t = diag(α_t) · S_{t-1} + k_t ⊗ v_t     — vector decay monoid recurrence
+o_t = q_t · S_t                              — state readout
+```
+
+This is a monoid because the binary operator `(log_α, S) ⊕ (log_β, X) = (log_α + log_β, exp(log_β)·S + X)` is **associative**, enabling O(T) parallel prefix scan for training and O(1) sequential update for inference.
+
+## Vector Decay — Per-Dimension Memory Lifetimes
+
+Unlike scalar decay (one α per head), Spartacus uses **vector decay**: each dimension of the d-vector has its own independent decay rate α_t[i] ∈ (0, 1]:
+
+```
+S_t[i,j] = α_t[i] · S_{t-1}[i,j] + k_t[i] · v_t[j]
+```
+
+This allows different feature dimensions to specialize:
+- **Fast-decaying dimensions** (α ≈ 0) — local syntax, punctuation, function words
+- **Slow-decaying dimensions** (α ≈ 1) — entity memory, topic tracking, long-range facts
+
+The decay gate uses **Negative Softplus** activation:
+
+```
+log α_t = -softplus(W·x_t + b)
 ```
 
-where `alpha_t = sigmoid(decay_proj(x_t))` is a learned, content-dependent decay gate that controls how fast past information fades.
+| Property | Value |
+|---|---|
+| Range | α ∈ (0, 1] — bounded, no explosion |
+| Perfect memory | W·x → -∞ ⟹ softplus → 0 ⟹ α → 1 (lossless retention) |
+| Full forgetting | W·x → +∞ ⟹ softplus → ∞ ⟹ α → 0 (complete reset) |
+| Stability | α ≤ 1 by construction — no divergence regardless of input magnitude |
+
+## Attention Mask — Padding-Aware Recurrence
 
-## Explicit Causal Modeling
+The monoid recurrence correctly handles `attention_mask` for padded batches (e.g., left-padding during `generate()`). For PAD positions (mask=0):
 
-Unlike Transformers where causality is a constraint imposed by masking, Spartacus makes causality a **first-class citizen**:
+```
+log_α = 0    →  α = 1  (preserve state unchanged)
+k = 0, v = 0 →  kv = 0 (no information injected)
+```
 
-- The decay gate `alpha_t` explicitly controls per-head information retention at every timestep
-- The model learns **when to forget** rather than encoding **where tokens are** (no positional encoding needed)
-- No attention mask required -- causality is structural, not enforced
+Net effect: `S_t = 1·S_{t-1} + 0 = S_{t-1}` — PAD acts as the **monoid identity element**, completely invisible to the recurrence. This ensures identical outputs whether inputs are padded or not.
 
 ## Design Choices
 
-- **SiLU-activated keys**: `k = SiLU(k_proj(x))` ensures non-negative keys, making the state matrix `S` positive semi-definite (PSD). This prevents "feature erasure" where one token's contribution cancels another's
-- **Log-space decay**: Working in log-space `log(alpha)` avoids numerical underflow when `alpha^T -> 0` for long sequences
-- **Learnable h0**: The initial state `S_0 = h0` is a learnable parameter (zero-initialized), acting as a compressed "system prompt"
+- **SiLU-activated keys**: `k = SiLU(k_proj(x))` ensures non-negative keys, making the state matrix S positive semi-definite (PSD). This prevents "feature erasure" where one token's contribution cancels another's
+- **QK-Norm**: RMSNorm on both q and k before readout, stabilizing the scale of q·S when the state matrix accumulates many outer products
+- **Log-space decay**: Working in log-space `log(α)` avoids numerical underflow when α^T → 0 for long sequences
+- **Learnable h0**: The initial state S₀ = h0 is a learnable parameter (zero-initialized), acting as a compressed "system prompt"
+- **Negative Softplus gate**: Ensures α ∈ (0, 1] by construction — allows perfect memory (α=1) while preventing state explosion (α>1)
+
+## Three Forward Paths
+
+| Path | Condition | Complexity | Description |
+|---|---|---|---|
+| Training | `use_cache=False` | O(T) parallel scan | Vectorized outer products → parallel prefix scan → vectorized readout |
+| Inference prefill | `use_cache=True, T>1` | O(T) parallel scan | Same as training + extracts final state S_T for cache |
+| Inference decode | `use_cache=True, T=1` | **O(1)** monoid_op | Single `monoid_op` to fold new token into state → one matmul readout |
 
 ## Model Details
 
@@ -151,7 +141,8 @@ Unlike Transformers where causality is a constraint imposed by masking, Spartacu
 | Layers | 16 |
 | Attention heads | 32 |
 | Head dimension | 64 |
-| State matrix per head | 64 x 64 = 4096 floats |
+| Decay gate | Vector decay, d=64 per head |
+| State matrix per head | 64 × 64 = 4,096 floats |
 | Vocabulary | 128,256 (Llama-3.2 tokenizer) |
 | Precision | bfloat16 |
 
@@ -167,7 +158,7 @@ Unlike Transformers where causality is a constraint imposed by masking, Spartacu
 
 ### Comparison with ~1B Baselines (acc_norm, 0-shot)
 
-| Task | Spartacus-1B-Instruct | TinyLlama-1.1B | Llama 3.2-1B | Mamba-1.4B | RWKV-6-1.6B |
+| Task | Spartacus-1B | TinyLlama-1.1B | Llama 3.2-1B | Mamba-1.4B | RWKV-6-1.6B |
 |---|---|---|---|---|---|
 | ARC-C | **0.3063** | 0.3268 | ~0.359 | 0.284 | ~0.301 |
 | ARC-E | **0.5518** | 0.5547 | ~0.752 | 0.512 | ~0.530 |
@@ -177,34 +168,15 @@ Unlike Transformers where causality is a constraint imposed by masking, Spartacu
 
 > Spartacus achieves competitive performance with sub-quadratic models (Mamba, RWKV) while maintaining **O(1) inference time and memory per token**. Scores marked with ~ are approximate community-reported values.
 
-## Training
-
-### Stage 1: General SFT
-
-- **Base weights**: Transferred from Llama-3.2-1B-Instruct (embeddings, MLP, norms)
-- **Data**: Capybara + smol-smoltalk (general conversation)
-- **Training**: Full-parameter SFT
-
-### Stage 2: Reasoning Enhancement
-
-- **Data mix**: 60% Qwen3-Short-Reasoning + 20% Capybara + 20% smol-smoltalk
-- **Steps**: 2,000
-- **Learning rate**: 2e-5 (cosine schedule, 50 warmup steps)
-- **Batch size**: 8
-- **Sequence length**: 2,048
-- **Precision**: bfloat16
-- **Optimizer**: AdamW (weight decay 0.01, max grad norm 1.0)
-
-The reasoning data uses structured "Thought + Solution" format to strengthen chain-of-thought capabilities while the general data prevents catastrophic forgetting.
-
 ## Parallel Scan Implementation
 
-The `monoid_scan_cuda.py` module provides a Triton JIT-compiled parallel prefix scan:
+The `monoid_scan_cuda.py` module provides a Triton JIT-compiled parallel prefix scan for the vector-decay monoid:
 
-- **Forward**: Sequential scan along T, parallelized across B x H x D on GPU via Triton kernels
-- **Backward**: Reverse-order adjoint scan computes gradients for both values and log-decay gates
+- **Grid**: `(B*H*D_k, ceil(D_v/BLOCK_DV))` — one program per state matrix row
+- **Forward**: Sequential scan along T per row, parallelized across all (batch, head, d_k) dimensions
+- **Backward**: Reverse-order adjoint scan with per-row D_v reduction (minimal atomic_add)
 - **Fallback**: Pure PyTorch sequential scan for CPU/MPS
-- **Auto-dispatch**: CUDA -> Triton kernel, otherwise -> PyTorch fallback
+- **Auto-dispatch**: CUDA → Triton kernel, otherwise → PyTorch fallback
 
 ## Usage
 
@@ -231,7 +203,7 @@ print(tokenizer.decode(outputs[0], skip_special_tokens=True))
 
 ```
 MonoidForCausalLM.py       # Model architecture (MonoidConfig, MonoidAttention, MonoidForCausalLM)
-monoid_scan_cuda.py        # Triton JIT parallel prefix scan + PyTorch fallback
+monoid_scan_cuda.py        # Triton JIT parallel prefix scan (vector decay) + PyTorch fallback
 model.safetensors          # Model weights (bfloat16)
 config.json                # Model configuration
 tokenizer.json             # Llama-3.2 tokenizer
@@ -245,7 +217,7 @@ tokenizer.json             # Llama-3.2 tokenizer
   author={NoesisLab},
   year={2025},
   url={https://huggingface.co/NoesisLab/Spartacus-1B-Instruct},
-  description={Replaces softmax attention with monoid state compression for constant-time, constant-memory autoregressive generation}
+  description={Replaces softmax attention with vector-decay monoid state compression for constant-time, constant-memory autoregressive generation}
 }
 ```
 

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a32bdaef8ff9aa5cba602faa280ab5d3526515a6abd97d411c3448f9b9ebcdc7
-size 2679277744
+oid sha256:d5cd463898c4ce262d12fe56c6227d0c1117680aa13892f9cac6e100a1db9077
+size 2811462896

monoid_scan_cuda.py

@@ -2,36 +2,33 @@
 monoid_scan_cuda.py — Triton CUDA JIT Accelerated Parallel Prefix Scan
 monoid_scan_cuda.py — Triton CUDA JIT 加速的并行前缀扫描
 
-This module implements the parallel prefix scan for the monoid recurrence:
-  y_t = exp(log_decay_t) · y_{t-1} + x_t
-本模块实现幺半群递推的并行前缀扫描:
-  y_t = exp(log_decay_t) · y_{t-1} + x_t
+This module implements the parallel prefix scan for the vector-decay monoid recurrence:
+  y_t[i,:] = exp(log_decay_t[i]) · y_{t-1}[i,:] + x_t[i,:]
+本模块实现向量衰减幺半群递推的并行前缀扫描:
+  y_t[i,:] = exp(log_decay_t[i]) · y_{t-1}[i,:] + x_t[i,:]
 
 This is the computational backbone of Monoid Attention's state compression.
 这是幺半群注意力状态压缩的计算骨干。
 
-Why parallel prefix scan matters / 并行前缀扫描为什么重要:
-  The monoid recurrence S_t = α_t·S_{t-1} + kv_t is inherently sequential.
-  However, because (log_α, S) ⊕ (log_β, X) = (log_α+log_β, exp(log_β)·S+X)
-  is ASSOCIATIVE, we can compute all prefix sums S_1..S_T via a parallel
-  reduction tree in O(log T) depth instead of O(T) sequential steps.
-  幺半群递推 S_t = α_t·S_{t-1} + kv_t 本质上是串行的。
-  但因为 (log_α, S) ⊕ (log_β, X) = (log_α+log_β, exp(log_β)·S+X)
-  满足结合律, 我们可以通过并行归约树在 O(log T) 深度内计算所有前缀和 S_1..S_T,
-  而非 O(T) 的串行步骤。
-
-  Training uses O(T) parallel scan (this file).
-  Inference uses O(1) sequential monoid_op (in MonoidForCausalLM.py).
-  训练使用 O(T) 并行扫描 (本文件)。
-  推理使用 O(1) 串行 monoid_op (在 MonoidForCausalLM.py 中)。
+Vector decay: each dimension of the D_k×D_v state matrix has its own
+per-dimension decay rate α_t ∈ ℝ^{D_k}, enabling different feature
+dimensions to have independent memory lifetimes (fast-decaying for
+local syntax, slow-decaying for global entity memory).
+向量衰减: D_k×D_v 状态矩阵的每个维度拥有独立的衰减率 α_t ∈ ℝ^{D_k},
+使不同特征维度拥有独立的记忆生命周期 (快速衰减用于局部语法, 慢速衰减用于全局实体记忆)。
 
 Implementation:
-  Forward: sequential scan along T, parallelized across B*H*D on GPU.
+  Forward: sequential scan along T, parallelized across B*H*D_k on GPU.
+           Each program handles one row of the state matrix (D_v elements)
+           with a scalar decay per row.
   Backward: reverse-order adjoint scan for gradient computation.
+            Per-row reduction for log_decay gradient (no atomic_add needed).
   Auto-dispatches: CUDA → Triton kernel, CPU/MPS → PyTorch fallback.
 
-  前向: 沿 T 维顺序扫描, 跨 B*H*D 在 GPU 上并行。
+  前向: 沿 T 维顺序扫描, 跨 B*H*D_k 在 GPU 上并行。
+        每个 program 处理状态矩阵的一行 (D_v 个元素), 每行一个标量衰减。
   反向: 逆序伴随变量扫描计算梯度。
+        逐行归约计算 log_decay 梯度 (无需 atomic_add)。
   自动分派: CUDA → Triton 核函数, CPU/MPS → PyTorch 回退。
 """
 
@@ -60,18 +57,18 @@ def _sequential_scan(log_decays: Tensor, values: Tensor) -> Tensor:
     Pure PyTorch sequential scan fallback (when no CUDA / Triton available).
     纯 PyTorch 串行扫描回退 (无 CUDA / Triton 时使用)。
 
-    Implements the monoid recurrence step by step:
+    Implements the vector-decay monoid recurrence step by step:
       acc_0 = 0
-      acc_t = exp(log_decay_t) · acc_{t-1} + values_t
+      acc_t[i,:] = exp(log_decay_t[i]) · acc_{t-1}[i,:] + values_t[i,:]
     This is O(T) sequential — correct but slow on GPU.
-    逐步实现幺半群递推:
+    逐步实现向量衰减幺半群递推:
       acc_0 = 0
-      acc_t = exp(log_decay_t) · acc_{t-1} + values_t
+      acc_t[i,:] = exp(log_decay_t[i]) · acc_{t-1}[i,:] + values_t[i,:]
     这是 O(T) 串行的 — 结果正确但在 GPU 上较慢。
 
     Args:
-        log_decays: [B, H, T, 1]     — log of per-head per-step decay gates
-                                        每头每步衰减门的对数
+        log_decays: [B, H, T, D_k]     — log of per-dimension per-step decay gates
+                                           每维度每步衰减门的对数
         values:     [B, H, T, D_k, D_v] — outer products k_t⊗v_t to accumulate
                                            待累积的外积 k_t⊗v_t
     Returns:
@@ -84,17 +81,17 @@ def _sequential_scan(log_decays: Tensor, values: Tensor) -> Tensor:
     # acc 代表 S_t — 时刻 t 的压缩因果状态
     acc = torch.zeros(B, H, D_k, D_v, device=values.device, dtype=values.dtype)
     for t in range(T):
-        # S_t = α_t · S_{t-1} + kv_t  (the core monoid recurrence)
-        # S_t = α_t · S_{t-1} + kv_t  (核心幺半群递推)
-        decay_t = torch.exp(log_decays[:, :, t]).unsqueeze(-1)  # [B,H,1,1]
+        # S_t = diag(α_t) · S_{t-1} + kv_t  (vector decay monoid recurrence)
+        # S_t = diag(α_t) · S_{t-1} + kv_t  (向量衰减幺半群递推)
+        decay_t = torch.exp(log_decays[:, :, t]).unsqueeze(-1)  # [B,H,D_k,1]
         acc = acc * decay_t + values[:, :, t]
         out[:, :, t] = acc
     return out
 
 
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-# Triton Kernels — GPU-accelerated scan
-# Triton 核函数 — GPU 加速扫描
+# Triton Kernels — GPU-accelerated scan (vector decay)
+# Triton 核函数 — GPU 加速扫描 (向量衰减)
 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 
 if HAS_TRITON:
@@ -102,147 +99,136 @@ if HAS_TRITON:
     @triton.jit
     def _scan_fwd_kernel(
         LD_ptr, V_ptr, O_ptr,
-        T, D,
-        s_ld_bh, s_ld_t,
-        s_v_bh, s_v_t, s_v_d,
-        s_o_bh, s_o_t, s_o_d,
-        BLOCK_D: tl.constexpr,
+        T, D_v,
+        s_ld_bhdk, s_ld_t,
+        s_v_bhdk, s_v_t, s_v_dv,
+        s_o_bhdk, s_o_t, s_o_dv,
+        BLOCK_DV: tl.constexpr,
     ):
         """
-        Forward scan kernel — computes all prefix states S_1..S_T.
-        前向扫描核函数 — 计算所有前缀状态 S_1..S_T。
+        Forward scan kernel — computes all prefix states S_1..S_T (vector decay).
+        前向扫描核函数 — 计算所有前缀状态 S_1..S_T (向量衰减)。
 
         Parallelization strategy / 并行化策略:
-          - program_id(0) = bh: one program per (batch, head) pair
-            每个 (batch, head) 对一个 program
-          - program_id(1) = db: one program per D-dimension block
-            每个 D 维 block 一个 program
+          - program_id(0) = bhdk: one program per (batch, head, d_k row) triple
+            每个 (batch, head, d_k 行) 三元组一个 program
+          - program_id(1) = dvb: one program per D_v-dimension block (typically 1 block)
+            每个 D_v 维 block 一个 program (通常只有 1 个 block)
           - Sequential loop over T (the causal recurrence is inherently sequential)
             沿 T 维串行循环 (因果递推本质上是串行的)
 
-        Each program computes: acc_t = exp(ld_t) * acc_{t-1} + val_t
-        for a BLOCK_D-wide slice of the flattened d_k*d_v state matrix.
-        每个 program 计算展平的 d_k*d_v 状态矩阵的一个 BLOCK_D 宽的切片。
+        Each program handles one row of the D_k×D_v state matrix, where the
+        decay is a single scalar per row. This eliminates the need for
+        row-index computation in the inner loop.
+        每个 program 处理 D_k×D_v 状态矩阵的一行, 该行的衰减是一个标量。
+        这消除了内循环中行索引计算的需要。
 
-        Note: while the T-loop is sequential within each program,
-        B*H*ceil(D/BLOCK_D) programs run in parallel on the GPU.
-        注意: 虽然 T 循环在每个 program 内是串行的,
-        但 B*H*ceil(D/BLOCK_D) 个 program 在 GPU 上并行运行。
+        Grid: (B*H*D_k, ceil(D_v/BLOCK_DV))
+        网格: (B*H*D_k, ceil(D_v/BLOCK_DV))
         """
-        bh = tl.program_id(0)
-        db = tl.program_id(1)
-        d_offs = db * BLOCK_D + tl.arange(0, BLOCK_D)
-        d_mask = d_offs < D
+        bhdk = tl.program_id(0)
+        dvb = tl.program_id(1)
+        dv_offs = dvb * BLOCK_DV + tl.arange(0, BLOCK_DV)
+        dv_mask = dv_offs < D_v
 
-        # acc = S_0 = 0 (identity element of the monoid)
-        # acc = S_0 = 0 (幺半群的单位元)
-        acc = tl.zeros([BLOCK_D], dtype=tl.float32)
+        # acc = S_0[row,:] = 0 (identity element of the monoid)
+        # acc = S_0[行,:] = 0 (幺半群的单位元)
+        acc = tl.zeros([BLOCK_DV], dtype=tl.float32)
 
-        ld_base = LD_ptr + bh * s_ld_bh
-        v_base = V_ptr + bh * s_v_bh
-        o_base = O_ptr + bh * s_o_bh
+        ld_base = LD_ptr + bhdk * s_ld_bhdk
+        v_base = V_ptr + bhdk * s_v_bhdk
+        o_base = O_ptr + bhdk * s_o_bhdk
 
         for t in range(T):
-            # Load log_decay and compute decay = exp(log_α_t)
-            # 加载 log_decay 并计算 decay = exp(log_α_t)
+            # Load scalar log_decay for this row at time t
+            # 加载此行在时刻 t 的标量 log_decay
             ld_val = tl.load(ld_base + t * s_ld_t).to(tl.float32)
             decay = tl.exp(ld_val)
 
-            # Load kv_t (a slice of the outer product k_t⊗v_t)
-            # 加载 kv_t (外积 k_t⊗v_t 的一个切片)
+            # Load kv_t[row, :] (one row of the outer product)
+            # 加载 kv_t[行, :] (外积的一行)
             val = tl.load(
-                v_base + t * s_v_t + d_offs * s_v_d,
-                mask=d_mask, other=0.0,
+                v_base + t * s_v_t + dv_offs * s_v_dv,
+                mask=dv_mask, other=0.0,
             ).to(tl.float32)
 
-            # Core recurrence: S_t = α_t · S_{t-1} + kv_t
-            # 核心递推: S_t = α_t · S_{t-1} + kv_t
+            # Core recurrence: S_t[i,:] = α_t[i] · S_{t-1}[i,:] + kv_t[i,:]
+            # 核心递推: S_t[i,:] = α_t[i] · S_{t-1}[i,:] + kv_t[i,:]
             acc = acc * decay + val
 
-            # Store S_t
+            # Store S_t[row, :]
             tl.store(
-                o_base + t * s_o_t + d_offs * s_o_d,
-                acc, mask=d_mask,
+                o_base + t * s_o_t + dv_offs * s_o_dv,
+                acc, mask=dv_mask,
             )
 
     @triton.jit
     def _scan_bwd_kernel(
         LD_ptr, O_ptr, GO_ptr, GV_ptr, GLD_ptr,
-        T, D,
-        s_ld_bh, s_ld_t,
-        s_o_bh, s_o_t, s_o_d,
-        s_go_bh, s_go_t, s_go_d,
-        s_gv_bh, s_gv_t, s_gv_d,
-        s_gld_bh, s_gld_t,
-        BLOCK_D: tl.constexpr,
+        T, D_v,
+        s_ld_bhdk, s_ld_t,
+        s_o_bhdk, s_o_t, s_o_dv,
+        s_go_bhdk, s_go_t, s_go_dv,
+        s_gv_bhdk, s_gv_t, s_gv_dv,
+        s_gld_bhdk, s_gld_t,
+        BLOCK_DV: tl.constexpr,
     ):
         """
-        Backward scan kernel — computes gradients via adjoint method.
-        反向扫描核函数 — 通过伴随方法计算梯度。
-
-        The forward recurrence is: y_t = a_t * y_{t-1} + x_t
-        前向递推: y_t = a_t * y_{t-1} + x_t
-
-        The adjoint (reverse-time) recurrence for the Lagrange multiplier λ:
-          λ_t = ∂L/∂y_t + a_{t+1} · λ_{t+1}    (backward in time)
-        伴随 (逆时间) 递推的拉格朗日乘子 λ:
-          λ_t = ∂L/∂y_t + a_{t+1} · λ_{t+1}    (时间反向)
-
-        Gradients / 梯度:
-          ∂L/∂x_t     = λ_t                       (gradient w.r.t. input values)
-                                                    (对输入值的梯度)
-          ∂L/∂log_a_t = a_t · Σ_D(λ_t · y_{t-1})  (gradient w.r.t. log-decay)
-                                                    (对对数衰减的梯度)
-
-        The gradient of log_decay is critical for training the decay gate:
-        it tells the model how to adjust each head's forgetting rate.
-        log_decay 的梯度对训练衰减门至关重要:
-        它告诉模型如何调整每个头的遗忘速率。
+        Backward scan kernel — computes gradients via adjoint method (vector decay).
+        反向扫描核函数 — 通过伴随方法计算梯度 (向量衰减)。
+
+        Each program handles one row of the state matrix (one d_k dimension).
+        The decay for this row is a scalar, so the log_decay gradient is:
+          ∂L/∂log_α_t[i] = α_t[i] · Σ_j(λ_t[i,j] · y_{t-1}[i,j])
+        The sum over j (D_v) is computed within this single program — no atomic_add.
+        每个 program 处理状态矩阵的一行 (一个 d_k 维度)。
+        该行的衰减是标量, 因此 log_decay 梯度为:
+          ∂L/∂log_α_t[i] = α_t[i] · Σ_j(λ_t[i,j] · y_{t-1}[i,j])
+        对 j (D_v) 的求和在单个 program 内完成 — 无需 atomic_add。
         """
-        bh = tl.program_id(0)
-        db = tl.program_id(1)
-        d_offs = db * BLOCK_D + tl.arange(0, BLOCK_D)
-        d_mask = d_offs < D
+        bhdk = tl.program_id(0)
+        dvb = tl.program_id(1)
+        dv_offs = dvb * BLOCK_DV + tl.arange(0, BLOCK_DV)
+        dv_mask = dv_offs < D_v
 
         # adj holds a_{t+1} · λ_{t+1}, initialized to 0 at the sequence end
         # adj 保存 a_{t+1} · λ_{t+1}, 在序列末尾初始化为 0
-        adj = tl.zeros([BLOCK_D], dtype=tl.float32)
+        adj = tl.zeros([BLOCK_DV], dtype=tl.float32)
 
         for t_rev in range(T):
             t = T - 1 - t_rev     # reverse time / 逆序时间
 
-            # Load ∂L/∂y_t (upstream gradient)
-            # 加载 ∂L/∂y_t (上游梯度)
+            # Load ∂L/∂y_t[row, :] (upstream gradient)
+            # 加载 ∂L/∂y_t[行, :] (上游梯度)
             go = tl.load(
-                GO_ptr + bh * s_go_bh + t * s_go_t + d_offs * s_go_d,
-                mask=d_mask, other=0.0,
+                GO_ptr + bhdk * s_go_bhdk + t * s_go_t + dv_offs * s_go_dv,
+                mask=dv_mask, other=0.0,
             ).to(tl.float32)
 
             # Adjoint: λ_t = ∂L/∂y_t + a_{t+1} · λ_{t+1}
             # 伴随: λ_t = ∂L/∂y_t + a_{t+1} · λ_{t+1}
             lam = go + adj
 
-            # ∂L/∂x_t = λ_t (gradient of values / 值的梯度)
+            # ∂L/∂x_t[row,:] = λ_t (gradient of values)
+            # ∂L/∂x_t[行,:] = λ_t (值的梯度)
             tl.store(
-                GV_ptr + bh * s_gv_bh + t * s_gv_t + d_offs * s_gv_d,
-                lam, mask=d_mask,
+                GV_ptr + bhdk * s_gv_bhdk + t * s_gv_t + dv_offs * s_gv_dv,
+                lam, mask=dv_mask,
             )
 
-            # ∂L/∂log_a_t = a_t · Σ_D(λ_t · y_{t-1})
-            # This gradient flows back to the decay gate (decay_proj),
-            # teaching the model how to control causal information retention.
-            # 此梯度回传到衰减门 (decay_proj),
-            # 教模型如何控制因果信息的保留。
-            ld_val = tl.load(LD_ptr + bh * s_ld_bh + t * s_ld_t).to(tl.float32)
+            # ∂L/∂log_α_t[i] = α_t[i] · Σ_j(λ_t[i,j] · y_{t-1}[i,j])
+            # Per-row scalar gradient: sum over D_v within this program.
+            # 逐行标量梯度: 在此 program 内对 D_v 求和。
+            ld_val = tl.load(LD_ptr + bhdk * s_ld_bhdk + t * s_ld_t).to(tl.float32)
             a_t = tl.exp(ld_val)
 
             if t > 0:
                 y_prev = tl.load(
-                    O_ptr + bh * s_o_bh + (t - 1) * s_o_t + d_offs * s_o_d,
-                    mask=d_mask, other=0.0,
+                    O_ptr + bhdk * s_o_bhdk + (t - 1) * s_o_t + dv_offs * s_o_dv,
+                    mask=dv_mask, other=0.0,
                 ).to(tl.float32)
-                grad_ld_partial = tl.sum(lam * y_prev) * a_t
-                tl.atomic_add(GLD_ptr + bh * s_gld_bh + t * s_gld_t, grad_ld_partial)
+                grad_ld = tl.sum(lam * y_prev) * a_t
+                tl.atomic_add(GLD_ptr + bhdk * s_gld_bhdk + t * s_gld_t, grad_ld)
 
             # Prepare for next step (t-1): adj = a_t · λ_t
             # 为下一步 (t-1) 准备: adj = a_t · λ_t
@@ -255,78 +241,91 @@ if HAS_TRITON:
 
     class _ParallelScanFn(Function):
         """
-        Custom autograd function for the parallel prefix scan.
-        并行前缀扫描的自定义 autograd 函数。
+        Custom autograd function for the parallel prefix scan (vector decay).
+        并行前缀扫描的自定义 autograd 函数 (向量衰减)。
 
         Forward: launches _scan_fwd_kernel to compute all prefix states.
+                 Grid: (B*H*D_k, ceil(D_v/BLOCK_DV)), one program per state row.
         Backward: launches _scan_bwd_kernel to compute gradients via adjoint method.
+                  Per-row reduction eliminates most atomic_add overhead.
 
         前向: 启动 _scan_fwd_kernel 计算所有前缀状态。
+              网格: (B*H*D_k, ceil(D_v/BLOCK_DV)), 每行状态一个 program。
         反向: 启动 _scan_bwd_kernel 通过伴随方法计算梯度。
+              逐行归约消除大部分 atomic_add 开销。
         """
         @staticmethod
         def forward(ctx, log_decays: Tensor, values: Tensor) -> Tensor:
             B, H, T, D_k, D_v = values.shape
-            D = D_k * D_v    # flattened state dimension / 展平的状态维度
 
-            # Flatten: [B,H,T,1] → [BH, T],  [B,H,T,Dk,Dv] → [BH, T, D]
-            # 展平: [B,H,T,1] → [BH, T],  [B,H,T,Dk,Dv] → [BH, T, D]
-            ld_flat = log_decays.squeeze(-1).contiguous().reshape(B * H, T)
-            v_flat = values.reshape(B * H, T, D).contiguous()
+            # Reshape for row-parallel kernel:
+            #   log_decays: [B, H, T, D_k] → permute to [B, H, D_k, T] → [B*H*D_k, T]
+            #   values:     [B, H, T, D_k, D_v] → permute to [B, H, D_k, T, D_v] → [B*H*D_k, T, D_v]
+            # 为行并行核函数重塑:
+            #   log_decays: [B, H, T, D_k] → 转置为 [B, H, D_k, T] → [B*H*D_k, T]
+            #   values:     [B, H, T, D_k, D_v] → 转置为 [B, H, D_k, T, D_v] → [B*H*D_k, T, D_v]
+            ld_flat = log_decays.permute(0, 1, 3, 2).contiguous().reshape(B * H * D_k, T)
+            v_flat = values.permute(0, 1, 3, 2, 4).contiguous().reshape(B * H * D_k, T, D_v)
             o_flat = torch.empty_like(v_flat)
 
-            BH = B * H
-            BLOCK_D = min(triton.next_power_of_2(D), 1024)
-            # Grid: (BH, ceil(D/BLOCK_D)) — one program per (batch*head, D-block)
-            # 网格: (BH, ceil(D/BLOCK_D)) — 每个 (batch*head, D-block) 一个 program
-            grid = (BH, triton.cdiv(D, BLOCK_D))
+            BHDK = B * H * D_k
+            BLOCK_DV = min(triton.next_power_of_2(D_v), 1024)
+            # Grid: (B*H*D_k, ceil(D_v/BLOCK_DV)) — one program per (batch, head, row, dv-block)
+            # 网格: (B*H*D_k, ceil(D_v/BLOCK_DV))
+            grid = (BHDK, triton.cdiv(D_v, BLOCK_DV))
 
             _scan_fwd_kernel[grid](
                 ld_flat, v_flat, o_flat,
-                T, D,
+                T, D_v,
                 ld_flat.stride(0), ld_flat.stride(1),
                 v_flat.stride(0), v_flat.stride(1), v_flat.stride(2),
                 o_flat.stride(0), o_flat.stride(1), o_flat.stride(2),
-                BLOCK_D=BLOCK_D,
+                BLOCK_DV=BLOCK_DV,
             )
 
             # Save for backward: need log_decays and forward outputs y_t
             # 为反向传播保存: 需要 log_decays 和前向输出 y_t
             ctx.save_for_backward(ld_flat, o_flat)
-            ctx.shape_info = (B, H, T, D_k, D_v, D, BH, BLOCK_D)
-            return o_flat.reshape(B, H, T, D_k, D_v)
+            ctx.shape_info = (B, H, T, D_k, D_v, BHDK, BLOCK_DV)
+            # Reshape back: [B*H*D_k, T, D_v] → [B, H, D_k, T, D_v] → [B, H, T, D_k, D_v]
+            return o_flat.reshape(B, H, D_k, T, D_v).permute(0, 1, 3, 2, 4).contiguous()
 
         @staticmethod
         def backward(ctx, grad_output: Tensor):
             ld_flat, o_flat = ctx.saved_tensors
-            B, H, T, D_k, D_v, D, BH, BLOCK_D = ctx.shape_info
+            B, H, T, D_k, D_v, BHDK, BLOCK_DV = ctx.shape_info
 
-            go_flat = grad_output.reshape(BH, T, D).contiguous()
+            # Permute grad_output to match row-parallel layout: [B,H,T,D_k,D_v] → [B*H*D_k, T, D_v]
+            go_flat = grad_output.permute(0, 1, 3, 2, 4).contiguous().reshape(BHDK, T, D_v)
             gv_flat = torch.empty_like(go_flat)
-            # Use f32 for atomic_add precision in gradient accumulation
-            # 使用 f32 保证 atomic_add 梯度累积的精度
-            gld_flat = torch.zeros(BH, T, device=ld_flat.device, dtype=torch.float32)
+            # Use f32 for gradient accumulation precision
+            # 使用 f32 保证梯度累积的精度
+            gld_flat = torch.zeros(BHDK, T, device=ld_flat.device, dtype=torch.float32)
 
-            grid = (BH, triton.cdiv(D, BLOCK_D))
+            grid = (BHDK, triton.cdiv(D_v, BLOCK_DV))
 
             _scan_bwd_kernel[grid](
                 ld_flat, o_flat, go_flat, gv_flat, gld_flat,
-                T, D,
+                T, D_v,
                 ld_flat.stride(0), ld_flat.stride(1),
                 o_flat.stride(0), o_flat.stride(1), o_flat.stride(2),
                 go_flat.stride(0), go_flat.stride(1), go_flat.stride(2),
                 gv_flat.stride(0), gv_flat.stride(1), gv_flat.stride(2),
                 gld_flat.stride(0), gld_flat.stride(1),
-                BLOCK_D=BLOCK_D,
+                BLOCK_DV=BLOCK_DV,
             )
 
-            grad_log_decays = gld_flat.to(grad_output.dtype).reshape(B, H, T, 1)
-            grad_values = gv_flat.reshape(B, H, T, D_k, D_v)
+            # Reshape gradients back to original layout
+            # 重塑梯度回原始布局
+            # gld: [B*H*D_k, T] → [B, H, D_k, T] → [B, H, T, D_k]
+            grad_log_decays = gld_flat.to(grad_output.dtype).reshape(B, H, D_k, T).permute(0, 1, 3, 2).contiguous()
+            # gv: [B*H*D_k, T, D_v] → [B, H, D_k, T, D_v] → [B, H, T, D_k, D_v]
+            grad_values = gv_flat.reshape(B, H, D_k, T, D_v).permute(0, 1, 3, 2, 4).contiguous()
             return grad_log_decays, grad_values
 
     def _triton_parallel_scan(log_decays: Tensor, values: Tensor) -> Tensor:
-        """Triton-accelerated parallel scan entry point.
-        Triton 加速的并行扫描入口。"""
+        """Triton-accelerated parallel scan entry point (vector decay).
+        Triton 加速的并行扫描入口 (向量衰减)。"""
         return _ParallelScanFn.apply(log_decays, values)
 
 else:
@@ -339,15 +338,16 @@ else:
 
 def parallel_scan(log_decays: Tensor, values: Tensor) -> Tensor:
     """
-    Parallel prefix scan — computes all prefix monoid sums.
-    并行前缀扫描 — 计算所有前缀幺半群和。
+    Parallel prefix scan — computes all prefix monoid sums (vector decay).
+    并行前缀扫描 — 计算所有前缀幺半群和 (向量衰减)。
 
     This is the training-time workhorse of Monoid Attention.
-    It computes S_1, S_2, ..., S_T where S_t = α_t·S_{t-1} + kv_t,
+    It computes S_1, S_2, ..., S_T where
+      S_t[i,:] = α_t[i]·S_{t-1}[i,:] + kv_t[i,:]
     for ALL timesteps simultaneously.
     这是幺半群注意力训练时的主力计算。
     它同时计算所有时间步的 S_1, S_2, ..., S_T,
-    其中 S_t = α_t·S_{t-1} + kv_t。
+    其中 S_t[i,:] = α_t[i]·S_{t-1}[i,:] + kv_t[i,:]。
 
     Auto-dispatches based on device:
       CUDA → Triton JIT kernel (fast, with custom backward)
@@ -357,8 +357,8 @@ def parallel_scan(log_decays: Tensor, values: Tensor) -> Tensor:
       CPU/MPS → PyTorch 串行扫描 (正确, 较慢)
 
     Args:
-        log_decays: [B, H, T, 1]        — log of decay gates α_t
-                                           衰减门 α_t 的对数
+        log_decays: [B, H, T, D_k]      — log of per-dimension decay gates α_t
+                                            每维度衰减门 α_t 的对数
         values:     [B, H, T, D_k, D_v] — outer products k_t⊗v_t
                                            外积 k_t⊗v_t
     Returns:
@@ -374,8 +374,8 @@ def parallel_scan_with_state(
     log_decays: Tensor, values: Tensor,
 ) -> Tuple[Tensor, Tuple[Tensor, Tensor]]:
     """
-    Parallel prefix scan + extract final state for inference handoff.
-    并行前缀扫描 + 提取最终状态用于推理切换。
+    Parallel prefix scan + extract final state for inference handoff (vector decay).
+    并行前缀扫描 + 提取最终状态用于推理切换 (向量衰减)。
 
     Used during prefill: compute all training-time prefix states,
     AND extract the final accumulated state S_T so that subsequent
@@ -389,22 +389,22 @@ def parallel_scan_with_state(
     这是训练模式 (并行扫描) 和推理模式 (串行 monoid_op) 之间的桥梁。
 
     Args:
-        log_decays: [B, H, T, 1]
+        log_decays: [B, H, T, D_k]
         values:     [B, H, T, D_k, D_v]
 
     Returns:
         output:      [B, H, T, D_k, D_v]  — all prefix states S_1..S_T
                                               所有前缀状态
         final_state: (log_acc, S_T) where
-            log_acc:     [B, H, 1]          — accumulated log-decay (for future monoid_op)
-                                              累积对数衰减 (供后续 monoid_op 使用)
-            final_state: [B, H, D_k, D_v]   — S_T, the compressed causal summary
-                                              S_T, 压缩的因果摘要
+            log_acc:     [B, H, D_k]         — accumulated log-decay vector (for future monoid_op)
+                                                累积对数衰减向量 (供后续 monoid_op 使用)
+            final_state: [B, H, D_k, D_v]    — S_T, the compressed causal summary
+                                                S_T, 压缩的因果摘要
     """
     output = parallel_scan(log_decays, values)
-    # Sum all log-decays to get the total accumulated decay
-    # 对所有 log-decay 求和得到总累积衰减
-    log_acc = log_decays.squeeze(-1).sum(dim=2, keepdim=True)  # [B, H, 1]
+    # Sum all log-decays over T to get the total accumulated decay per dimension
+    # 对所有 log-decay 沿 T 求和得到每个维度的总累积衰减
+    log_acc = log_decays.sum(dim=2)  # [B, H, D_k]
     # The last timestep's state IS the full causal summary
     # 最后一个时间步的状态就是完整的因果摘要
     final_state = output[:, :, -1]  # [B, H, D_k, D_v]

training_args.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:785ade819d78869db7f691e74573b3ce4183646d4e45ad2e1d0a940564b6b20f
-size 6289
+oid sha256:32938ebb880dc58fa7d6f8e45383c55e1d5d4352618531d62a28069918595445
+size 6417