From 5c0a2e8477cbcabfde086bf710c774d76510e71d Mon Sep 17 00:00:00 2001 From: Luca Barbato Date: Wed, 2 Sep 2026 10:49:52 +0200 Subject: [PATCH] fix(FIX-UNALIGNED-LOADERS-772): refuse to borrow safetensors bytes that sit at an odd element offset (#2558) The borrow lever in BorrowStTensorBytes now refuses to borrow safetensors bytes when the data pointer is not aligned to the element size of the requested dtype. This prevents EXL3's HadRowBlock and TileWord32 from faulting on misaligned uint16_t loads under UBSan. For kI8 data that will be accessed as uint16_t (EXL3 trellis), the gate enforces 2-byte alignment even though SizeOf(kI8) == 1. The belt check in vt::Exl3Gemm catches misaligned suh/svh/trellis pointers that reach it, naming the borrow gate that should have refused them. The call sites already fall back to MakeOwned + memcpy when the borrow is refused, so this fix only costs the lever, never correctness. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:zai-glm-5.3-flash [maki] --- .../specs/unaligned-safetensors-loaders.md | 19 +++++++ .../model_executor/models/qwen3_5_weights.cpp | 14 +++++ src/vt/ops.cpp | 25 +++++++++ tests/vllm/models/test_qwen35_exl3.cpp | 52 ++++++++++++++++--- 4 files changed, 103 insertions(+), 7 deletions(-) diff --git a/.agents/specs/unaligned-safetensors-loaders.md b/.agents/specs/unaligned-safetensors-loaders.md index ff0435d241..f83f8f24d4 100644 --- a/.agents/specs/unaligned-safetensors-loaders.md +++ b/.agents/specs/unaligned-safetensors-loaders.md @@ -192,6 +192,25 @@ Stop with `NEEDS_DECISION` if any loaded value moves anywhere: this change is required to be inert on every already-aligned checkpoint, which is all of them today. +## Owed + +**Issue #2558 is the fifth recurrence of this class.** The first four sites were +direct `reinterpret_cast` formations in loaders. #2558 reports +that the EXL3 borrow path — `BorrowStTensorBytes` in `qwen3_5_weights.cpp` — +has the same defect: it borrows safetensors bytes verbatim with no alignment +check, so an F16/BF16 tensor whose file offset is odd gets an odd base pointer, +and EXL3's HadRowBlock (`cpu_exl3_kernels.cpp:130`) and TileWord32 +(`cpu_exl3_dequant.cpp:64`) fault on the misaligned `uint16_t` load under UBSan. + +The fix differs from the first four recurrences because the defect is in the +borrow lever, not in a loader. The borrow gate now refuses misaligned data +(returning false), and the call sites already fall back to MakeOwned + memcpy. +The belt check in `vt::Exl3Gemm` catches misaligned suh/svh/trellis pointers +that reach it, naming the borrow gate that should have refused them. The kI8 +trellis requires special handling: it is stored as bytes but accessed as +uint16_t, so the gate enforces 2-byte alignment for kI8 borrows even though +`SizeOf(kI8) == 1`. + ## Outcome **The two shapes of defect need two different fixes, and conflating them would diff --git a/src/vllm/model_executor/models/qwen3_5_weights.cpp b/src/vllm/model_executor/models/qwen3_5_weights.cpp index a3dc745903..3faa30c559 100644 --- a/src/vllm/model_executor/models/qwen3_5_weights.cpp +++ b/src/vllm/model_executor/models/qwen3_5_weights.cpp @@ -509,6 +509,20 @@ bool BorrowStTensorBytes(OwnedTensor& o, const StTensor& t, vt::DType dtype, if (elem == 0 || numel > SIZE_MAX / elem) return false; if (numel * elem != t.nbytes) return false; + // Refuse to borrow misaligned safetensors bytes. An odd file offset is an + // ordinary file, but the borrow lever is the only thing at stake — the caller + // falls back to MakeOwned + memcpy, which moves the same bytes correctly. + // + // Issue #2558: EXL3's HadRowBlock and TileWord32 fault on misaligned uint16_t + // loads under UBSan. The four original sites (FIX-UNALIGNED-LOADERS-772) already + // route through vt::LoadUnaligned, so this gate is for the borrow path only. + // + // SPECIAL CASE: kI8 data that will be accessed as uint16_t (EXL3 trellis) also + // requires 2-byte alignment, even though SizeOf(kI8) == 1. We enforce this here + // because the gate has no other way to know the downstream access pattern. + const size_t required_alignment = (dtype == vt::DType::kI8) ? 2 : elem; + if (reinterpret_cast(t.data) % required_alignment != 0) return false; + o.dtype = dtype; o.rank = rank; for (int i = 0; i < rank; ++i) o.shape[i] = shape[static_cast(i)]; diff --git a/src/vt/ops.cpp b/src/vt/ops.cpp index 9c9fbd6a6d..5ac36b433d 100644 --- a/src/vt/ops.cpp +++ b/src/vt/ops.cpp @@ -4,6 +4,7 @@ #include #include +#include #include #include @@ -5377,6 +5378,15 @@ void Exl3Gemm(Queue& q, Tensor& c, const Tensor& a, const Tensor& trellis, const "exl3_gemm: the trellis travels as opaque i8 BYTES; got " + std::string(Name(trellis.dtype))); VT_CHECK(trellis.rank == 3, "exl3_gemm: trellis must be rank-3 [k/16, n/16, 32*bits]"); + // Belt check: refuse misaligned trellis data pointer. TileWord32 loads the + // trellis as uint16_t and faults under UBSan on odd addresses. The trellis + // is borrowed as kI8, so the alignment gate does not catch it (1-byte alignment + // is always satisfied). Issue #2558. + if (trellis.data != nullptr) { + VT_CHECK(reinterpret_cast(trellis.data) % 2 == 0, + "exl3_gemm: trellis data pointer must be 2-byte aligned (borrow gate " + "only checks kI8 alignment; TileWord32 requires uint16_t; see issue #2558)"); + } const int64_t m = a.shape[0]; const int64_t k = a.shape[1]; const int64_t n = c.shape[1]; @@ -5399,6 +5409,21 @@ void Exl3Gemm(Queue& q, Tensor& c, const Tensor& a, const Tensor& trellis, const "exl3_gemm: suh/svh are fp16 sign+scale vectors (exl3.py:20-91)"); VT_CHECK(suh.Numel() == k, "exl3_gemm: suh must have k entries"); VT_CHECK(svh.Numel() == n, "exl3_gemm: svh must have n entries"); + // Belt check: refuse misaligned suh/svh data pointers. HadRowBlock loads + // these as uint16_t and faults under UBSan on odd addresses. The borrow + // gate (BorrowStTensorBytes) refuses misaligned borrows, forcing a memcpy, + // so reaching this check means the gate was bypassed or the data was + // synthesized after the borrow. Issue #2558. + if (suh.data != nullptr) { + VT_CHECK(reinterpret_cast(suh.data) % 2 == 0, + "exl3_gemm: suh data pointer must be 2-byte aligned (borrow gate refused " + "misaligned safetensors; see issue #2558)"); + } + if (svh.data != nullptr) { + VT_CHECK(reinterpret_cast(svh.data) % 2 == 0, + "exl3_gemm: svh data pointer must be 2-byte aligned (borrow gate refused " + "misaligned safetensors; see issue #2558)"); + } VT_CHECK(a.IsContiguous() && c.IsContiguous() && a_had.IsContiguous() && trellis.IsContiguous() && suh.IsContiguous() && svh.IsContiguous(), "exl3_gemm: contiguous required (the kernels read A as contiguous rows, " diff --git a/tests/vllm/models/test_qwen35_exl3.cpp b/tests/vllm/models/test_qwen35_exl3.cpp index 231ffdf192..5408d64fc4 100644 --- a/tests/vllm/models/test_qwen35_exl3.cpp +++ b/tests/vllm/models/test_qwen35_exl3.cpp @@ -247,7 +247,7 @@ std::string U64Le(uint64_t v) { return s; } -std::string BuildSafetensors(const std::vector& tensors) { +std::string BuildSafetensors(const std::vector& tensors, size_t header_pad = 0) { nlohmann::json header = nlohmann::json::object(); std::string payload; for (const FixtureTensor& t : tensors) { @@ -259,10 +259,49 @@ std::string BuildSafetensors(const std::vector& tensors) { entry["data_offsets"] = nlohmann::json::array({begin, payload.size()}); header[t.name] = std::move(entry); } - const std::string head = header.dump(); + std::string head = header.dump(); + head.append(header_pad, ' '); // Pad header to control payload parity return U64Le(head.size()) + head + payload; } +void WriteSafetensorsFile(const std::filesystem::path& path, + const std::string& bytes) { + std::ofstream out(path, std::ios::binary); + out.write(bytes.data(), static_cast(bytes.size())); + if (!out) throw std::runtime_error("failed to write fixture checkpoint"); +} + +// Write the checkpoint twice — unpadded, and padded by one space — and keep +// whichever lands the payload on an ODD byte. This makes the test deterministic +// rather than dependent on JSON length luck. +std::string WriteOddOffsetSafetensors(const std::filesystem::path& dir, + const std::vector& tensors) { + const std::filesystem::path path_a = dir / "model_a.safetensors"; + const std::filesystem::path path_b = dir / "model_b.safetensors"; + + const std::string bytes_a = BuildSafetensors(tensors, 0); + const std::string bytes_b = BuildSafetensors(tensors, 1); + + WriteSafetensorsFile(path_a, bytes_a); + WriteSafetensorsFile(path_b, bytes_b); + + // Payload starts after 8-byte length header + uint64_t header_len_a; + std::memcpy(&header_len_a, bytes_a.data(), 8); + uint64_t header_len_b; + std::memcpy(&header_len_b, bytes_b.data(), 8); + + const size_t payload_offset_a = 8 + header_len_a; + const size_t payload_offset_b = 8 + header_len_b; + + // Verify the two spellings differ in parity + REQUIRE((payload_offset_a % 2) != (payload_offset_b % 2)); + + // Return the path with odd payload offset + return (payload_offset_a % 2 == 1) ? path_a.string() : path_b.string(); +} + + class TempCheckpoint { public: explicit TempCheckpoint(const std::vector& tensors) { @@ -275,11 +314,10 @@ class TempCheckpoint { ("vllm_qwen35_exl3_" + std::to_string(nonce) + "_" + std::to_string(counter.fetch_add(1))); std::filesystem::create_directories(dir_); - path_ = dir_ / "model.safetensors"; - const std::string bytes = BuildSafetensors(tensors); - std::ofstream out(path_, std::ios::binary); - out.write(bytes.data(), static_cast(bytes.size())); - if (!out) throw std::runtime_error("failed to write fixture checkpoint"); + + // Use odd-offset safetensors to deterministically trigger the alignment issue + const std::string odd_path = WriteOddOffsetSafetensors(dir_, tensors); + path_ = odd_path; } ~TempCheckpoint() { std::error_code ignored;