From 15bacc34e0274be7060917d70015dab64d059f0b Mon Sep 17 00:00:00 2001 From: Loki Chen Date: Sat, 15 Aug 2026 08:27:32 -0700 Subject: [PATCH] weight-sync: fast-path host detile for the bf16 two-tile TPU layout 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. --- tpu_sync/weight_sync/tiling_utils.cc | 73 +++++++++++++++++++++++ tpu_sync/weight_sync/tiling_utils_test.cc | 40 +++++++++++++ 2 files changed, 113 insertions(+) diff --git a/tpu_sync/weight_sync/tiling_utils.cc b/tpu_sync/weight_sync/tiling_utils.cc index 868e9b0e..eccd9ff2 100644 --- a/tpu_sync/weight_sync/tiling_utils.cc +++ b/tpu_sync/weight_sync/tiling_utils.cc @@ -63,6 +63,70 @@ bool IsStandardRowMajorTiled(const xla::Shape& shape, return true; } +// --- BF16 {Tile(8,128),Tile(2,1)} fast detile (PR-E) --------------------- +// Real bf16 weights carry a nested layout {Tile(8,128), Tile(2,1)} rather +// than the single-tile layout IsStandardRowMajorTiled() fast-paths, so they +// otherwise fall through to the scalar ForEachIndexNoStatus() detile. The +// inner Tile(2,1) interleaves pairs of adjacent rows: within each 8x128 tile +// the physical element order is [rowpair(4)][col(128)][parity(2)], i.e. the +// two rows of a pair are stored with a fixed stride of 2. This is a pure +// fixed-stride-2 row de-interleave, so it can be done with tight contiguous +// output writes and one strided read per element -- byte-exact with, but far +// faster than, the scalar per-element index-math fallback. +// +// Matches ONLY BF16, rank-2, standard row-major (minor_to_major = {1,0}), +// tiles == {Tile(8,128), Tile(2,1)}. Every other layout returns false and +// falls through to the existing paths unchanged. +bool IsBf16SubTiled_8_128_2_1(const xla::Shape& shape, + const xla::Layout& layout) { + if (shape.element_type() != xla::PrimitiveType::BF16) return false; + if (shape.dimensions().size() != 2) return false; + if (layout.minor_to_major().size() != 2) return false; + if (layout.minor_to_major(0) != 1 || layout.minor_to_major(1) != 0) { + return false; + } + if (layout.tiles().size() != 2) return false; + const xla::Tile& t0 = layout.tiles(0); + const xla::Tile& t1 = layout.tiles(1); + if (t0.dimensions().size() != 2 || t1.dimensions().size() != 2) return false; + if (t0.dimension(0) != 8 || t0.dimension(1) != 128) return false; + if (t1.dimension(0) != 2 || t1.dimension(1) != 1) return false; + return true; +} + +// Byte-exact fast detile for the BF16 {Tile(8,128),Tile(2,1)} layout. The +// caller has already verified the layout via IsBf16SubTiled_8_128_2_1(). +void DetileBufferBf16SubTiled(const uint8_t* src_tiled, uint8_t* dst_linear, + int64_t H, int64_t W) { + constexpr int64_t kTileH = 8; + constexpr int64_t kTileW = 128; + constexpr int64_t kTileElems = kTileH * kTileW; // 1024 + const int64_t col_tiles = (W + kTileW - 1) / kTileW; + const uint16_t* src = reinterpret_cast(src_tiled); + uint16_t* dst = reinterpret_cast(dst_linear); + + for (int64_t r = 0; r < H; ++r) { + const int64_t tr = r / kTileH; + const int64_t rin = r % kTileH; + const int64_t rowpair = rin / 2; // 0..3 + const int64_t parity = rin % 2; // 0..1 + uint16_t* dst_row = dst + r * W; + for (int64_t tc = 0; tc < col_tiles; ++tc) { + const int64_t c0 = tc * kTileW; + const int64_t cw = std::min(kTileW, W - c0); + const int64_t tile_index = tr * col_tiles + tc; + // Physical base of this row's elements within the tile: + // tile_index*1024 + rowpair*256 + parity, then stride 2 across cols. + const uint16_t* src_row = + src + tile_index * kTileElems + rowpair * (kTileW * 2) + parity; + uint16_t* out = dst_row + c0; + for (int64_t cin = 0; cin < cw; ++cin) { + out[cin] = src_row[cin * 2]; + } + } + } +} + // Dispatches row copy operations to compile-time specialized fixed-width // vector copy loops for common TPU tile row sizes (FP8, BF16, FP32 with tile_W // 128 or 8), falling back to generic std::memcpy for arbitrary dimensions. @@ -562,6 +626,15 @@ absl::Status DetileBuffer(const uint8_t* src_tiled, uint8_t* dst_linear, return DetileBufferNDOptimized(src_tiled, dst_linear, shape, layout); } + // Fast byte-exact path for the real bf16 weight layout + // {Tile(8,128), Tile(2,1)}, which otherwise hits the scalar fallback below. + // Auto-dispatched (no env gate); any non-matching layout falls through. + if (IsBf16SubTiled_8_128_2_1(shape, layout)) { + DetileBufferBf16SubTiled(src_tiled, dst_linear, shape.dimensions(0), + shape.dimensions(1)); + return absl::OkStatus(); + } + int64_t itemsize = xla::ShapeUtil::ByteSizeOfPrimitiveType(shape.element_type()); diff --git a/tpu_sync/weight_sync/tiling_utils_test.cc b/tpu_sync/weight_sync/tiling_utils_test.cc index c6186048..00f7bb3c 100644 --- a/tpu_sync/weight_sync/tiling_utils_test.cc +++ b/tpu_sync/weight_sync/tiling_utils_test.cc @@ -15,6 +15,7 @@ #include "tpu_sync/weight_sync/tiling_utils.h" #include +#include #include #include @@ -780,5 +781,44 @@ TEST(TilingUtilsTest, SpecializedRowBytes_MultiBatch_AllTypes) { } } + +// PR-E: exercises the fast byte-exact DetileBuffer path for the real bf16 +// weight layout {Tile(8,128), Tile(2,1)} on shapes that are NOT tile-aligned +// (H not a multiple of 8, W not a multiple of 128), which the scalar fallback +// previously handled. The tiled buffer is produced by the (scalar) TileBuffer +// and inverted by DetileBuffer, so equality proves the fast detile matches the +// canonical index-math result bit-for-bit, including padded tiles. +TEST(TilingUtilsTest, Bf16SubTiledFastPathByteExact) { + for (const auto& hw : std::vector>{ + {16, 256}, {8, 128}, {24, 384}, {17, 257}, {10, 130}, {1, 128}}) { + const int64_t H = hw.first; + const int64_t W = hw.second; + xla::Shape shape = xla::ShapeUtil::MakeShapeWithDenseLayout( + xla::PrimitiveType::BF16, {H, W}, {1, 0}, + {xla::Tile({8, 128}), xla::Tile({2, 1})}); + + const int64_t num_elements = H * W; + std::vector src_linear(num_elements); + for (int64_t i = 0; i < num_elements; ++i) { + src_linear[i] = static_cast(i * 7 + 1); + } + const int64_t tiled_elems = GetTiledBufferElements(shape); + std::vector dst_tiled(tiled_elems * sizeof(uint16_t), 0); + ASSERT_TRUE(TileBuffer(reinterpret_cast(src_linear.data()), + dst_tiled.data(), shape, shape.layout()) + .ok()); + + std::vector dst_linear(num_elements, 0); + ASSERT_TRUE(DetileBuffer(dst_tiled.data(), + reinterpret_cast(dst_linear.data()), + shape, shape.layout()) + .ok()); + for (int64_t i = 0; i < num_elements; ++i) { + EXPECT_EQ(dst_linear[i], src_linear[i]) + << "H=" << H << " W=" << W << " mismatch at " << i; + } + } +} + } // namespace } // namespace tpu_raiden::weight_sync