Fix bf16 binary search in topk_mask

tallax/vllm/binary_search.py

@@ -5,6 +5,7 @@
 """
 
 from collections.abc import Callable
+import functools
 import jax
 import jax.numpy as jnp
 from jax import lax
@@ -50,14 +51,18 @@ def monotonic_u32_to_f32(x: jax.Array) -> jax.Array:
   return lax.bitcast_convert_type(x_bits, jnp.float32)
 
 
-def interp(l: jax.Array, r: jax.Array) -> jax.Array:
+def interp(l: jax.Array, r: jax.Array, underlying_dtype=None) -> jax.Array:
   """Interpolate between two float32 values in the monotonic u32 space.
 
   Computes the midpoint in uint32 space (avoiding overflow) and converts back.
 
   Args:
     l: Left boundary (float32 or int32)
     r: Right boundary (float32 or int32)
+    underlying_dtype: If the values represent a lower-precision dtype (e.g.
+      bfloat16) cast to float32, pass the original dtype here. The midpoint
+      will be snapped to that dtype's representable grid by zeroing the
+      unused lower mantissa bits in f32.
 
   Returns:
     Midpoint value (float32 or int32)
@@ -73,6 +78,10 @@ def interp(l: jax.Array, r: jax.Array) -> jax.Array:
   pivot = (l // 2) + (r // 2) + ((l & one) + (r & one)) // 2
   if floating:
     pivot = monotonic_u32_to_f32(pivot)
+    if underlying_dtype == jnp.bfloat16:
+      # bf16 occupies the upper 16 bits of f32; zero the lower 16 to snap to bf16 grid
+      pivot_bits = lax.bitcast_convert_type(pivot, jnp.uint32)
+      pivot = lax.bitcast_convert_type(pivot_bits & jnp.uint32(0xFFFF0000), jnp.float32)
   return pivot
 
 
@@ -101,6 +110,7 @@ def binary_search(
   lower_bound: jax.Array = None,
   upper_bound: jax.Array = None,
   num_iter: int | None = None,
+  underlying_dtype=None,
 ) -> jax.Array:
   """Find threshold using binary search with custom predicate.
 
@@ -112,18 +122,22 @@ def binary_search(
     lower_bound: Lower bound for search
     upper_bound: Upper bound for search
     num_iter: Number of iterations (required, typically dtype.itemsize * 8)
+    underlying_dtype: If searching over values from a lower-precision dtype
+      (e.g. bfloat16) cast to float32, pass the original dtype so midpoints
+      are snapped to that dtype's representable grid.
 
   Returns:
     Tuple of (lower_bound, threshold, next_pivot) from final search state
   """
+  _interp = functools.partial(interp, underlying_dtype=underlying_dtype)
   # Binary search finds LARGEST value where predicate is FALSE
 
   @jax.jit
   def loop_body(state):
     l, r, pivot = state
 
     # We pre-compute two possible pivots of next iter to reduce latency, then select later.
-    next_pivots = (interp(l, pivot), interp(pivot, r))
+    next_pivots = (_interp(l, pivot), _interp(pivot, r))
 
     # Evaluate predicate at midpoint
     predicate_true = predicate_fn(pivot)
@@ -144,7 +158,7 @@ def cond(state):
     return jnp.any(next_pivot != l)
 
   # Run binary search, user decides if they need l or r
-  state = (lower_bound, upper_bound, interp(lower_bound, upper_bound))
+  state = (lower_bound, upper_bound, _interp(lower_bound, upper_bound))
   if num_iter is not None:
     return jax.lax.fori_loop(0, num_iter, lambda _, carry: loop_body(carry), init_val=state)
   else:

tallax/vllm/topk_mask.py

@@ -201,6 +201,7 @@ def topk_mask_ref_inputs(
     predicate_fn,
     *(jnp.full(bound_shape, v, logits.dtype) for v in (finfo.min, finfo.max)),
     num_iter=logits_ref.dtype.itemsize * 8, # 32 for f32, 16 for bf16
+    underlying_dtype=logits_ref.dtype,
   )
 
   assert logits.shape[1] % NUM_LANES == 0