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" 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 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",