ATen codegen fixes: matmul column-stride, FFN weight row-align, single-pass softmax

asm_templates/ffn_asm.py

@@ -556,8 +556,15 @@ def _emit_ffn_projection_chunk(
                 )
             lines.append(f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n")
         else:
+            # Sub-column tile within an MLEN block. The WEIGHT-read pointer must
+            # advance by blen*mlen (row-aligned) so each sub-column reads a DISTINCT
+            # MRAM row: the RTL matrix SRAM read is row-granular (raddr >> log2(MLEN)
+            # drops the low log2(MLEN) bits), so a plain +blen would collapse tiles
+            # 1..(mlen/blen-1) onto tile 0's row. The OUTPUT-write pointer stays at
+            # +blen (the low bits select the 4-wide column lane of the VLEN row).
+            # This matches the working projection_asm (_emit_projection_chunk).
             lines.append(
-                f"S_ADDI_INT gp{w_actual_register}, gp0, {(weight_row % (mlen // blen)) * blen} \n"
+                f"S_ADDI_INT gp{w_actual_register}, gp0, {(weight_row % (mlen // blen)) * blen * mlen} \n"
             )
             lines.append(
                 f"S_ADDI_INT gp{intermediate_register}, gp{result_base_register}, {(weight_row % (mlen // blen)) * blen} \n"

aten/ops/plena/softmax_ops.py

@@ -1,47 +1,60 @@
 """
 PLENA backend implementation for softmax operator.
 
-This encapsulates the online softmax algorithm that was previously
-written inline in fpvar_softmax_test.py (Tilelang-style).
-
 ATen-style: the function receives a PlenaCompiler context and TensorVar
 arguments, orchestrates PlenaCompiler calls, and returns the result TensorVar.
 """
 
 from __future__ import annotations
 
+from compiler.aten.isa_builder import IsaBuilder, fp, gp
+
 
 def softmax_plena(prog, input_var, scale: float = 1.0):
     """
-    PLENA backend: online softmax (numerically stable, row-wise).
+    PLENA backend: single-pass, per-row, register-direct row-wise softmax.
 
-    Algorithm (matches fpvar_softmax_test.py):
-        Initialize: m_old[row] = -inf, l_old[row] = 0
-        for row in range(mlen):
-            m_old_saved = m_old[row]
-            S[row] *= scale
-            row_max = max(S[row])
-            m_old[row] = max(m_old[row], row_max)  # m_curr
-            m_res[row] = exp(m_old_saved - m_curr)
-            S[row] -= m_curr
-            P[row] = exp(S[row])
-            sum_p = sum(P[row])
-            l_old[row] = l_old[row] * m_res[row] + sum_p
-        P[row] /= l_old[row]  # final normalization
+    The input is one (mlen, mlen) block whose rows each fit in a single vector
+    tile (mlen == VLEN), so the whole softmax can be computed in one pass over
+    the row without the streaming running-max / running-sum machinery that the
+    flash-attention path (``_online_softmax_asm``) needs for multi-tile
+    sequences.
+
+    This mirrors the hand-written softmax_asm template EXACTLY: it emits, per
+    row, the 6-op sequence with the reduction results (max, sum) held in FP
+    REGISTERS and fed DIRECTLY to their immediate vector consumers -- no
+    ``S_ST_FP`` / ``S_LD_FP`` round-trip through FP SRAM:
+
+        V_RED_MAX f2, row, 0          ; row max -> FP register f2
+        V_SUB_VF  row, row, f2, 0, 0  ; x - max   (f2 used directly)
+        V_EXP_V   row, row, 0         ; exp(x - max)
+        S_ADD_FP  f3, f0, f0          ; clear sum accumulator
+        V_RED_SUM f3, row             ; row sum -> FP register f3
+        S_RECI_FP f3, f3, 0           ; 1 / sum   (in register)
+        V_MUL_VF  row, row, f3, 0     ; normalize (f3 used directly)
+
+    Batching the reductions and round-tripping max/sum through FP SRAM (the
+    ``tile_row_max`` / ``tile_row_sum`` helpers) trips a reduction-result
+    capture hazard on back-to-back reductions (the stored sum comes out ~half
+    on alternate rows -> result ~2x golden). Keeping the reduction result in a
+    register and consuming it immediately avoids that hazard entirely.
+
+    No scalar running-max (no ``S_MAX_FP``, which is a silent no-op in RTL) and
+    no correction-factor bookkeeping.
 
     Args:
         prog:      PlenaCompiler instance (compilation context)
         input_var: BatchVar or VRAMMatrixVar — the input matrix in VRAM
                    Shape: (mlen, mlen)
-        scale:     Multiplicative scale applied before softmax (default 1.0)
+        scale:     Multiplicative scale applied before softmax (default 1.0).
+                   When != 1.0 it is read from FP SRAM slot 1.
 
     Returns:
-        The same input_var (in-place modification) after softmax is applied.
-        The result is stored in S (VRAM), which is also returned.
+        The output VRAM matrix ``S`` (the softmax result is written in place
+        into S, tile (0, 0), across all mlen rows).
 
     Note:
-        fp_preload layout expected: [0]=0.0, [1]=scale, [2]=-inf
-        The caller must set fp_preload = [0.0, scale, float('-inf'), ...]
+        fp_preload layout expected: [0]=0.0, [1]=scale.
     """
     mlen = prog.mlen
 
@@ -52,23 +65,37 @@ def softmax_plena(prog, input_var, scale: float = 1.0):
     # simulator, matching the previous implementation's S += input behavior.
     prog.vram_add(S, input_var)
 
-    # Reuse the flash-attention online-softmax state layout instead of
-    # allocating separate FPVars for m_old/m_res/l_old/inv_l.  At MLEN=256 the
-    # old FPVar-heavy lowering needed 1028 slots plus reserved constants and
-    # overflowed the 1024-slot FPRAM.  The shared layout uses exactly 3*MLEN
-    # slots at _ONLINE_SOFTMAX_FPSRAM_BASE and scalar FP registers for temps.
-    fp_sram_start = prog._ONLINE_SOFTMAX_FPSRAM_BASE
-    prog.emit(prog._reset_fpsram_asm(fp_sram_start, mlen, 2))  # m_old = -inf
-    prog.emit(prog._reset_fpsram_asm(fp_sram_start + 2 * mlen, mlen, 0))  # l_old = 0
-    prog.emit(
-        prog._online_softmax_asm(
-            mlen=mlen,
-            s_address=prog.get_vram_addr(S.name),
-            m_start_address=fp_sram_start,
-            scale=scale,
-            rows=mlen,
-        )
-    )
-    prog.final_scale_o(0, S, rows=mlen)
+    s_addr = prog.get_vram_addr(S.name)
+
+    # f2 = row max, f3 = row sum; each reduction result stays in its register and
+    # feeds the next vector op directly (no S_ST_FP/S_LD_FP round-trip).
+    fp_max = 2
+    fp_sum = 3
+
+    (gp_row,) = prog._reg.allocate_gp(1)
+    try:
+        asm = IsaBuilder().comment("=== Single-pass per-row softmax on S (register-direct) ===")
+        for row in range(mlen):
+            row_addr = s_addr + row * mlen
+            asm.comment(f"row {row}: softmax over VRAM[{row_addr}:{row_addr + mlen}]")
+            asm.instr("S_ADDI_INT", gp(gp_row), gp(0), row_addr)
+
+            # Optional pre-scale: row *= scale (FP SRAM slot 1 holds scale).
+            if scale != 1.0:
+                asm.instr("S_LD_FP", fp(fp_max), gp(0), 1)
+                asm.instr("V_MUL_VF", gp(gp_row), gp(gp_row), fp(fp_max), 0)
+
+            # max -> sub -> exp -> (clear acc) sum -> reci -> mul, register-direct
+            asm.instr("V_RED_MAX", fp(fp_max), gp(gp_row), 0)
+            asm.instr("V_SUB_VF", gp(gp_row), gp(gp_row), fp(fp_max), 0, 0)
+            asm.instr("V_EXP_V", gp(gp_row), gp(gp_row), 0)
+            asm.instr("S_ADD_FP", fp(fp_sum), fp(0), fp(0))
+            asm.instr("V_RED_SUM", fp(fp_sum), gp(gp_row))
+            asm.instr("S_RECI_FP", fp(fp_sum), fp(fp_sum), 0)
+            asm.instr("V_MUL_VF", gp(gp_row), gp(gp_row), fp(fp_sum), 0)
+
+        prog.emit(asm)
+    finally:
+        prog._reg.free_gp([gp_row])
 
     return S

aten/plena/isa_matrix.py

@@ -294,7 +294,12 @@ def vram_sub_projection_asm(
             result_vram_addr=result_vram_addr,
             full_batch=full_batch,
             num_hidden_blocks=num_hidden_blocks,
-            mat_col_stride=self.blen,
+            # Row-aligned column stride: the matrix SRAM read is row-granular
+            # (raddr >> log2(MLEN)), so successive BLEN-column output sub-tiles
+            # must be a full MLEN row apart, else they collapse onto the same
+            # MRAM rows and the output columns replicate ([a,b,c,d]x4). blen
+            # alone is sub-row and only works on an element-granular model.
+            mat_col_stride=self.blen * self.mlen,
             transposed=False,
             gp_regs=gp_regs,
             caller_name="vram_sub_projection_asm",