Skip to content

Fix FFN down-projection HBM prefetch advance for multi-K-tile #49

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
GeorgeWu1204 merged 1 commit into from
May 28, 2026
Merged

Fix FFN down-projection HBM prefetch advance for multi-K-tile #49

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 21 additions & 4 deletions asm_templates/ffn_asm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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 += (
Expand Down Expand Up @@ -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 += (
Expand Down
Loading