[Perf][Frontend] Vectorize Cohere binary embedding bit-packing (#52277)
Signed-off-by: Fangchen Li <fangchen.li@outlook.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Fangchen Li committed
Aug 14, 2026 at 02:15 UTC
103c419f048474d757c752524abb1e155dd01b9a
2 files changed
+57
-24
tests/entrypoints/pooling/embed/test_protocol.py
+35
@@ -75,6 +75,41 @@ class TestBuildTypedEmbeddingsBinary:
75
# 0.0 >= 0 is True, so bit=1 for all => 127 (signed)
76
assert result.binary[0] == [127]
77
78
+ def test_negative_zero_treated_as_positive(self):
79
+ """Packing follows ``value >= 0``, not the IEEE-754 sign bit.
80
+
81
+ ``-0.0 >= 0`` is True, so these bits must be set. An implementation
82
+ reading the sign bit instead would pack -128 here.
83
+ """
84
+ result = build_typed_embeddings([[-0.0] * 8], ["binary"])
85
+ assert result.binary is not None
86
+ assert result.binary[0] == [127]
87
+
88
+ def test_negative_denormal_treated_as_negative(self):
89
+ """Sign is decided at float64 precision.
90
+
91
+ The smallest float64 denormal underflows to -0.0 in float32, where
92
+ ``>= 0`` is True; narrowing the comparison would wrongly pack 127.
93
+ """
94
+ result = build_typed_embeddings([[-5e-324] * 8], ["binary"])
95
+ assert result.binary is not None
96
+ assert result.binary[0] == [-128]
97
+
98
+ def test_empty_input(self):
99
+ result = build_typed_embeddings([], ["binary", "ubinary"])
100
+ assert result.binary == []
101
+ assert result.ubinary == []
102
+
103
+ @pytest.mark.parametrize("embs", [[0.1] * 8, [[[0.1] * 8] * 2]], ids=["1d", "3d"])
104
+ def test_non_2d_input_raises(self, embs):
105
+ """A malformed batch must raise rather than pack along the last axis.
106
+
107
+ Packing a 1D or 3D input would silently emit wrongly shaped
108
+ ``binary``/``ubinary`` values instead of reporting an error.
109
+ """
110
+ with pytest.raises(ValueError, match="2D batch"):
111
+ build_typed_embeddings(embs, ["binary"])
112
+
113
def test_non_multiple_of_8_raises(self):
114
embs = [[0.1] * 7]
115
with pytest.raises(ValueError, match="multiple of 8"):
vllm/entrypoints/pooling/embed/protocol.py
+22
-24
@@ -12,6 +12,7 @@ import time
12
from collections.abc import Sequence
13
from typing import Annotated, Any, Literal, TypeAlias
14
15
+import numpy as np
16
import pybase64 as base64
17
from pydantic import BaseModel, BeforeValidator, Field, model_validator
18
@@ -298,31 +299,28 @@ def _pack_binary_embeddings(
299
) -> list[list[int]]:
300
"""Bit-pack float embeddings: positive -> 1, negative -> 0.
301
301
- Each bit is shifted left by ``7 - idx%8``, and every 8 bits are packed
302
- into one byte.
302
+ Bits are packed MSB-first, eight per byte.
303
"""
304
- result: list[list[int]] = []
305
- for embedding in float_embeddings:
306
- dim = len(embedding)
307
- if dim % 8 != 0:
308
- raise ValueError(
309
- "Embedding dimension must be a multiple of 8 for binary "
310
- f"embedding types, but got {dim}."
311
- )
312
- packed_len = dim // 8
313
- packed: list[int] = []
314
- byte_val = 0
315
- for idx, value in enumerate(embedding):
316
- bit = 1 if value >= 0 else 0
317
- byte_val += bit << (7 - idx % 8)
318
- if (idx + 1) % 8 == 0:
319
- if signed:
320
- byte_val -= _UNSIGNED_TO_SIGNED_DIFF
321
- packed.append(byte_val)
322
- byte_val = 0
323
- assert len(packed) == packed_len
324
- result.append(packed)
325
- return result
304
+ if not float_embeddings:
305
+ return []
306
+
307
+ array = np.asarray(float_embeddings, dtype=np.float64)
308
+ if array.ndim != 2:
309
+ raise ValueError(
310
+ f"Expected a 2D batch of embeddings, but got {array.ndim}D input."
311
+ )
312
+
313
+ dim = array.shape[1]
314
+ if dim % 8 != 0:
315
+ raise ValueError(
316
+ "Embedding dimension must be a multiple of 8 for binary "
317
+ f"embedding types, but got {dim}."
318
+ )
319
+
320
+ packed = np.packbits(array >= 0, axis=-1)
321
+ if signed:
322
+ packed = packed.astype(np.int16) - _UNSIGNED_TO_SIGNED_DIFF
323
+ return packed.tolist()
324
325
326
def _encode_base64_embeddings(