From 790489bb500a697b312669f831bf64a1b8e9e392 Mon Sep 17 00:00:00 2001 From: booth-algo Date: Sat, 4 Jul 2026 22:01:48 +0000 Subject: [PATCH 1/3] fix(matrix): row-align non-transposed weight-column stride vram_sub_projection (non-transposed M_MM) strode the weight-column sub-tile base by blen (sub-row). The RTL matrix SRAM read is row-granular (raddr >> log2(MLEN)), so successive BLEN-column output sub-tiles collapsed onto the same MRAM rows and the output columns replicated ([a,b,c,d]x4) -> ATen linear = 27% on RTL. Stride by blen*mlen (one full row per sub-tile, matching the projection_asm template and the transposed M_TMM path). ATen linear 27% -> 95.17% on RTL. Prefetch is block-granular (mlen*mlen), so no prefetch change is needed. Note: the transactional emulator's M_MM weight model is element-granular (matrix_machine.rs asserts mat_offset < mlen) and rejects this row-aligned stride; that emulator model is infaithful to the RTL and is a separate fix. --- aten/plena/isa_matrix.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/aten/plena/isa_matrix.py b/aten/plena/isa_matrix.py index a65e90a..1196c1e 100644 --- a/aten/plena/isa_matrix.py +++ b/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", From 5f058c663bbf280c52624c70e13b31f172f26959 Mon Sep 17 00:00:00 2001 From: booth-algo Date: Sat, 4 Jul 2026 22:32:46 +0000 Subject: [PATCH 2/3] fix(ffn): row-align weight-read pointer in FFN projection _emit_ffn_projection_chunk advanced the weight-read pointer by blen between sub-column tiles. The RTL matrix SRAM read is row-granular (raddr >> log2(MLEN)), so +blen collapsed sub-columns 1..mlen/blen-1 onto MRAM row 0 (output cols 4-15 = 0). Stride the weight pointer by blen*mlen (row-aligned, matching projection_asm); the output-write pointer stays at +blen. FFN UP projection 7.6% -> 94.3% on RTL. Same class as the vram_sub_projection mat_col_stride fix. --- asm_templates/ffn_asm.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/asm_templates/ffn_asm.py b/asm_templates/ffn_asm.py index 04a2e56..f7b158b 100644 --- a/asm_templates/ffn_asm.py +++ b/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" From e8c150710aaa4781f69c8fdca5d788e4b0a37c38 Mon Sep 17 00:00:00 2001 From: booth-algo Date: Sun, 5 Jul 2026 23:23:11 +0000 Subject: [PATCH 3/3] fix(softmax): single-pass register-direct row softmax Rework softmax_plena to compute each row in one vector pass, keeping the reduction results in FP registers and feeding them directly to the next vector op (V_RED_MAX f2 -> V_SUB_VF ..f2 ; V_RED_SUM f3 -> S_RECI_FP f3 -> V_MUL_VF ..f3). Drops the online/streaming path that relied on S_MAX_FP and round-tripped max/sum through FP SRAM, which tripped a reduction-result capture hazard (softmax ATen 51.95% -> 100%). Attention's online-softmax is unchanged. --- aten/ops/plena/softmax_ops.py | 107 +++++++++++++++++++++------------- 1 file changed, 67 insertions(+), 40 deletions(-) diff --git a/aten/ops/plena/softmax_ops.py b/aten/ops/plena/softmax_ops.py index 7a3a024..19214de 100644 --- a/aten/ops/plena/softmax_ops.py +++ b/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