Skip to content

Avoid unaligned writes in CPU RandomBits#3891

Open
markjouh wants to merge 1 commit into
ml-explore:mainfrom
markjouh:fix/random-bits-unaligned-writes
Open

Avoid unaligned writes in CPU RandomBits#3891
markjouh wants to merge 1 commit into
ml-explore:mainfrom
markjouh:fix/random-bits-unaligned-writes

Conversation

@markjouh

Copy link
Copy Markdown

Summary

  • avoid unaligned 32-bit stores in the CPU RandomBits implementation
  • preserve the generated bytes and Threefry sequence by using memcpy for complete words
  • leave the existing partial-word copy path unchanged

Root cause

For batched keys, each key's output begins bytes_per_key bytes after the
previous one. When that stride is not divisible by four, casting the start of a
later key's output to uint32_t* and writing through it is undefined behavior.

The existing test random bits case covers an affected layout: vmapping five
uint16 values over two keys makes each per-key output 10 bytes wide. Before
this change, that test aborts under UBSan with a misaligned uint32_t store.

Validation

  • UBSAN_OPTIONS=halt_on_error=1:print_stacktrace=1 build-ubsan/tests/tests -tc='test random bits' (1 test, 28 assertions)
  • release C++ suite (244 tests, 3,238 assertions)
  • UBSan suite excluding two unrelated pre-existing GGUF alignment failures (242 tests, 3,187 assertions)
  • clang-format --dry-run --Werror mlx/backend/cpu/primitives.cpp
  • git diff --check

@markjouh
markjouh marked this pull request as ready for review July 22, 2026 19:38
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.

1 participant