Refactor debug output to return full intermediate arrays instead of checksums

README.md

@@ -18,7 +18,7 @@ vLLM    ███████████████ 390μs
 tallax  █ 25μs
 
      🔥 15× AVERAGE SPEEDUP
-     ⚡ 10× WORST-CASE SPEEDUP (36μs)
+     ⚡ 10× WORST-CASE SPEEDUP (40μs)
 ```
 
 #### 📦📦📦 Large Batch (128)
@@ -28,7 +28,7 @@ vLLM    ████████████████████████
 tallax  █ 150μs
 
      🔥 75× AVERAGE SPEEDUP
-     ⚡ 45× WORST-CASE SPEEDUP (240μs)
+     ⚡ 45× WORST-CASE SPEEDUP (250μs)
 ```
 
 ### 🎯 Scenario 2: Speculative Decoding Top-k

pyproject.toml

@@ -35,24 +35,6 @@ indent-width = 2
 target-version = "py39"
 
 [tool.ruff.lint]
-ignore = [
-    # Unnecessary collection call
-    "C408",
-    # Unnecessary map usage
-    "C417",
-    # Unnecessary dict comprehension for iterable
-    "C420",
-    # Object names too complex
-    "C901",
-    # Local variable is assigned to but never used
-    "F841",
-    # Class could be dataclass or namedtuple
-    "B903",
-    # Raise with from clause inside except block
-    "B904",
-    # Zip without explicit strict parameter
-    "B905",
-]
 select = [
     "B9",
     "C",

tallax/constants.py

@@ -0,0 +1,3 @@
+REPLACE_VAL = -1e12
+SAMPLING_EPS = 1e-5
+SCALE_BITS = 24

tallax/tax/divide_and_filter_topk/topk.py

@@ -64,7 +64,7 @@ def nan_to_min(x):
   """
   if jnp.issubdtype(x.dtype, jnp.floating):
     x = jnp.where(jnp.isnan(x), get_dtype_info(x).min, x)
-  return x 
+  return x
 
 
 def to_comparison_dtype(x):
@@ -80,7 +80,8 @@ def to_comparison_dtype(x):
 
 
 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.
+  """Top-k of arrays, first operand of value, second of indices, possibly with further operands.
+  Packing bf16 and indices into single 32-bit dtype if possible.
 
   Args:
     operands: List of arrays to find top-k from.
@@ -91,19 +92,36 @@ def bitonic_topk_arrays(operands, k, num_keys, val_dtype=None, max_index=None):
   Returns:
     List of top-k arrays.
   """
-  if val_dtype is None or max_index is None:
-    return _bitonic_topk_arrays(operands, k=k, num_keys=num_keys)
-  assert len(operands) == 2 and type(operands) == list
+  assert type(operands) == list
+  stable = num_keys > 1
+  if not jnp.issubdtype(operands[1].dtype, jnp.integer):
+    raise ValueError("Expected second operand to be indices, so integer dtype")
   dtypes = [x.dtype for x in operands]
-  pack = val_dtype == jnp.bfloat16 and max_index <= 2**16
+  pack = (
+    val_dtype is not None
+    and max_index is not None
+    and (val_dtype == jnp.bfloat16 and max_index <= 2**16)
+  )
+  if num_keys == 2 and not pack:
+    # sort vals descending, indices ascending so we negate indices
+    operands[1] = -operands[1]
   if pack:
-    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))
+    operands[1] = 2**16 - 1 - operands[1]
+    operands = [
+      pack_bf16_u16_to_i32(*operands[:2], stable=stable),
+      *operands[2:],
+    ]
+    num_keys = max(1, num_keys - 1)
+  operands = _bitonic_topk_arrays(operands, k=k, num_keys=num_keys)
   if pack:
-    assert len(operands) == 1
-    operands = list(unpack_bf16_u16_from_i32(operands[0], stable=False))
-    assert len(operands) == 2
-  # convert back dtypes (incase theyve changed)
+    operands = (
+      list(unpack_bf16_u16_from_i32(operands[0], stable=stable)) + operands[1:]
+    )
+    operands[1] = 2**16 - 1 - operands[1]
+  if num_keys == 2 and not pack:
+    # invert back so sort vals descending, indices ascending
+    operands[1] = -operands[1]
+  # convert back dtypes (incase they've changed)
   return [x.astype(dtype) for x, dtype in zip(operands, dtypes, strict=True)]
 
 
@@ -178,16 +196,18 @@ def compute_idxs(i):
     )
 
   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
+    i = num_full_slices - 1 - i
     vals = logits[..., pl.dslice(num_bins * i, num_bins)]
     idxs = compute_idxs(i)
     return update_bins_topk(vals, idxs, *bins_topk_outs)
 
   # 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
+  bins_topk_outs = (bins_topk_vals, bins_topk_idxs)
   rem_vals = _extract_remainder_slice(logits, num_bins)
   if rem_vals is not None:
     # Create idxs for the final segment
@@ -198,7 +218,7 @@ def loop_body(i, bins_topk_outs):
   bins_topk_outs = unrolled_fori_loop(
     num_full_slices,
     loop_body,
-    (bins_topk_vals, bins_topk_idxs),
+    bins_topk_outs,
     unroll=unroll,
   )
 
@@ -225,15 +245,22 @@ def _merge_unconverged_bins_topk(
   num_packed_bins = pl.cdiv(max_k, m) - 1
   if num_packed_bins > NUM_LANES or num_packed_bins > num_bins:
     raise NotImplementedError
-  bin_vals = bins_topm_vals_ref[:, pl.dslice((m - 1) * num_bins, num_bins)]
+  bin_vals, bin_idxs = (
+    ref[:, pl.dslice((m - 1) * num_bins, num_bins)]
+    for ref in (bins_topm_vals_ref, bins_topm_idxs_ref)
+  )
   # Use bitonic_topk_arrays descending to get bin indices ordered by contribution count
   bin_indices = jax.lax.broadcasted_iota(jnp.int32, (block_token, num_bins), 1)
-  # Sort descending by num_gt_k to get top NUM_LANES bin indices
-  _, sorted_bin_indices = bitonic_topk_arrays(
-    [bin_vals, bin_indices],
+  bin_operands = [bin_vals, bin_indices]
+  if stable:
+    # stable sort comparison using indices
+    bin_operands.insert(1, bin_idxs)
+  # Sort descending by m'th largest value of each bin to get bins which may contribute to top-k
+  sorted_bin_indices = bitonic_topk_arrays(
+    bin_operands,
     k=num_packed_bins,
-    num_keys=1+int(stable),
-  )
+    num_keys=1 + int(stable),
+  )[-1]
   sorted_bin_indices = pad(sorted_bin_indices, (NUM_SUBLANES, NUM_LANES))
 
   # Repeat first num_packed_bins values across NUM_LANES positions to create packing permutation
@@ -338,7 +365,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),
+      num_keys=1 + int(stable),
     )
   )
 
@@ -379,9 +406,9 @@ def dynamic_topk_refs(
   block_token = logits_ref.shape[0]
   shape = (block_token, bins_topm_vals_ref.shape[1])
   block_topk = bins_topm_vals_ref.shape[0]
-  assert block_topk % block_token == 0, (
-    "block_topk must be a multiple of block_token"
-  )
+  assert (
+    block_topk % block_token == 0
+  ), "block_topk must be a multiple of block_token"
 
   pid = pl.program_id(0)
 
@@ -440,11 +467,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 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(_comp(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)
       )
@@ -542,7 +569,7 @@ def _():
           k=max_k,
           val_dtype=logits_ref.dtype,
           max_index=logits_ref.shape[1],
-          num_keys=1+int(stable),
+          num_keys=1 + int(stable),
         )
         topk_vals_ref[...], topk_idxs_ref[...] = (
           vals.astype(topk_vals_ref.dtype),
@@ -567,6 +594,7 @@ def _():
     "guarantee_convergence",
     "replace_val",
     "interpret",
+    "stable",
   ),
 )
 def _top_bounded_k(
@@ -581,6 +609,7 @@ def _top_bounded_k(
   guarantee_convergence: bool = True,
   replace_val: float | int | None = None,
   interpret: bool = False,
+  stable: bool = True,
 ):
   """
   High-level interface for adaptive binned top-k computation on TPU.
@@ -723,6 +752,7 @@ def _top_bounded_k(
       bins_topm_schedule=bins_topm_schedule,
       guarantee_convergence=guarantee_convergence,
       replace_val=replace_val,
+      stable=stable,
     ),
     in_specs=(
       pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0)),
@@ -767,6 +797,7 @@ def _top_bounded_k(
     "guarantee_convergence",
     "replace_val",
     "interpret",
+    "stable",
   ),
 )
 @functools.wraps(_top_bounded_k)
@@ -782,6 +813,7 @@ def top_bounded_k(
   guarantee_convergence: bool = False,
   replace_val: float | int | None = None,
   interpret: bool = False,
+  stable: bool = True,
 ):
   def _closed_topk(logits: jax.Array, k: jax.Array):
     return _top_bounded_k(
@@ -796,6 +828,7 @@ def _closed_topk(logits: jax.Array, k: jax.Array):
       guarantee_convergence=guarantee_convergence,
       replace_val=replace_val,
       interpret=interpret,
+      stable=stable,
     )
 
   @custom_partitioning

tallax/vllm/__init__.py

@@ -8,14 +8,12 @@
   - reference: Pure JAX reference implementation (no Pallas).
 """
 
-from tallax.vllm.sampling import topk_topp_and_sample
-from tallax.vllm.bounded_k import top_p_and_sample
-from tallax.vllm.arbitrary_k import topk_topp_mask_and_sample
-from tallax.vllm.reference import reference_topk_topp_mask_and_sample
+from tallax.vllm.sampling import bounded_topk_topp_and_sample
+from tallax.vllm.arbitrary_k import arbitrary_topk_topp_and_sample
+from tallax.vllm.reference import reference_topk_topp_and_sample
 
 __all__ = [
-  "topk_topp_and_sample",
-  "top_p_and_sample",
-  "topk_topp_mask_and_sample",
-  "reference_topk_topp_mask_and_sample",
+  "bounded_topk_topp_and_sample",
+  "arbitrary_topk_topp_and_sample",
+  "reference_topk_topp_and_sample",
 ]

tallax/vllm/arbitrary_k/__init__.py

@@ -6,6 +6,6 @@
 Suitable for large vocabularies where sorting would be prohibitively expensive.
 """
 
-from tallax.vllm.arbitrary_k.kernel import topk_topp_mask_and_sample
+from tallax.vllm.arbitrary_k.kernel import arbitrary_topk_topp_and_sample
 
-__all__ = ["topk_topp_mask_and_sample"]
+__all__ = ["arbitrary_topk_topp_and_sample"]