Skip to content

fix: dot/schedule_bench A_stride_m for transpose_a kernels#10004

Open
kasper0406 wants to merge 1 commit intogoogle:masterfrom
kasper0406:kn/dot-bench-fix
Open

fix: dot/schedule_bench A_stride_m for transpose_a kernels#10004
kasper0406 wants to merge 1 commit intogoogle:masterfrom
kasper0406:kn/dot-bench-fix

Conversation

@kasper0406
Copy link
Copy Markdown

Summary

  • schedule_bench was passing a_stride_m (the {i, tile_k} intra-row stride) as the 6th positional arg to kernels that set dot_flag::transpose_a, but for those kernels the 6th arg is consumed as the stride along the k1 dimension of the packed tensor (the advance per inner-k step)
  • subgraph/dot.cc::call_kernel already does the correct swap (transposed_a ? a_k_strides[0] : a_stride_m). This change mirrors it in schedule_bench so the bench exercises the same work the production path executes
  • The built-in correctness check (A = B = 1, assert c == k everywhere) can't catch this class of bug because the wrong-stride loads still return 1s. Benchmark GFLOPS numbers reported against this bench were inflated by artificial cache hits on overlapping reads

Test plan

  • bazel test //ynnpack/kernels/dot/... passes (including schedule_bench_test, schedule_test, consistent_arithmetic_test, test)

@google-cla
Copy link
Copy Markdown

google-cla Bot commented Apr 18, 2026

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@kasper0406 kasper0406 marked this pull request as ready for review April 19, 2026 13:16
@kasper0406 kasper0406 marked this pull request as draft April 19, 2026 13:16
Correctly supply the A stride for transpose_a kernels
@kasper0406 kasper0406 marked this pull request as ready for review April 19, 2026 13:31
Copy link
Copy Markdown
Collaborator

@dsharlet dsharlet left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the fix!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants