Reorganize sampling kernels into clear code paths with shared utils

tallax/tax/divide_and_filter_topk/topk.py

@@ -1,6 +1,7 @@
 """Divide-and-filter top-k algorithm implementation."""
 
 import functools
+import operator
 import jax
 import jax.numpy as jnp
 from jax import jit
@@ -61,7 +62,9 @@ def nan_to_min(x):
   Returns:
     Array with NaNs replaced by minimum value.
   """
-  return jnp.where(jnp.isnan(x), get_dtype_info(x).min, x)
+  if jnp.issubdtype(x.dtype, jnp.floating):
+    x = jnp.where(jnp.isnan(x), get_dtype_info(x).min, x)
+  return x 
 
 
 def to_comparison_dtype(x):
@@ -76,7 +79,7 @@ def to_comparison_dtype(x):
   return x.astype(to_32bit_dtype(x.dtype))
 
 
-def bitonic_topk_arrays(operands, k, val_dtype=None, max_index=None):
+def bitonic_topk_arrays(operands, k, num_keys, val_dtype=None, max_index=None):
   """Top-k of arrays, packing bf16 and indices into single 32-bit dtype if possible.
 
   Args:
@@ -89,13 +92,13 @@ def bitonic_topk_arrays(operands, k, val_dtype=None, max_index=None):
     List of top-k arrays.
   """
   if val_dtype is None or max_index is None:
-    return _bitonic_topk_arrays(operands, k=k)
+    return _bitonic_topk_arrays(operands, k=k, num_keys=num_keys)
   assert len(operands) == 2 and type(operands) == list
   dtypes = [x.dtype for x in operands]
   pack = val_dtype == jnp.bfloat16 and max_index <= 2**16
   if pack:
-    operands = [pack_bf16_u16_to_i32(*operands, stable=False)]
-  operands = _bitonic_topk_arrays(operands, k=k)
+    operands = [pack_bf16_u16_to_i32(*operands, stable=(num_keys > 1))]
+  operands = _bitonic_topk_arrays(operands, k=k, num_keys=min(len(operands), num_keys))
   if pack:
     assert len(operands) == 1
     operands = list(unpack_bf16_u16_from_i32(operands[0], stable=False))
@@ -158,7 +161,7 @@ def update_bins_topk(
     for i in range(completed_k, k):
       # Exchange with stored top-k
       # Only perform the swap if the value is larger
-      mask = bubble_vals > bins_topk_vals[i]
+      mask = bubble_vals >= bins_topk_vals[i]
       bins_topk_vals[i], bubble_vals = (
         jnp.where(m, bubble_vals, bins_topk_vals[i]) for m in (mask, ~mask)
       )
@@ -174,26 +177,31 @@ def compute_idxs(i):
       jnp.int32, shape, 1
     )
 
+  num_full_slices = vocab_size // num_bins
   def loop_body(i, bins_topk_outs):
+    # Process in reverse order, due to >= swap
+    # this leads to stable index ordering
+    i = num_full_bins - 1 - i
     vals = logits[..., pl.dslice(num_bins * i, num_bins)]
     idxs = compute_idxs(i)
     return update_bins_topk(vals, idxs, *bins_topk_outs)
 
-  num_full_slices = vocab_size // num_bins
-  bins_topk_outs = unrolled_fori_loop(
-    num_full_slices,
-    loop_body,
-    (bins_topk_vals, bins_topk_idxs),
-    unroll=unroll,
-  )
-
   # Handle remaining elements if vocab_size doesn't divide num_bins
+  # We do this first as we loop the array backwards, due to >= swap this leads to stable index ordering and cases where the top-k includes -inf
   rem_vals = _extract_remainder_slice(logits, num_bins)
   if rem_vals is not None:
     # Create idxs for the final segment
     rem_idxs = compute_idxs(num_full_slices)
     # Update bins topk with the overspill
     bins_topk_outs = update_bins_topk(rem_vals, rem_idxs, *bins_topk_outs)
+
+  bins_topk_outs = unrolled_fori_loop(
+    num_full_slices,
+    loop_body,
+    (bins_topk_vals, bins_topk_idxs),
+    unroll=unroll,
+  )
+
   return bins_topk_outs
 
 
@@ -205,6 +213,7 @@ def _merge_unconverged_bins_topk(
   num_bins: int,
   m: int,
   max_k: int,
+  stable: bool = False,
 ):
   """Compute top-k from most active bins and merge with unconverged bins."""
 
@@ -223,6 +232,7 @@ def _merge_unconverged_bins_topk(
   _, sorted_bin_indices = bitonic_topk_arrays(
     [bin_vals, bin_indices],
     k=num_packed_bins,
+    num_keys=1+int(stable),
   )
   sorted_bin_indices = pad(sorted_bin_indices, (NUM_SUBLANES, NUM_LANES))
 
@@ -328,6 +338,7 @@ def _merge_unconverged_bins_topk(
       k=max_k,
       val_dtype=logits_ref.dtype,
       max_index=logits_ref.shape[1],
+      num_keys=1+int(stable),
     )
   )
 
@@ -352,6 +363,7 @@ def dynamic_topk_refs(
   bins_topm_schedule: tuple[int, ...],
   guarantee_convergence: bool,
   replace_val: float | int | None,
+  stable: bool,
 ):
   """
   Pallas kernel for computing binned top-k supersets until global top-k is guaranteed.
@@ -428,8 +440,11 @@ def _():
       # the largest m-th largest value, then top-k is guaranteed to be in bins
       # top-(m-1) collated
       pivot = bins_topm_vals[m - 1].max(-1, keepdims=True)
+      # if stable sort, we must include values at the boundary in to the bitonic stable sort 
+      # if not we just need k values greater than or equal to
+      _comp = operator.gt if stable else operator.ge
       num_larger = (
-        sum((v >= pivot) for v in bins_topm_vals[: m - 1])
+        sum(_comp(v, pivot) for v in bins_topm_vals[: m - 1])
         .astype(jnp.float32)
         .sum(-1)
       )
@@ -466,6 +481,7 @@ def _():
         num_bins=num_bins,
         m=m_final,
         max_k=max_k,
+        stable=stable,
       )
 
   # early on bins_topm_schedule are convergence checks so we go to bins-top-(m-1). For final bins-top-(m_max) for convergence guaranteed we only need to consider top-(m_max-1), if not must cover bins-top-(m_max)
@@ -526,6 +542,7 @@ def _():
           k=max_k,
           val_dtype=logits_ref.dtype,
           max_index=logits_ref.shape[1],
+          num_keys=1+int(stable),
         )
         topk_vals_ref[...], topk_idxs_ref[...] = (
           vals.astype(topk_vals_ref.dtype),
@@ -837,6 +854,7 @@ def shmap_fn(logits, k):
     "bins_topm_unroll",
     "bins_topm_schedule",
     "interpret",
+    "stable",
   ),
 )
 def topk(
@@ -848,13 +866,13 @@ def topk(
   bins_topm_unroll: int = 64,
   bins_topm_schedule: tuple[int, ...] | None = None,
   interpret: bool = False,
+  stable: bool = True,
 ):
   """
   Compute top-k element.
 
   Behavior differences to jax.lax.top_k:
       - Handling of NaNs is different to jax.lax.top_k, here NaNs are never part of top-k.
-      - Any output where k values are larger than or equal to the k'th largest value is considered valid, unlike jax.lax.top_k which in case of ties in value considers lower-index elements larger.
   If you wish exactly the same behavior, use `tallax.tax.bitonic_top_k(x, k=k, is_stable=True)` instead.
 
   Sharding is supported in either/both dimensions
@@ -869,6 +887,7 @@ def topk(
       bins_topm_schedule: Optional custom search schedule. If None, automatically
           computed.
       interpret: If True, run in CPU interpret mode (default: False).
+      stable: Whether sort should be stable, faster if False. Note: jax.lax.top_k is stable sort. (default: True)
 
   Returns:
       Tuple of (topk_vals, topk_idxs):
@@ -886,4 +905,5 @@ def topk(
     bins_topm_schedule=bins_topm_schedule,
     guarantee_convergence=True,
     interpret=interpret,
+    stable=stable,
   )

tallax/tax/utils.py

@@ -507,3 +507,70 @@ def jit_f(*args, **kwargs):
     return fs[key](*args, **kwargs)
 
   return jit_f
+
+
+def map_reduce(
+  vals,
+  map_fn=lambda x: x,
+  reduce_fn="sum",
+  num_parallel: int = 7,
+  apply_post_partial_sums_fn=None,
+):
+  """Computes unary_reduce(map_fn(vals)) efficiently using parallel chunks.
+
+  The input vals is chunked along axis 1 using NUM_LANES.
+  The map_fn is expected to already close over any required parameters.
+
+  reduce_fn can be a string ('sum', 'max', 'min') or a tuple
+  (binary_reduce, unary_reduce).
+  """
+  if reduce_fn == "sum":
+    binary_reduce, unary_reduce = (lambda x, y: x + y), (lambda x: x.sum(1, keepdims=True))
+  elif reduce_fn == "max":
+    binary_reduce, unary_reduce = jnp.maximum, lambda x: x.max(1, keepdims=True)
+  elif reduce_fn == "min":
+    binary_reduce, unary_reduce = jnp.minimum, lambda x: x.min(1, keepdims=True)
+  else:
+    binary_reduce, unary_reduce = reduce_fn
+
+  vals_shape_1 = vals.shape[1]
+
+  # Fast path for small inputs
+  if num_parallel == 0 or vals_shape_1 < NUM_LANES:
+    return unary_reduce(map_fn(vals))
+
+  num_chunks = vals_shape_1 // NUM_LANES
+  chunks = [
+    vals[:, i * NUM_LANES : (i + 1) * NUM_LANES] for i in range(num_chunks)
+  ]
+  num_parallel = min(num_parallel, len(chunks))
+
+  # Compute partial results
+  partial_results = [map_fn(chunk) for chunk in chunks[:num_parallel]]
+  for i, chunk in enumerate(chunks[num_parallel:]):
+    partial_results[i % num_parallel] = binary_reduce(
+      partial_results[i % num_parallel], map_fn(chunk)
+    )
+
+  if apply_post_partial_sums_fn is not None:
+    partial_results = list(map(apply_post_partial_sums_fn, partial_results))
+
+  # Tree reduction
+  while len(partial_results) > 1:
+    n = len(partial_results)
+    for i in range(n // 2):
+      partial_results[i] = binary_reduce(
+        partial_results[i], partial_results[i + (n + 1) // 2]
+      )
+    partial_results = partial_results[: (n + 1) // 2]
+
+  full_result = unary_reduce(partial_results[0])
+
+  has_remainder = vals_shape_1 % NUM_LANES
+  if has_remainder:
+    remainder_vals = vals[:, num_chunks * NUM_LANES :]
+    full_result = binary_reduce(
+      full_result, unary_reduce(map_fn(remainder_vals))
+    )
+
+  return full_result