weight-sync: fast-path host detile for the bf16 two-tile TPU layout - #665
Open
lokic233 wants to merge 1 commit into
Open
weight-sync: fast-path host detile for the bf16 two-tile TPU layout#665lokic233 wants to merge 1 commit into
lokic233 wants to merge 1 commit into
Conversation
DetileBuffer fast-paths only IsStandardRowMajorTiled() (a single tile). Real bf16
weights carry a two-tile layout {Tile(8,128),Tile(2,1)} and fall to the per-element
scalar ForEachIndexNoStatus path (measured 0.067 GB/s on a 1192 MB tree).
Add IsBf16SubTiled_8_128_2_1() (matches ONLY BF16 rank-2 row-major with tiles ==
{Tile(8,128),Tile(2,1)}) + DetileBufferBf16SubTiled() (fixed-stride-2 row
de-interleave), auto-dispatched after IsStandardRowMajorTiled and before the scalar
fallback. Byte-exact; every non-matching layout falls through unchanged. No API
change, no opt-in.
Byte-exactness validated against XLA's LinearIndexForNestedTiling on 9 shapes
(tile-aligned and padded): 16x256, 8x128, 24x384, 17x257, 10x130, 1x128, 8x1, 6x3,
2x1. Full-state detile on a 1192 MB Qwen3-0.6B tree drops ~17.7s -> ~0.11s.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Removes the scalar-fallback detile for a verified bf16 TPU layout while preserving completion semantics.
DetileBufferfast-paths onlyIsStandardRowMajorTiled()(a single tile). Real bf16 weights carry a two-tile layout{Tile(8,128), Tile(2,1)}and therefore fall to the per-element scalarForEachIndexNoStatuspath (measured ~0.067 GB/s on a 1192 MB weight tree ≈ 17.7 s of full-state detile).Change (byte-exact, auto-dispatched, no API change, no opt-in)
IsBf16SubTiled_8_128_2_1()— matches only BF16, rank-2, row-major (minor_to_major = {1,0}), withtiles() == {Tile(8,128), Tile(2,1)}. Any other layout returns false.DetileBufferBf16SubTiled()— a fixed-stride-2 row de-interleave (the innerTile(2,1)interleaves adjacent row pairs).DetileBufferafterIsStandardRowMajorTiledand before the scalar fallback. Every non-matching layout falls through unchanged.Full-state detile on a 1192 MB Qwen3-0.6B tree drops from ~17.7 s to ~0.11 s.
Byte-exactness
The physical↔logical mapping was derived from XLA's
LinearIndexForNestedTilingsemantics and verified equal to the scalar reference on 9 shapes (tile-aligned and padded):16×256, 8×128, 24×384, 17×257, 10×130, 1×128, 8×1, 6×3, 2×1.A gtest (
TilingUtilsTest.Bf16SubTiledFastPathByteExact) round-tripping scalarTileBuffer→ fastDetileBufferis included in this PR. Caveat: in my offline build environment googletest could not be fetched, so the unit test was validated via an equivalent standalone harness rather than the in-repo gtest runner; CI with network access will build and run it normally.Not in this PR
int8/fp8 sub-tiled
{Tile(...), Tile(4,1)}layouts (different element size / stride) — left as a documented follow-up. This PR is independent of and separate from the D2H event-completion correctness fix.Authored with agentic assistance (Navi) and validated on TPU silicon before submission.