Skip to content
Merged
Changes from all commits
Commits
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
9 changes: 6 additions & 3 deletions aiter/ops/flydsl/kernels/preshuffle_gemm.py
Original file line number Diff line number Diff line change
Expand Up @@ -256,8 +256,9 @@ def kernel_gemm(

# Bound A (read) and C (store) to the actual M extent so blocks covering
# rows past M (ragged M) drop their OOB loads/stores at the descriptor
# instead of faulting / writing past the allocation. B and scales are
# exact-multiple in N and stay max_size.
# instead of faulting / writing past the allocation. B and scale_b/bias
# are per-N (exact multiple) and stay max_size; scale_a is per-row (M) and
# is bounded the same way below (see its epilogue load).
gA = fx.rocdl.make_buffer_tensor(
arg_a,
max_size=False,
Expand Down Expand Up @@ -638,7 +639,9 @@ def load_epi_operands():
for ni in range_constexpr(num_acc_n)
]
scale_a_rsrc = buffer_ops.create_buffer_resource(
arg_scale_a, max_size=True
arg_scale_a,
max_size=False,
num_records_bytes=fx.Int64(i32_m) * fx.Int64(4),
)
s_a = [
Vec(
Expand Down
Loading