From 8da8785e35c0a052879f41616ed15cf46ec00949 Mon Sep 17 00:00:00 2001 From: booth-algo Date: Thu, 28 May 2026 15:23:01 +0100 Subject: [PATCH] Fix FFN down-projection HBM prefetch advance for multi-K-tile The down-projection weight prefetch in `_ffn_asm_with_loops` and `_ffn_asm_fused_up_gate` used `_load_large_int(a_actual, mlen*hidden_size)` to advance the HBM offset register between K-tile prefetches. This overwrote the register instead of adding to it, so: * At K_tiles >= 3 the prefetch for tiles 2..N-1 all read from HBM[base + mlen*hidden_size] instead of advancing by that stride each iteration. * At K_tiles == 2 combined with multi-outer-iter (out_size > mlen), K-tile 1 of outer iters > 0 read from the wrong absolute address because the per-outer-iter base offset got dropped. The up-projection in the same function already used `_addi_large_int` correctly; the down path now matches. Fix applied at all three down-projection prefetch sites (lines ~1333, 1699, 1777). Also: `_emit_ffn_projection_unrolled` switched the K-split V_ADD_VV accumulator loop count from integer division to `math.ceil` so the tail vector when `output_elements` is not a multiple of `vlen` is no longer truncated; `import math` added. Verified with ATen FFN testbench at MLEN=64 and MLEN=128 across hidden/inter combinations exercising single-tile, multi-K-tile, and multi-outer-iter codegen paths (all pass). Co-Authored-By: Claude Opus 4.7 --- asm_templates/ffn_asm.py | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/asm_templates/ffn_asm.py b/asm_templates/ffn_asm.py index ac55c18..87491a6 100644 --- a/asm_templates/ffn_asm.py +++ b/asm_templates/ffn_asm.py @@ -1,5 +1,7 @@ from __future__ import annotations +import math + from ._imm import addi_large_int_str as _addi_large_int from ._imm import load_large_int_str as _load_large_int from ._k_split import k_chunks as _k_chunks @@ -381,7 +383,7 @@ def _emit_ffn_projection_unrolled( chunks = _k_chunks(num_k_tiles, max_k_tiles) # Total output region size (elements) for the VRAM accumulator pass. output_elements = out_size * batch * seq_len - per_vlen_adds = output_elements // vlen + per_vlen_adds = math.ceil(output_elements / vlen) for chunk_idx, (k_start, k_count) in enumerate(chunks): lines.append( @@ -1328,7 +1330,12 @@ def _ffn_asm_with_loops( generated_code += ( f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen}\n" ) - generated_code += _load_large_int(a_actual_register, mlen * hidden_size) + generated_code += _addi_large_int( + a_actual_register, + a_actual_register, + mlen * hidden_size, + w_temp_register, + ) # Reset for compute phase generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0\n" @@ -1689,7 +1696,12 @@ def _ffn_asm_fused_up_gate( generated_code += f"; Prefetch DOWN weight tile {prefetch_idx}\n" generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{a_actual_register}, a{down_weight_hbm_offset_reg}, 1, 0\n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen}\n" - generated_code += _load_large_int(a_actual_register, mlen * hidden_size) + generated_code += _addi_large_int( + a_actual_register, + a_actual_register, + mlen * hidden_size, + w_temp_register, + ) # Downsize linear (first block already prefetched) generated_code += ( @@ -1767,7 +1779,12 @@ def _ffn_asm_fused_up_gate( for weight_col in range(num_down_weight_tiles): generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{a_actual_register}, a{down_weight_hbm_offset_reg}, 1, 0\n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen}\n" - generated_code += _load_large_int(a_actual_register, mlen * hidden_size) + generated_code += _addi_large_int( + a_actual_register, + a_actual_register, + mlen * hidden_size, + w_temp_register, + ) generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0\n" generated_code += (