Skip to content

Allow mmap for blob-aligned checkpoints - #1032

Open
Mikyx-1 wants to merge 2 commits into
google:devfrom
Mikyx-1:fix/mmap-partial-page-length
Open

Mikyx-1 wants to merge 2 commits into
google:devfrom
Mikyx-1:fix/mmap-partial-page-length

Conversation

@Mikyx-1

@Mikyx-1 Mikyx-1 commented Sep 13, 2026

Copy link
Copy Markdown

Fixes #1031.

Checkpoints padded to 256-byte blob boundaries currently fall back to reading when their sizes are not multiples of the system page size. Relax the loader guard to kBlobAlign so these checkpoints can use mmap while retaining the alignment requirement for kernels that read ahead.

  • Share kBlobAlign between the blob store and weight loader, and correct the comment about page-aligned file lengths.
  • Keep the read fallback for lengths that are not multiples of 256, including explicitly requested mapping.
  • Test a real blob-aligned partial-page mapping, explicit and automatic mode selection, unaligned fallback, and BF16 conversion precedence.

Validation

  • weights_test: 6 tests passed.
  • blob_store_test: 2 tests passed.
  • git diff --check and formatting checks passed.

@jan-wassenberg

Copy link
Copy Markdown
Member

Thanks, I agree mmap and Windows will not fail with an unaligned length.
But we still have a potential issue: our kernels can read multiple vectors ahead. If files are unaligned, say length 4090, most HW/OSes will round up to 4 KiB, but this is insufficient for loading a full vector.
How about we relax the alignment requirement to kBlobAlign (256)?

@Mikyx-1 Mikyx-1 changed the title Allow mmap for non-page-aligned files Allow mmap for blob-aligned checkpoints Sep 15, 2026
@Mikyx-1

Mikyx-1 commented Sep 15, 2026

Copy link
Copy Markdown
Author

Thanks @jan-wassenberg! Updated the PR to require kBlobAlign (256-byte) alignment, as you suggested, and added regression coverage; all 8 loader and blob-store tests pass. Could you take another look when you have a chance?

@jan-wassenberg jan-wassenberg left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Nice, thanks for making the change!

@jan-wassenberg jan-wassenberg added the copybara-import Trigger Copybara for merging pull requests label Sep 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

copybara-import Trigger Copybara for merging pull requests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants