From 2a98bd8f641928588e4ad0cf201b537b8d8dc332 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Sun, 15 Feb 2026 16:03:47 +0000
Subject: [PATCH 01/33] Refactor debug output to use OrderedDict with
 standardized fields

Changes:
- Updated reference.py to return OrderedDict with exactly 5 fields:
  greedy_sampled, topk_logits_unsorted, topk_topp_unnorm_probs_i32_unsorted,
  random_unnorm_cdf_sampled, next_tokens
- Modified arbitrary_k/kernel.py to return same OrderedDict structure
  with full intermediate arrays when debug=True
- Added debug support to bounded_k/top_p_and_sample.py with OrderedDict
  containing k-slice versions of debug outputs
- Updated tests to use new debug field names and verify all intermediates

The bounded_k implementation returns debug values in the reduced k-slice
format (batch, k) rather than full vocabulary size, matching the input
dimensions.

https://claude.ai/code/session_01NJso6RoxkCzL1d3hSTnU5R
---
 tallax/vllm/arbitrary_k/kernel.py         | 104 +++++++++++-----------
 tallax/vllm/bounded_k/top_p_and_sample.py |  71 +++++++++++----
 tallax/vllm/reference.py                  |  16 ++--
 tests/vllm/arbitrary_k/kernel_test.py     |  23 +++--
 4 files changed, 131 insertions(+), 83 deletions(-)

diff --git a/tallax/vllm/arbitrary_k/kernel.py b/tallax/vllm/arbitrary_k/kernel.py
index 236c81f..b05b4ce 100644
--- a/tallax/vllm/arbitrary_k/kernel.py
+++ b/tallax/vllm/arbitrary_k/kernel.py
@@ -14,6 +14,7 @@
 for verifying that each intermediate matches the reference implementation.
 """
 
+from collections import OrderedDict
 import functools
 import jax
 import jax.numpy as jnp
@@ -63,7 +64,7 @@ def topk_topp_mask_and_sample_kernel(
   p_ref,
   temperature_ref,
   sampled_tokens_ref,
-  debug_results_ref=None,
+  debug_arrays_ref=None,
   *,
   stable: bool,
   replace_val: float,
@@ -78,7 +79,7 @@ def topk_topp_mask_and_sample_kernel(
     p_ref: Top-p values [block_token, 1]
     temperature_ref: Temperature values [block_token, 1]
     sampled_tokens_ref: Output sampled tokens [block_token, 1]
-    debug_results_ref: Optional dict of SMEM refs for intermediate checks
+    debug_arrays_ref: Optional dict of refs for full intermediate arrays
     stable: Whether to use stable masking
     replace_val: Replacement value for masked elements
     underlying_logits_dtype: Original dtype if logits were cast
@@ -93,7 +94,7 @@ def scoped_body(scoped_ref):
         p_ref,
         temperature_ref,
         sampled_tokens_ref,
-        debug_results_ref=debug_results_ref,
+        debug_arrays_ref=debug_arrays_ref,
         stable=stable,
         replace_val=replace_val,
         underlying_logits_dtype=logits_ref.dtype,
@@ -102,6 +103,7 @@ def scoped_body(scoped_ref):
     return pl.run_scoped(scoped_body, pltpu.VMEM(logits_ref.shape, jnp.float32))
 
   batch_size = logits_ref.shape[0]
+  vocab_size = logits_ref.shape[1]
   logits_max = map_reduce(logits_ref, reduce_fn="max")
   logits_max_lanes = jnp.broadcast_to(logits_max, (batch_size, NUM_LANES))
 
@@ -136,51 +138,35 @@ def scoped_body(scoped_ref):
   )
 
   # Stage 5: Sample
-  next_tokens = sample_probs(
-    unnorm_probs_i32, [ref[...] for ref in random_u128_in_u32s_ref]
+  random_u128_parts = [ref[...] for ref in random_u128_in_u32s_ref]
+  total_sum_u48 = U48.map_reduce_sum(unnorm_probs_i32, max_val=2**24 - 1)
+  sampled_u64_in_u32s = modulo_u128_u64(
+    random_u128_parts,
+    total_sum_u48.to_u64_in_u32s(),
+  )
+  random_unnorm_cdf_sampled = U48.from_u64_in_u32s(sampled_u64_in_u32s)
+  next_tokens = find_boundary_idx(
+    unnorm_probs_i32,
+    map_fn=lambda x: U48(x, max_val=2**24 - 1),
+    target=random_unnorm_cdf_sampled,
   )
 
   sampled_tokens_ref[...] = jnp.where(
     temperature < _SAMPLING_EPS, greedy_sampled, next_tokens
   )
 
-  # Debug: write intermediate match checks if debug refs are provided
-  if debug_results_ref is not None:
-    _write_debug_results(
-      debug_results_ref,
-      greedy_sampled=greedy_sampled,
-      topk_logits=topk_logits,
-      unnorm_probs_i32=unnorm_probs_i32,
-      next_tokens=next_tokens,
+  # Debug: write full intermediate arrays if debug refs are provided
+  if debug_arrays_ref is not None:
+    debug_arrays_ref["greedy_sampled"][...] = greedy_sampled
+    debug_arrays_ref["topk_logits_unsorted"][...] = topk_logits
+    debug_arrays_ref["topk_topp_unnorm_probs_i32_unsorted"][...] = unnorm_probs_i32
+    # Extract the low 32 bits of the sampled value for consistency with reference
+    random_unnorm_cdf_sampled_i64 = (
+      (random_unnorm_cdf_sampled.high.astype(jnp.int64) << 32) +
+      random_unnorm_cdf_sampled.low.astype(jnp.int64)
     )
-
-
-def _write_debug_results(
-  debug_refs,
-  *,
-  greedy_sampled,
-  topk_logits,
-  unnorm_probs_i32,
-  next_tokens,
-):
-  """Write debug intermediate values to SMEM refs.
-
-  Each ref is int32[1]: 1 if the intermediate matches the expected value,
-  0 if not. The expected values are written by the test harness.
-  """
-  if "greedy_sampled" in debug_refs:
-    debug_refs["greedy_sampled"][...] = greedy_sampled[:1, :1].astype(jnp.int32)
-  if "topk_logits_hash" in debug_refs:
-    # Store a simple checksum: sum of sign bits as a fingerprint
-    debug_refs["topk_logits_hash"][...] = (
-      (topk_logits != 0).astype(jnp.int32).sum(keepdims=True)[:1, :1]
-    )
-  if "topp_nonzero_count" in debug_refs:
-    debug_refs["topp_nonzero_count"][...] = (
-      (unnorm_probs_i32 != 0).astype(jnp.int32).sum(keepdims=True)[:1, :1]
-    )
-  if "next_tokens" in debug_refs:
-    debug_refs["next_tokens"][...] = next_tokens[:1, :1].astype(jnp.int32)
+    debug_arrays_ref["random_unnorm_cdf_sampled"][...] = random_unnorm_cdf_sampled_i64
+    debug_arrays_ref["next_tokens"][...] = next_tokens
 
 
 @functools.partial(
@@ -262,12 +248,6 @@ def topk_topp_mask_and_sample(
     sample_random_u128_in_u32s(rng_key, (padded_batch, 1))
   )
 
-  # Build debug output specs if requested
-  debug_shapes = {}
-  if debug:
-    for name in ["greedy_sampled", "topk_logits_hash", "topp_nonzero_count", "next_tokens"]:
-      debug_shapes[name] = jax.ShapeDtypeStruct((1, 1), jnp.int32)
-
   output_shape = jax.ShapeDtypeStruct((padded_batch, 1), jnp.int32)
 
   if not debug:
@@ -296,7 +276,14 @@ def topk_topp_mask_and_sample(
 
     return result[:batch_size, 0]
   else:
-    # Debug mode: also output debug intermediates via SMEM
+    # Debug mode: also output debug intermediates as full arrays
+    debug_shapes = OrderedDict([
+      ("greedy_sampled", jax.ShapeDtypeStruct((padded_batch, 1), jnp.int32)),
+      ("topk_logits_unsorted", jax.ShapeDtypeStruct((padded_batch, vocab_size), jnp.float32)),
+      ("topk_topp_unnorm_probs_i32_unsorted", jax.ShapeDtypeStruct((padded_batch, vocab_size), jnp.int32)),
+      ("random_unnorm_cdf_sampled", jax.ShapeDtypeStruct((padded_batch, 1), jnp.int64)),
+      ("next_tokens", jax.ShapeDtypeStruct((padded_batch, 1), jnp.int32)),
+    ])
     out_shapes = (output_shape, debug_shapes)
     result, debug_results = pl.pallas_call(
       functools.partial(
@@ -318,13 +305,24 @@ def topk_topp_mask_and_sample(
       ],
       out_specs=(
         pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
-        jax.tree.map(
-          lambda _: pl.BlockSpec(memory_space=pltpu.SMEM),
-          debug_shapes,
-        ),
+        OrderedDict([
+          ("greedy_sampled", pl.BlockSpec((block_token, 1), lambda i: (i, 0))),
+          ("topk_logits_unsorted", pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0))),
+          ("topk_topp_unnorm_probs_i32_unsorted", pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0))),
+          ("random_unnorm_cdf_sampled", pl.BlockSpec((block_token, 1), lambda i: (i, 0))),
+          ("next_tokens", pl.BlockSpec((block_token, 1), lambda i: (i, 0))),
+        ]),
       ),
       compiler_params=pltpu.CompilerParams(vmem_limit_bytes=int(0.9 * 2**27)),
       interpret=interpret,
     )(logits, random_u128_in_u32s, k, p, temperature)
 
-    return result[:batch_size, 0], debug_results
+    # Trim padding and reshape debug outputs to match reference format
+    debug_results_trimmed = OrderedDict([
+      ("greedy_sampled", debug_results["greedy_sampled"][:batch_size, 0]),
+      ("topk_logits_unsorted", debug_results["topk_logits_unsorted"][:batch_size, :]),
+      ("topk_topp_unnorm_probs_i32_unsorted", debug_results["topk_topp_unnorm_probs_i32_unsorted"][:batch_size, :]),
+      ("random_unnorm_cdf_sampled", debug_results["random_unnorm_cdf_sampled"][:batch_size, 0]),
+      ("next_tokens", debug_results["next_tokens"][:batch_size, 0]),
+    ])
+    return result[:batch_size, 0], debug_results_trimmed
diff --git a/tallax/vllm/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index e3abd89..30c41bc 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -10,6 +10,7 @@
 Only supports vocab_size (k) <= 128 due to i32 overflow constraints.
 """
 
+from collections import OrderedDict
 import functools
 import jax
 import jax.numpy as jnp
@@ -113,6 +114,7 @@ def top_p_and_sample_arrays(
   random_u128_in_u32s,
   top_p,
   temperature,
+  debug=False,
 ):
   """Fused top-p filtering + sampling on pre-sorted top-k logits.
 
@@ -122,42 +124,69 @@ def top_p_and_sample_arrays(
     random_u128_in_u32s: List of 4 u32 arrays for random sampling
     top_p: Top-p threshold values, shape (batch_size,)
     temperature: Temperature values, shape (batch_size,)
+    debug: If True, return (tokens, debug_results) with intermediate values
 
   Returns:
-    Sampled tokens of shape (batch_size,)
+    Sampled tokens of shape (batch_size,), or (tokens, debug_dict) if debug=True
   """
   topk_logits = topk_logits.astype(jnp.float32)
 
+  # Store original shape for debug
+  batch_size, k = topk_logits.shape
+
   # Shift to dim 0 for sublane-based reductions
-  topk_logits = topk_logits.T
-  topk_idx = topk_idx.T
-  random_u128_in_u32s = [x.T for x in random_u128_in_u32s]
-  shape = topk_logits.shape
+  topk_logits_transposed = topk_logits.T
+  topk_idx_transposed = topk_idx.T
+  random_u128_in_u32s_transposed = [x.T for x in random_u128_in_u32s]
+  shape = topk_logits_transposed.shape
 
-  topk_logits = topk_logits / temperature[None, :].astype(topk_logits.dtype)
+  # Greedy sample (before temperature scaling)
+  greedy_sampled = topk_idx[:, 0]
 
-  unnorm_probs_i32 = top_p_integer_mask(
-    topk_logits=topk_logits, p=top_p, axis=0
+  # Temperature scaling
+  topk_logits_scaled = topk_logits_transposed / temperature[None, :].astype(topk_logits_transposed.dtype)
+
+  # Top-p masking in i32 space
+  unnorm_probs_i32_sorted = top_p_integer_mask(
+    topk_logits=topk_logits_scaled, p=top_p, axis=0
   )
 
   # Re-sort back to original index order for bitwise-matching sampling
-  inverted_idxs, unnorm_probs_i32 = bitonic_topk_arrays(
-    [-topk_idx, unnorm_probs_i32], k=topk_idx.shape[0], axis=0, num_keys=1
+  inverted_idxs, unnorm_probs_i32_unsorted = bitonic_topk_arrays(
+    [-topk_idx_transposed, unnorm_probs_i32_sorted], k=topk_idx_transposed.shape[0], axis=0, num_keys=1
   )
   idxs = -inverted_idxs
+
+  # Sample from the unnormalized probabilities
   target_cumsum = modulo_u128_u64(
-    random_u128_in_u32s,
+    random_u128_in_u32s_transposed,
     [
       jnp.zeros((1, shape[1]), dtype=jnp.uint32),
-      unnorm_probs_i32.sum(0, keepdims=True).astype(jnp.uint32),
+      unnorm_probs_i32_unsorted.sum(0, keepdims=True).astype(jnp.uint32),
     ],
   )[1]
-  cumsums = cumsum_arrays(unnorm_probs_i32, axis=0)
+  cumsums = cumsum_arrays(unnorm_probs_i32_unsorted, axis=0)
   threshold_local_idx = sum((c < target_cumsum) for c in cumsums)
 
   next_tokens = (idxs * (jax.lax.broadcasted_iota(jnp.int32, idxs.shape, 0) == threshold_local_idx)).sum(0)
-  greedy_sampled = topk_idx[0, :]
-  return jnp.where(temperature < _SAMPLING_EPS, greedy_sampled, next_tokens)
+  result = jnp.where(temperature < _SAMPLING_EPS, greedy_sampled, next_tokens)
+
+  if not debug:
+    return result
+
+  # Build debug output matching the reference format but for the k-slice
+  # Convert sampled CDF value to int64 for consistency
+  random_unnorm_cdf_sampled = target_cumsum.astype(jnp.int64)
+
+  debug_results = OrderedDict([
+    ("greedy_sampled", greedy_sampled),
+    ("topk_logits_unsorted", topk_logits),  # Original sorted logits in batch-first format
+    ("topk_topp_unnorm_probs_i32_unsorted", unnorm_probs_i32_unsorted.T),  # Transpose back to [batch, k]
+    ("random_unnorm_cdf_sampled", random_unnorm_cdf_sampled),
+    ("next_tokens", next_tokens),
+  ])
+
+  return result, debug_results
 
 
 def top_p_and_sample_refs(
@@ -195,6 +224,7 @@ def _top_p_and_sample(
   replace_val: float,
   interpret: bool = False,
   dim0_offset: int = 0,
+  debug: bool = False,
 ) -> jax.Array:
   """Fused TPU kernel for sampling with top-p filtering and temperature scaling.
 
@@ -208,9 +238,10 @@ def _top_p_and_sample(
     replace_val: Value to replace filtered logits with
     interpret: If True, run in CPU interpret mode
     dim0_offset: Offset for dim0 (batch) axis, for sharding
+    debug: If True, return (tokens, debug_results) with intermediate values
 
   Returns:
-    Sampled tokens of shape (batch_size,)
+    Sampled tokens of shape (batch_size,), or (tokens, debug_dict) if debug=True
   """
   batch_size = topk_logits.shape[0]
 
@@ -226,6 +257,7 @@ def _top_p_and_sample(
     random_u128_in_u32s=random_u128_in_u32s,
     top_p=top_p,
     temperature=temperature,
+    debug=debug,
   )
 
 
@@ -235,6 +267,7 @@ def _top_p_and_sample(
     "vocab_size",
     "replace_val",
     "interpret",
+    "debug",
   ),
 )
 def top_p_and_sample(
@@ -247,6 +280,7 @@ def top_p_and_sample(
   vocab_size: int,
   replace_val: float,
   interpret: bool = False,
+  debug: bool = False,
 ) -> jax.Array:
   """Sharded wrapper for top-p sampling with custom partitioning.
 
@@ -261,9 +295,10 @@ def top_p_and_sample(
     vocab_size: Total vocabulary size.
     replace_val: Value to replace filtered logits with.
     interpret: If True, run in CPU interpret mode.
+    debug: If True, return (tokens, debug_results) with intermediate values.
 
   Returns:
-    Sampled tokens of shape (batch_size,).
+    Sampled tokens of shape (batch_size,), or (tokens, debug_dict) if debug=True.
   """
 
   @custom_partitioning
@@ -279,6 +314,7 @@ def sharded_top_p_and_sample(
       vocab_size=vocab_size,
       replace_val=replace_val,
       interpret=interpret,
+      debug=debug,
     )
 
   def infer_sharding_from_operands(mesh, arg_shapes, result_shape):
@@ -305,6 +341,7 @@ def shmap_fn(topk_logits, topk_idx, rng_key, top_p, temperature):
         replace_val=replace_val,
         interpret=interpret,
         dim0_offset=dim0_offset,
+        debug=debug,
       )
 
     return mesh, shmap_fn, out_shardings, arg_shardings
diff --git a/tallax/vllm/reference.py b/tallax/vllm/reference.py
index 4120266..154eac2 100644
--- a/tallax/vllm/reference.py
+++ b/tallax/vllm/reference.py
@@ -7,6 +7,7 @@
 should match (modulo the intermediate representation).
 """
 
+from collections import OrderedDict
 import numpy as np
 import jax
 import jax.numpy as jnp
@@ -132,12 +133,11 @@ def reference_topk_topp_mask_and_sample(
     if not debug:
       return result
 
-    debug_results = {
-      "greedy_sampled": greedy_sampled,
-      "topk_logits_unsorted": topk_logits_unsorted,
-      "topp_unnorm_probs_i32": unnorm_probs_i32,
-      "topp_nonzero_count": (unnorm_probs_i32 != 0).sum(axis=1),
-      "total_sum": total_sum,
-      "next_tokens": next_tokens,
-    }
+    debug_results = OrderedDict([
+      ("greedy_sampled", greedy_sampled),
+      ("topk_logits_unsorted", topk_logits_unsorted),
+      ("topk_topp_unnorm_probs_i32_unsorted", unnorm_probs_i32),
+      ("random_unnorm_cdf_sampled", sampled_total),
+      ("next_tokens", next_tokens),
+    ])
     return result, debug_results
diff --git a/tests/vllm/arbitrary_k/kernel_test.py b/tests/vllm/arbitrary_k/kernel_test.py
index f59f166..81eae59 100644
--- a/tests/vllm/arbitrary_k/kernel_test.py
+++ b/tests/vllm/arbitrary_k/kernel_test.py
@@ -43,18 +43,31 @@ def test_arbitrary_k_vs_reference(seed):
 
   # Greedy argmax matches
   np.testing.assert_array_equal(
-    kernel_debug["greedy_sampled"][0, 0],
+    kernel_debug["greedy_sampled"][0],
     ref_debug["greedy_sampled"][0],
   )
 
   # Next tokens match
   np.testing.assert_array_equal(
-    kernel_debug["next_tokens"][0, 0],
+    kernel_debug["next_tokens"][0],
     ref_debug["next_tokens"][0],
   )
 
-  # Top-p nonzero count matches
+  # Top-k logits match (unsorted order)
+  np.testing.assert_allclose(
+    kernel_debug["topk_logits_unsorted"][0],
+    ref_debug["topk_logits_unsorted"][0],
+    rtol=1e-5,
+  )
+
+  # Top-p unnormalized probs match (unsorted order)
+  np.testing.assert_array_equal(
+    kernel_debug["topk_topp_unnorm_probs_i32_unsorted"][0],
+    ref_debug["topk_topp_unnorm_probs_i32_unsorted"][0],
+  )
+
+  # Random sampled CDF value matches
   np.testing.assert_array_equal(
-    kernel_debug["topp_nonzero_count"][0, 0],
-    ref_debug["topp_nonzero_count"][0],
+    kernel_debug["random_unnorm_cdf_sampled"][0],
+    ref_debug["random_unnorm_cdf_sampled"][0],
   )

From ff744b2dc8b1a57c2658e13d96855ab04041aec5 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Sun, 15 Feb 2026 16:08:36 +0000
Subject: [PATCH 02/33] Return random_unnorm_cdf_sampled as tuple of (high_u32,
 low_u32)

Changed random_unnorm_cdf_sampled from single i64 value to tuple of
two u32 values (high, low) in all debug outputs:
- reference.py: Split i64 into high and low u32 components
- arbitrary_k/kernel.py: Store U48 high/low as separate arrays
- bounded_k/top_p_and_sample.py: Split target_cumsum into high/low
- Updated test to verify both high and low components match

https://claude.ai/code/session_01NJso6RoxkCzL1d3hSTnU5R
---
 tallax/vllm/arbitrary_k/kernel.py         | 20 +++++++++++---------
 tallax/vllm/bounded_k/top_p_and_sample.py |  8 +++++---
 tallax/vllm/reference.py                  |  6 +++++-
 tests/vllm/arbitrary_k/kernel_test.py     | 10 +++++++---
 4 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/tallax/vllm/arbitrary_k/kernel.py b/tallax/vllm/arbitrary_k/kernel.py
index b05b4ce..425f920 100644
--- a/tallax/vllm/arbitrary_k/kernel.py
+++ b/tallax/vllm/arbitrary_k/kernel.py
@@ -160,12 +160,9 @@ def scoped_body(scoped_ref):
     debug_arrays_ref["greedy_sampled"][...] = greedy_sampled
     debug_arrays_ref["topk_logits_unsorted"][...] = topk_logits
     debug_arrays_ref["topk_topp_unnorm_probs_i32_unsorted"][...] = unnorm_probs_i32
-    # Extract the low 32 bits of the sampled value for consistency with reference
-    random_unnorm_cdf_sampled_i64 = (
-      (random_unnorm_cdf_sampled.high.astype(jnp.int64) << 32) +
-      random_unnorm_cdf_sampled.low.astype(jnp.int64)
-    )
-    debug_arrays_ref["random_unnorm_cdf_sampled"][...] = random_unnorm_cdf_sampled_i64
+    # Store as tuple of (high_u32, low_u32)
+    debug_arrays_ref["random_unnorm_cdf_sampled_high"][...] = random_unnorm_cdf_sampled.high
+    debug_arrays_ref["random_unnorm_cdf_sampled_low"][...] = random_unnorm_cdf_sampled.low
     debug_arrays_ref["next_tokens"][...] = next_tokens
 
 
@@ -281,7 +278,8 @@ def topk_topp_mask_and_sample(
       ("greedy_sampled", jax.ShapeDtypeStruct((padded_batch, 1), jnp.int32)),
       ("topk_logits_unsorted", jax.ShapeDtypeStruct((padded_batch, vocab_size), jnp.float32)),
       ("topk_topp_unnorm_probs_i32_unsorted", jax.ShapeDtypeStruct((padded_batch, vocab_size), jnp.int32)),
-      ("random_unnorm_cdf_sampled", jax.ShapeDtypeStruct((padded_batch, 1), jnp.int64)),
+      ("random_unnorm_cdf_sampled_high", jax.ShapeDtypeStruct((padded_batch, 1), jnp.uint32)),
+      ("random_unnorm_cdf_sampled_low", jax.ShapeDtypeStruct((padded_batch, 1), jnp.uint32)),
       ("next_tokens", jax.ShapeDtypeStruct((padded_batch, 1), jnp.int32)),
     ])
     out_shapes = (output_shape, debug_shapes)
@@ -309,7 +307,8 @@ def topk_topp_mask_and_sample(
           ("greedy_sampled", pl.BlockSpec((block_token, 1), lambda i: (i, 0))),
           ("topk_logits_unsorted", pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0))),
           ("topk_topp_unnorm_probs_i32_unsorted", pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0))),
-          ("random_unnorm_cdf_sampled", pl.BlockSpec((block_token, 1), lambda i: (i, 0))),
+          ("random_unnorm_cdf_sampled_high", pl.BlockSpec((block_token, 1), lambda i: (i, 0))),
+          ("random_unnorm_cdf_sampled_low", pl.BlockSpec((block_token, 1), lambda i: (i, 0))),
           ("next_tokens", pl.BlockSpec((block_token, 1), lambda i: (i, 0))),
         ]),
       ),
@@ -322,7 +321,10 @@ def topk_topp_mask_and_sample(
       ("greedy_sampled", debug_results["greedy_sampled"][:batch_size, 0]),
       ("topk_logits_unsorted", debug_results["topk_logits_unsorted"][:batch_size, :]),
       ("topk_topp_unnorm_probs_i32_unsorted", debug_results["topk_topp_unnorm_probs_i32_unsorted"][:batch_size, :]),
-      ("random_unnorm_cdf_sampled", debug_results["random_unnorm_cdf_sampled"][:batch_size, 0]),
+      ("random_unnorm_cdf_sampled", (
+        debug_results["random_unnorm_cdf_sampled_high"][:batch_size, 0],
+        debug_results["random_unnorm_cdf_sampled_low"][:batch_size, 0],
+      )),
       ("next_tokens", debug_results["next_tokens"][:batch_size, 0]),
     ])
     return result[:batch_size, 0], debug_results_trimmed
diff --git a/tallax/vllm/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index 30c41bc..9e61089 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -175,14 +175,16 @@ def top_p_and_sample_arrays(
     return result
 
   # Build debug output matching the reference format but for the k-slice
-  # Convert sampled CDF value to int64 for consistency
-  random_unnorm_cdf_sampled = target_cumsum.astype(jnp.int64)
+  # Convert sampled CDF value to tuple of (high_u32, low_u32)
+  target_cumsum_i64 = target_cumsum.astype(jnp.int64)
+  random_unnorm_cdf_sampled_high = (target_cumsum_i64 >> 32).astype(jnp.uint32)
+  random_unnorm_cdf_sampled_low = (target_cumsum_i64 & 0xFFFFFFFF).astype(jnp.uint32)
 
   debug_results = OrderedDict([
     ("greedy_sampled", greedy_sampled),
     ("topk_logits_unsorted", topk_logits),  # Original sorted logits in batch-first format
     ("topk_topp_unnorm_probs_i32_unsorted", unnorm_probs_i32_unsorted.T),  # Transpose back to [batch, k]
-    ("random_unnorm_cdf_sampled", random_unnorm_cdf_sampled),
+    ("random_unnorm_cdf_sampled", (random_unnorm_cdf_sampled_high, random_unnorm_cdf_sampled_low)),
     ("next_tokens", next_tokens),
   ])
 
diff --git a/tallax/vllm/reference.py b/tallax/vllm/reference.py
index 154eac2..3331bb3 100644
--- a/tallax/vllm/reference.py
+++ b/tallax/vllm/reference.py
@@ -133,11 +133,15 @@ def reference_topk_topp_mask_and_sample(
     if not debug:
       return result
 
+    # Convert sampled_total (i64) to tuple of (high_u32, low_u32)
+    sampled_total_high = (sampled_total >> 32).astype(jnp.uint32)
+    sampled_total_low = (sampled_total & 0xFFFFFFFF).astype(jnp.uint32)
+
     debug_results = OrderedDict([
       ("greedy_sampled", greedy_sampled),
       ("topk_logits_unsorted", topk_logits_unsorted),
       ("topk_topp_unnorm_probs_i32_unsorted", unnorm_probs_i32),
-      ("random_unnorm_cdf_sampled", sampled_total),
+      ("random_unnorm_cdf_sampled", (sampled_total_high, sampled_total_low)),
       ("next_tokens", next_tokens),
     ])
     return result, debug_results
diff --git a/tests/vllm/arbitrary_k/kernel_test.py b/tests/vllm/arbitrary_k/kernel_test.py
index 81eae59..23a7540 100644
--- a/tests/vllm/arbitrary_k/kernel_test.py
+++ b/tests/vllm/arbitrary_k/kernel_test.py
@@ -66,8 +66,12 @@ def test_arbitrary_k_vs_reference(seed):
     ref_debug["topk_topp_unnorm_probs_i32_unsorted"][0],
   )
 
-  # Random sampled CDF value matches
+  # Random sampled CDF value matches (tuple of high and low u32)
   np.testing.assert_array_equal(
-    kernel_debug["random_unnorm_cdf_sampled"][0],
-    ref_debug["random_unnorm_cdf_sampled"][0],
+    kernel_debug["random_unnorm_cdf_sampled"][0][0],
+    ref_debug["random_unnorm_cdf_sampled"][0][0],
+  )
+  np.testing.assert_array_equal(
+    kernel_debug["random_unnorm_cdf_sampled"][1][0],
+    ref_debug["random_unnorm_cdf_sampled"][1][0],
   )

From 69ac95aaa4f2dd25ca502755f59b9831a19899ee Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Sun, 15 Feb 2026 16:14:26 +0000
Subject: [PATCH 03/33] Normalize function names and signatures across
 implementations
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Renamed functions:
- reference_topk_topp_mask_and_sample → reference_topk_topp_and_sample
- topk_topp_mask_and_sample → arbitrary_topk_topp_and_sample
- top_p_and_sample → bounded_topk_topp_and_sample

Changes:
- Added max_k parameter to bounded_topk_topp_and_sample
- Moved random_u128_in_u32s generation outside sharding rules in
  bounded implementation for better control over RNG state
- Updated all imports and exports in __init__.py files
- Updated test imports to use new function names
- Normalized sharding rules to pass random_u128_in_u32s explicitly

https://claude.ai/code/session_01NJso6RoxkCzL1d3hSTnU5R
---
 tallax/vllm/__init__.py                   | 12 ++++-----
 tallax/vllm/arbitrary_k/__init__.py       |  4 +--
 tallax/vllm/arbitrary_k/kernel.py         | 10 ++++----
 tallax/vllm/bounded_k/__init__.py         |  4 +--
 tallax/vllm/bounded_k/top_p_and_sample.py | 31 ++++++++++++-----------
 tallax/vllm/reference.py                  |  2 +-
 tests/vllm/arbitrary_k/kernel_test.py     |  8 +++---
 7 files changed, 36 insertions(+), 35 deletions(-)

diff --git a/tallax/vllm/__init__.py b/tallax/vllm/__init__.py
index 8f15360..9ce2d21 100644
--- a/tallax/vllm/__init__.py
+++ b/tallax/vllm/__init__.py
@@ -9,13 +9,13 @@
 """
 
 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.bounded_k 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",
 ]
diff --git a/tallax/vllm/arbitrary_k/__init__.py b/tallax/vllm/arbitrary_k/__init__.py
index bac7c21..4181f1b 100644
--- a/tallax/vllm/arbitrary_k/__init__.py
+++ b/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"]
diff --git a/tallax/vllm/arbitrary_k/kernel.py b/tallax/vllm/arbitrary_k/kernel.py
index 425f920..dcbdd02 100644
--- a/tallax/vllm/arbitrary_k/kernel.py
+++ b/tallax/vllm/arbitrary_k/kernel.py
@@ -57,7 +57,7 @@ def sample_probs(unnorm_probs_i32, random_u128_in_u32s, max_val=2**24 - 1):
   )
 
 
-def topk_topp_mask_and_sample_kernel(
+def arbitrary_topk_topp_and_sample_kernel(
   logits_ref,
   random_u128_in_u32s_ref,
   k_ref,
@@ -87,7 +87,7 @@ def topk_topp_mask_and_sample_kernel(
   if logits_ref.dtype != jnp.float32:
     def scoped_body(scoped_ref):
       scoped_ref[...] = logits_ref[...].astype(jnp.float32)
-      return topk_topp_mask_and_sample_kernel(
+      return arbitrary_topk_topp_and_sample_kernel(
         scoped_ref,
         random_u128_in_u32s_ref,
         k_ref,
@@ -176,7 +176,7 @@ def scoped_body(scoped_ref):
     "debug",
   ],
 )
-def topk_topp_mask_and_sample(
+def arbitrary_topk_topp_and_sample(
   logits: jax.Array,
   rng_key: jax.Array,
   k: jax.Array,
@@ -250,7 +250,7 @@ def topk_topp_mask_and_sample(
   if not debug:
     result = pl.pallas_call(
       functools.partial(
-        topk_topp_mask_and_sample_kernel,
+        arbitrary_topk_topp_and_sample_kernel,
         stable=stable,
         replace_val=replace_val,
       ),
@@ -285,7 +285,7 @@ def topk_topp_mask_and_sample(
     out_shapes = (output_shape, debug_shapes)
     result, debug_results = pl.pallas_call(
       functools.partial(
-        topk_topp_mask_and_sample_kernel,
+        arbitrary_topk_topp_and_sample_kernel,
         stable=stable,
         replace_val=replace_val,
       ),
diff --git a/tallax/vllm/bounded_k/__init__.py b/tallax/vllm/bounded_k/__init__.py
index f042b46..4032cfb 100644
--- a/tallax/vllm/bounded_k/__init__.py
+++ b/tallax/vllm/bounded_k/__init__.py
@@ -8,6 +8,6 @@
 enabled.
 """
 
-from tallax.vllm.bounded_k.top_p_and_sample import top_p_and_sample
+from tallax.vllm.bounded_k.top_p_and_sample import bounded_topk_topp_and_sample
 
-__all__ = ["top_p_and_sample"]
+__all__ = ["bounded_topk_topp_and_sample"]
diff --git a/tallax/vllm/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index 9e61089..329a7bd 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -218,7 +218,7 @@ def top_p_and_sample_refs(
 def _top_p_and_sample(
   topk_logits: jax.Array,
   topk_idx: jax.Array,
-  rng_key: jax.Array,
+  random_u128_in_u32s: list,
   top_p: jax.Array,
   temperature: jax.Array,
   *,
@@ -233,7 +233,7 @@ def _top_p_and_sample(
   Args:
     topk_logits: Sorted logits of shape (batch_size, k)
     topk_idx: Indices corresponding to sorted logits (batch_size, k)
-    rng_key: RNG key for sampling, shape (2,)
+    random_u128_in_u32s: List of 4 u32 arrays for random sampling
     top_p: Top-p threshold values
     temperature: Temperature values
     vocab_size: Vocabulary size for sampling
@@ -245,12 +245,6 @@ def _top_p_and_sample(
   Returns:
     Sampled tokens of shape (batch_size,), or (tokens, debug_dict) if debug=True
   """
-  batch_size = topk_logits.shape[0]
-
-  # Generate random u128 outside kernel
-  from tallax.vllm.utils.high_precision_uint import sample_random_u128_in_u32s
-  random_u128_in_u32s = list(sample_random_u128_in_u32s(rng_key, (batch_size, 1)))
-
   # Run the arrays-based implementation directly (no inner pallas_call needed
   # since the outer top_p_and_sample already wraps this)
   return top_p_and_sample_arrays(
@@ -267,12 +261,13 @@ def _top_p_and_sample(
   jit,
   static_argnames=(
     "vocab_size",
+    "max_k",
     "replace_val",
     "interpret",
     "debug",
   ),
 )
-def top_p_and_sample(
+def bounded_topk_topp_and_sample(
   topk_logits: jax.Array,
   topk_idx: jax.Array,
   rng_key: jax.Array,
@@ -280,6 +275,7 @@ def top_p_and_sample(
   temperature: jax.Array,
   *,
   vocab_size: int,
+  max_k: int,
   replace_val: float,
   interpret: bool = False,
   debug: bool = False,
@@ -295,6 +291,7 @@ def top_p_and_sample(
     top_p: Top-p threshold values.
     temperature: Temperature values.
     vocab_size: Total vocabulary size.
+    max_k: Maximum k value (bounded implementation supports k <= 128).
     replace_val: Value to replace filtered logits with.
     interpret: If True, run in CPU interpret mode.
     debug: If True, return (tokens, debug_results) with intermediate values.
@@ -302,15 +299,19 @@ def top_p_and_sample(
   Returns:
     Sampled tokens of shape (batch_size,), or (tokens, debug_dict) if debug=True.
   """
+  # Generate random u128 outside the sharded function
+  from tallax.vllm.utils.high_precision_uint import sample_random_u128_in_u32s
+  batch_size = topk_logits.shape[0]
+  random_u128_in_u32s = list(sample_random_u128_in_u32s(rng_key, (batch_size, 1)))
 
   @custom_partitioning
   def sharded_top_p_and_sample(
-    topk_logits, topk_idx, rng_key, top_p, temperature
+    topk_logits, topk_idx, random_u128_in_u32s, top_p, temperature
   ):
     return _top_p_and_sample(
       topk_logits,
       topk_idx,
-      rng_key,
+      random_u128_in_u32s,
       top_p,
       temperature,
       vocab_size=vocab_size,
@@ -329,14 +330,14 @@ def partition(mesh, arg_shapes, out_shapes):
     )
     batch_axis_name = arg_shardings[0].spec[0]
 
-    def shmap_fn(topk_logits, topk_idx, rng_key, top_p, temperature):
+    def shmap_fn(topk_logits, topk_idx, random_u128_in_u32s, top_p, temperature):
       dim0_offset = 0
       if batch_axis_name is not None:
         dim0_offset = jax.lax.axis_index(batch_axis_name) * topk_logits.shape[0]
       return _top_p_and_sample(
         topk_logits,
         topk_idx,
-        rng_key,
+        random_u128_in_u32s,
         top_p,
         temperature,
         vocab_size=vocab_size,
@@ -351,10 +352,10 @@ def shmap_fn(topk_logits, topk_idx, rng_key, top_p, temperature):
   sharded_top_p_and_sample.def_partition(
     infer_sharding_from_operands=infer_sharding_from_operands,
     partition=partition,
-    sharding_rule="b k, b k, r, b, b -> b",
+    sharding_rule="b k, b k, r r r r, b, b -> b",
     need_replication_factors=("k", "r"),
   )
 
   return sharded_top_p_and_sample(
-    topk_logits, topk_idx, rng_key, top_p, temperature
+    topk_logits, topk_idx, random_u128_in_u32s, top_p, temperature
   )
diff --git a/tallax/vllm/reference.py b/tallax/vllm/reference.py
index 3331bb3..aa4b9e7 100644
--- a/tallax/vllm/reference.py
+++ b/tallax/vllm/reference.py
@@ -44,7 +44,7 @@ def u128_modulo_u64_pure_callback(r_parts, m_parts):
   return (high.astype(jnp.int64) << 32) + low.astype(jnp.int64)
 
 
-def reference_topk_topp_mask_and_sample(
+def reference_topk_topp_and_sample(
   logits: jax.Array,
   rng_key: jax.Array,
   k: jax.Array,
diff --git a/tests/vllm/arbitrary_k/kernel_test.py b/tests/vllm/arbitrary_k/kernel_test.py
index 23a7540..7f38cd1 100644
--- a/tests/vllm/arbitrary_k/kernel_test.py
+++ b/tests/vllm/arbitrary_k/kernel_test.py
@@ -8,8 +8,8 @@
 import jax.numpy as jnp
 import numpy as np
 
-from tallax.vllm.reference import reference_topk_topp_mask_and_sample
-from tallax.vllm.arbitrary_k.kernel import topk_topp_mask_and_sample as arbitrary_k_sample
+from tallax.vllm.reference import reference_topk_topp_and_sample
+from tallax.vllm.arbitrary_k.kernel import arbitrary_topk_topp_and_sample
 from tallax.tax.utils import is_cpu_platform
 
 
@@ -31,10 +31,10 @@ def test_arbitrary_k_vs_reference(seed):
   """arbitrary_k kernel matches reference output and debug intermediates."""
   logits, k, p, temperature, rng_key = _make_inputs(seed, 4, 256)
 
-  ref_result, ref_debug = reference_topk_topp_mask_and_sample(
+  ref_result, ref_debug = reference_topk_topp_and_sample(
     logits, rng_key, k, p, temperature, debug=True
   )
-  kernel_result, kernel_debug = arbitrary_k_sample(
+  kernel_result, kernel_debug = arbitrary_topk_topp_and_sample(
     logits, rng_key, k, p, temperature, stable=True, debug=True
   )
 

From 5c8bd6de5b74607c81434393461906f4acdbdf2b Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 17:10:38 +0000
Subject: [PATCH 04/33] Fix U48 -> two u32s

---
 pyproject.toml                    | 18 --------
 tallax/vllm/arbitrary_k/kernel.py | 74 ++++++++++++++++++++++++-------
 2 files changed, 57 insertions(+), 35 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 8476466..37342ed 100644
--- a/pyproject.toml
+++ b/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",
diff --git a/tallax/vllm/arbitrary_k/kernel.py b/tallax/vllm/arbitrary_k/kernel.py
index 425f920..9023515 100644
--- a/tallax/vllm/arbitrary_k/kernel.py
+++ b/tallax/vllm/arbitrary_k/kernel.py
@@ -85,6 +85,7 @@ def topk_topp_mask_and_sample_kernel(
     underlying_logits_dtype: Original dtype if logits were cast
   """
   if logits_ref.dtype != jnp.float32:
+
     def scoped_body(scoped_ref):
       scoped_ref[...] = logits_ref[...].astype(jnp.float32)
       return topk_topp_mask_and_sample_kernel(
@@ -159,10 +160,16 @@ def scoped_body(scoped_ref):
   if debug_arrays_ref is not None:
     debug_arrays_ref["greedy_sampled"][...] = greedy_sampled
     debug_arrays_ref["topk_logits_unsorted"][...] = topk_logits
-    debug_arrays_ref["topk_topp_unnorm_probs_i32_unsorted"][...] = unnorm_probs_i32
+    debug_arrays_ref["topk_topp_unnorm_probs_i32_unsorted"][...] = (
+      unnorm_probs_i32
+    )
     # Store as tuple of (high_u32, low_u32)
-    debug_arrays_ref["random_unnorm_cdf_sampled_high"][...] = random_unnorm_cdf_sampled.high
-    debug_arrays_ref["random_unnorm_cdf_sampled_low"][...] = random_unnorm_cdf_sampled.low
+    for ref, val in zip(
+      debug_arrays_ref["random_unnorm_cdf_sampled"],
+      random_unnorm_cdf_sampled.to_u64_in_u32s(),
+      strict=True,
+    ):
+      ref[...] = val
     debug_arrays_ref["next_tokens"][...] = next_tokens
 
 
@@ -276,10 +283,22 @@ def topk_topp_mask_and_sample(
     # Debug mode: also output debug intermediates as full arrays
     debug_shapes = OrderedDict([
       ("greedy_sampled", jax.ShapeDtypeStruct((padded_batch, 1), jnp.int32)),
-      ("topk_logits_unsorted", jax.ShapeDtypeStruct((padded_batch, vocab_size), jnp.float32)),
-      ("topk_topp_unnorm_probs_i32_unsorted", jax.ShapeDtypeStruct((padded_batch, vocab_size), jnp.int32)),
-      ("random_unnorm_cdf_sampled_high", jax.ShapeDtypeStruct((padded_batch, 1), jnp.uint32)),
-      ("random_unnorm_cdf_sampled_low", jax.ShapeDtypeStruct((padded_batch, 1), jnp.uint32)),
+      (
+        "topk_logits_unsorted",
+        jax.ShapeDtypeStruct((padded_batch, vocab_size), jnp.float32),
+      ),
+      (
+        "topk_topp_unnorm_probs_i32_unsorted",
+        jax.ShapeDtypeStruct((padded_batch, vocab_size), jnp.int32),
+      ),
+      (
+        "random_unnorm_cdf_sampled_high",
+        jax.ShapeDtypeStruct((padded_batch, 1), jnp.uint32),
+      ),
+      (
+        "random_unnorm_cdf_sampled_low",
+        jax.ShapeDtypeStruct((padded_batch, 1), jnp.uint32),
+      ),
       ("next_tokens", jax.ShapeDtypeStruct((padded_batch, 1), jnp.int32)),
     ])
     out_shapes = (output_shape, debug_shapes)
@@ -305,10 +324,22 @@ def topk_topp_mask_and_sample(
         pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
         OrderedDict([
           ("greedy_sampled", pl.BlockSpec((block_token, 1), lambda i: (i, 0))),
-          ("topk_logits_unsorted", pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0))),
-          ("topk_topp_unnorm_probs_i32_unsorted", pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0))),
-          ("random_unnorm_cdf_sampled_high", pl.BlockSpec((block_token, 1), lambda i: (i, 0))),
-          ("random_unnorm_cdf_sampled_low", pl.BlockSpec((block_token, 1), lambda i: (i, 0))),
+          (
+            "topk_logits_unsorted",
+            pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0)),
+          ),
+          (
+            "topk_topp_unnorm_probs_i32_unsorted",
+            pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0)),
+          ),
+          (
+            "random_unnorm_cdf_sampled_high",
+            pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
+          ),
+          (
+            "random_unnorm_cdf_sampled_low",
+            pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
+          ),
           ("next_tokens", pl.BlockSpec((block_token, 1), lambda i: (i, 0))),
         ]),
       ),
@@ -319,12 +350,21 @@ def topk_topp_mask_and_sample(
     # Trim padding and reshape debug outputs to match reference format
     debug_results_trimmed = OrderedDict([
       ("greedy_sampled", debug_results["greedy_sampled"][:batch_size, 0]),
-      ("topk_logits_unsorted", debug_results["topk_logits_unsorted"][:batch_size, :]),
-      ("topk_topp_unnorm_probs_i32_unsorted", debug_results["topk_topp_unnorm_probs_i32_unsorted"][:batch_size, :]),
-      ("random_unnorm_cdf_sampled", (
-        debug_results["random_unnorm_cdf_sampled_high"][:batch_size, 0],
-        debug_results["random_unnorm_cdf_sampled_low"][:batch_size, 0],
-      )),
+      (
+        "topk_logits_unsorted",
+        debug_results["topk_logits_unsorted"][:batch_size, :],
+      ),
+      (
+        "topk_topp_unnorm_probs_i32_unsorted",
+        debug_results["topk_topp_unnorm_probs_i32_unsorted"][:batch_size, :],
+      ),
+      (
+        "random_unnorm_cdf_sampled",
+        (
+          debug_results["random_unnorm_cdf_sampled_high"][:batch_size, 0],
+          debug_results["random_unnorm_cdf_sampled_low"][:batch_size, 0],
+        ),
+      ),
       ("next_tokens", debug_results["next_tokens"][:batch_size, 0]),
     ])
     return result[:batch_size, 0], debug_results_trimmed

From d2c1dc6d3233d73e56ccf6438f5fdaa713597e9b Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 17:22:07 +0000
Subject: [PATCH 05/33] remove legacy OrderedDict and simplify debug path

---
 tallax/vllm/arbitrary_k/kernel.py         | 170 ++++++++--------------
 tallax/vllm/bounded_k/top_p_and_sample.py |  15 +-
 tallax/vllm/reference.py                  |  15 +-
 3 files changed, 76 insertions(+), 124 deletions(-)

diff --git a/tallax/vllm/arbitrary_k/kernel.py b/tallax/vllm/arbitrary_k/kernel.py
index 12395dd..47f2d14 100644
--- a/tallax/vllm/arbitrary_k/kernel.py
+++ b/tallax/vllm/arbitrary_k/kernel.py
@@ -14,7 +14,6 @@
 for verifying that each intermediate matches the reference implementation.
 """
 
-from collections import OrderedDict
 import functools
 import jax
 import jax.numpy as jnp
@@ -254,117 +253,72 @@ def arbitrary_topk_topp_and_sample(
 
   output_shape = jax.ShapeDtypeStruct((padded_batch, 1), jnp.int32)
 
-  if not debug:
-    result = pl.pallas_call(
-      functools.partial(
-        arbitrary_topk_topp_and_sample_kernel,
-        stable=stable,
-        replace_val=replace_val,
+  if debug:
+    debug_out_shapes = {
+      "greedy_sampled": jax.ShapeDtypeStruct((padded_batch, 1), jnp.int32),
+      "topk_logits_unsorted": jax.ShapeDtypeStruct(
+        (padded_batch, vocab_size), jnp.float32
       ),
-      out_shape=output_shape,
-      grid=(num_blocks,),
-      in_specs=[
-        pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0)),
-        jax.tree.map(
-          lambda x: pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
-          random_u128_in_u32s,
-        ),
-        pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
-        pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
-        pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
-      ],
-      out_specs=pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
-      compiler_params=pltpu.CompilerParams(vmem_limit_bytes=int(0.9 * 2**27)),
-      interpret=interpret,
-    )(logits, random_u128_in_u32s, k, p, temperature)
-
-    return result[:batch_size, 0]
-  else:
-    # Debug mode: also output debug intermediates as full arrays
-    debug_shapes = OrderedDict([
-      ("greedy_sampled", jax.ShapeDtypeStruct((padded_batch, 1), jnp.int32)),
-      (
-        "topk_logits_unsorted",
-        jax.ShapeDtypeStruct((padded_batch, vocab_size), jnp.float32),
-      ),
-      (
-        "topk_topp_unnorm_probs_i32_unsorted",
-        jax.ShapeDtypeStruct((padded_batch, vocab_size), jnp.int32),
+      "topk_topp_unnorm_probs_i32_unsorted": jax.ShapeDtypeStruct(
+        (padded_batch, vocab_size), jnp.int32
       ),
-      (
-        "random_unnorm_cdf_sampled_high",
+      "random_unnorm_cdf_sampled": (
         jax.ShapeDtypeStruct((padded_batch, 1), jnp.uint32),
+      )
+      * 2,
+      "next_tokens": jax.ShapeDtypeStruct((padded_batch, 1), jnp.int32),
+    }
+    debug_out_specs = {
+      "greedy_sampled": pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
+      "topk_logits_unsorted": pl.BlockSpec(
+        (block_token, vocab_size), lambda i: (i, 0)
       ),
-      (
-        "random_unnorm_cdf_sampled_low",
-        jax.ShapeDtypeStruct((padded_batch, 1), jnp.uint32),
+      "topk_topp_unnorm_probs_i32_unsorted": pl.BlockSpec(
+        (block_token, vocab_size), lambda i: (i, 0)
       ),
-      ("next_tokens", jax.ShapeDtypeStruct((padded_batch, 1), jnp.int32)),
-    ])
-    out_shapes = (output_shape, debug_shapes)
-    result, debug_results = pl.pallas_call(
-      functools.partial(
-        arbitrary_topk_topp_and_sample_kernel,
-        stable=stable,
-        replace_val=replace_val,
-      ),
-      out_shape=out_shapes,
-      grid=(num_blocks,),
-      in_specs=[
-        pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0)),
-        jax.tree.map(
-          lambda x: pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
-          random_u128_in_u32s,
-        ),
-        pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
-        pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
-        pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
-      ],
-      out_specs=(
+      "random_unnorm_cdf_sampled": (
         pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
-        OrderedDict([
-          ("greedy_sampled", pl.BlockSpec((block_token, 1), lambda i: (i, 0))),
-          (
-            "topk_logits_unsorted",
-            pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0)),
-          ),
-          (
-            "topk_topp_unnorm_probs_i32_unsorted",
-            pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0)),
-          ),
-          (
-            "random_unnorm_cdf_sampled_high",
-            pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
-          ),
-          (
-            "random_unnorm_cdf_sampled_low",
-            pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
-          ),
-          ("next_tokens", pl.BlockSpec((block_token, 1), lambda i: (i, 0))),
-        ]),
-      ),
-      compiler_params=pltpu.CompilerParams(vmem_limit_bytes=int(0.9 * 2**27)),
-      interpret=interpret,
-    )(logits, random_u128_in_u32s, k, p, temperature)
-
-    # Trim padding and reshape debug outputs to match reference format
-    debug_results_trimmed = OrderedDict([
-      ("greedy_sampled", debug_results["greedy_sampled"][:batch_size, 0]),
-      (
-        "topk_logits_unsorted",
-        debug_results["topk_logits_unsorted"][:batch_size, :],
-      ),
-      (
-        "topk_topp_unnorm_probs_i32_unsorted",
-        debug_results["topk_topp_unnorm_probs_i32_unsorted"][:batch_size, :],
-      ),
-      (
-        "random_unnorm_cdf_sampled",
-        (
-          debug_results["random_unnorm_cdf_sampled_high"][:batch_size, 0],
-          debug_results["random_unnorm_cdf_sampled_low"][:batch_size, 0],
-        ),
+      )
+      * 2,
+      "next_tokens": pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
+    }
+  else:
+    debug_out_shapes = None
+    debug_out_specs = None
+
+  results = pl.pallas_call(
+    functools.partial(
+      arbitrary_topk_topp_and_sample_kernel,
+      stable=stable,
+      replace_val=replace_val,
+    ),
+    out_shape=(output_shape, debug_out_shapes),
+    grid=(num_blocks,),
+    in_specs=[
+      pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0)),
+      jax.tree.map(
+        lambda x: pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
+        random_u128_in_u32s,
       ),
-      ("next_tokens", debug_results["next_tokens"][:batch_size, 0]),
-    ])
-    return result[:batch_size, 0], debug_results_trimmed
+      pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
+      pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
+      pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
+    ],
+    out_specs=(
+      pl.BlockSpec((block_token, 1), lambda i: (i, 0)),
+      debug_out_specs,
+    ),
+    compiler_params=pltpu.CompilerParams(vmem_limit_bytes=int(0.9 * 2**27)),
+    interpret=interpret,
+  )(logits, random_u128_in_u32s, k, p, temperature)
+
+  def trim(x):
+    if x.shape[1] == 1:
+      x = x.squeeze(1)
+    return x[:batch_size]
+
+  sampled_tokens, debug_aux = jax.tree.map(trim, results)
+
+  if debug:
+    return sampled_tokens, debug_aux
+  return sampled_tokens
diff --git a/tallax/vllm/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index 329a7bd..f91342e 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -10,7 +10,6 @@
 Only supports vocab_size (k) <= 128 due to i32 overflow constraints.
 """
 
-from collections import OrderedDict
 import functools
 import jax
 import jax.numpy as jnp
@@ -180,13 +179,13 @@ def top_p_and_sample_arrays(
   random_unnorm_cdf_sampled_high = (target_cumsum_i64 >> 32).astype(jnp.uint32)
   random_unnorm_cdf_sampled_low = (target_cumsum_i64 & 0xFFFFFFFF).astype(jnp.uint32)
 
-  debug_results = OrderedDict([
-    ("greedy_sampled", greedy_sampled),
-    ("topk_logits_unsorted", topk_logits),  # Original sorted logits in batch-first format
-    ("topk_topp_unnorm_probs_i32_unsorted", unnorm_probs_i32_unsorted.T),  # Transpose back to [batch, k]
-    ("random_unnorm_cdf_sampled", (random_unnorm_cdf_sampled_high, random_unnorm_cdf_sampled_low)),
-    ("next_tokens", next_tokens),
-  ])
+  debug_results = {
+    "greedy_sampled": greedy_sampled,
+    "topk_logits_unsorted": topk_logits,  # Original sorted logits in batch-first format
+    "topk_topp_unnorm_probs_i32_unsorted": unnorm_probs_i32_unsorted.T,  # Transpose back to [batch, k]
+    "random_unnorm_cdf_sampled": (random_unnorm_cdf_sampled_high, random_unnorm_cdf_sampled_low),
+    "next_tokens": next_tokens,
+  }
 
   return result, debug_results
 
diff --git a/tallax/vllm/reference.py b/tallax/vllm/reference.py
index aa4b9e7..c3e58b1 100644
--- a/tallax/vllm/reference.py
+++ b/tallax/vllm/reference.py
@@ -7,7 +7,6 @@
 should match (modulo the intermediate representation).
 """
 
-from collections import OrderedDict
 import numpy as np
 import jax
 import jax.numpy as jnp
@@ -137,11 +136,11 @@ def reference_topk_topp_and_sample(
     sampled_total_high = (sampled_total >> 32).astype(jnp.uint32)
     sampled_total_low = (sampled_total & 0xFFFFFFFF).astype(jnp.uint32)
 
-    debug_results = OrderedDict([
-      ("greedy_sampled", greedy_sampled),
-      ("topk_logits_unsorted", topk_logits_unsorted),
-      ("topk_topp_unnorm_probs_i32_unsorted", unnorm_probs_i32),
-      ("random_unnorm_cdf_sampled", (sampled_total_high, sampled_total_low)),
-      ("next_tokens", next_tokens),
-    ])
+    debug_results = {
+      "greedy_sampled": greedy_sampled,
+      "topk_logits_unsorted": topk_logits_unsorted,
+      "topk_topp_unnorm_probs_i32_unsorted": unnorm_probs_i32,
+      "random_unnorm_cdf_sampled": (sampled_total_high, sampled_total_low),
+      "next_tokens": next_tokens,
+    }
     return result, debug_results

From 1ba187da8e6956975d7e65e423ce862535ea4678 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 17:30:10 +0000
Subject: [PATCH 06/33] align reference debug values

---
 tallax/vllm/arbitrary_k/kernel.py | 3 ++-
 tallax/vllm/reference.py          | 6 +++++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/tallax/vllm/arbitrary_k/kernel.py b/tallax/vllm/arbitrary_k/kernel.py
index 47f2d14..cad30fe 100644
--- a/tallax/vllm/arbitrary_k/kernel.py
+++ b/tallax/vllm/arbitrary_k/kernel.py
@@ -127,6 +127,7 @@ def scoped_body(scoped_ref):
   )
 
   # Stage 3: Temperature scaling
+  topk_logits_pre_temperature = topk_logits
   topk_logits /= temperature
   logits_max /= temperature
 
@@ -158,7 +159,7 @@ def scoped_body(scoped_ref):
   # Debug: write full intermediate arrays if debug refs are provided
   if debug_arrays_ref is not None:
     debug_arrays_ref["greedy_sampled"][...] = greedy_sampled
-    debug_arrays_ref["topk_logits_unsorted"][...] = topk_logits
+    debug_arrays_ref["topk_logits_unsorted"][...] = topk_logits_pre_temperature
     debug_arrays_ref["topk_topp_unnorm_probs_i32_unsorted"][...] = (
       unnorm_probs_i32
     )
diff --git a/tallax/vllm/reference.py b/tallax/vllm/reference.py
index c3e58b1..466393c 100644
--- a/tallax/vllm/reference.py
+++ b/tallax/vllm/reference.py
@@ -39,6 +39,7 @@ def u128_modulo_u64_pure_callback(r_parts, m_parts):
     (jax.ShapeDtypeStruct(r_parts[0].shape, jnp.uint32),) * 2,
     *r_parts,
     *m_parts,
+    vectorized=True,
   )
   return (high.astype(jnp.int64) << 32) + low.astype(jnp.int64)
 
@@ -140,7 +141,10 @@ def reference_topk_topp_and_sample(
       "greedy_sampled": greedy_sampled,
       "topk_logits_unsorted": topk_logits_unsorted,
       "topk_topp_unnorm_probs_i32_unsorted": unnorm_probs_i32,
-      "random_unnorm_cdf_sampled": (sampled_total_high, sampled_total_low),
+      "random_unnorm_cdf_sampled": (
+        sampled_total_high.squeeze(1),
+        sampled_total_low.squeeze(1),
+      ),
       "next_tokens": next_tokens,
     }
     return result, debug_results

From 5c9cf6d2d1f0a9a4a666dbd9405d0c01e2c54121 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 17:37:05 +0000
Subject: [PATCH 07/33] fix vmap method

---
 tallax/vllm/reference.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tallax/vllm/reference.py b/tallax/vllm/reference.py
index 466393c..2b2909b 100644
--- a/tallax/vllm/reference.py
+++ b/tallax/vllm/reference.py
@@ -39,7 +39,7 @@ def u128_modulo_u64_pure_callback(r_parts, m_parts):
     (jax.ShapeDtypeStruct(r_parts[0].shape, jnp.uint32),) * 2,
     *r_parts,
     *m_parts,
-    vectorized=True,
+    vmap_method="legacy_vectorized",
   )
   return (high.astype(jnp.int64) << 32) + low.astype(jnp.int64)
 

From 18a35c2090b2e23dfc84a92208a6dcd0203ac144 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 19:40:04 +0000
Subject: [PATCH 08/33] WIP bounded top_p_and_sample pallas kernel fixes

---
 tallax/vllm/arbitrary_k/kernel.py         |  63 +++---
 tallax/vllm/bounded_k/top_p_and_sample.py | 235 +++++++++++++---------
 2 files changed, 176 insertions(+), 122 deletions(-)

diff --git a/tallax/vllm/arbitrary_k/kernel.py b/tallax/vllm/arbitrary_k/kernel.py
index cad30fe..5f434be 100644
--- a/tallax/vllm/arbitrary_k/kernel.py
+++ b/tallax/vllm/arbitrary_k/kernel.py
@@ -32,6 +32,32 @@
 _SAMPLING_EPS = 1e-5
 
 
+def canonicalize_args(batch_size, block_token, *args_with_pad):
+  """Canonicalize per-sample args and pad batch to block_token multiple.
+
+  Each arg is (value, pad_constant). 1D/scalar args are broadcast to batch_size
+  and reshaped to (batch_size, 1). All args are padded to the next multiple
+  of block_token.
+
+  Returns (padded_batch, num_blocks, *padded_args)
+  """
+  num_blocks = pl.cdiv(batch_size, block_token)
+  padded_batch = num_blocks * block_token
+  pad_size = padded_batch - batch_size
+
+  result = []
+  for val, pad_val in args_with_pad:
+    val = jnp.atleast_1d(val)
+    if val.ndim == 1:
+      if val.shape[0] == 1:
+        val = jnp.broadcast_to(val, (batch_size,))
+      val = jnp.reshape(val, (batch_size, 1))
+    if pad_size > 0:
+      val = jnp.pad(val, ((0, pad_size), (0, 0)), constant_values=pad_val)
+    result.append(val)
+  return (padded_batch, num_blocks, *result)
+
+
 def sample_probs(unnorm_probs_i32, random_u128_in_u32s, max_val=2**24 - 1):
   """Sample from unnormalized i32 probabilities using U48 arithmetic.
 
@@ -218,35 +244,14 @@ def arbitrary_topk_topp_and_sample(
   """
   batch_size, vocab_size = logits.shape
 
-  k = jnp.atleast_1d(k)
-  p = jnp.atleast_1d(p)
-  temperature = jnp.atleast_1d(temperature)
-
-  if k.shape[0] == 1:
-    k = jnp.broadcast_to(k, (batch_size,))
-  if p.shape[0] == 1:
-    p = jnp.broadcast_to(p, (batch_size,))
-  if temperature.shape[0] == 1:
-    temperature = jnp.broadcast_to(temperature, (batch_size,))
-
-  k = jnp.reshape(k, (batch_size, 1))
-  p = jnp.reshape(p, (batch_size, 1))
-  temperature = jnp.reshape(temperature, (batch_size, 1))
-
-  # Pad batch to multiple of block_token
-  num_blocks = pl.cdiv(batch_size, block_token)
-  padded_batch = num_blocks * block_token
-
-  if padded_batch != batch_size:
-    pad_size = padded_batch - batch_size
-    logits = jnp.pad(
-      logits, ((0, pad_size), (0, 0)), constant_values=replace_val
-    )
-    k = jnp.pad(k, ((0, pad_size), (0, 0)), constant_values=1)
-    p = jnp.pad(p, ((0, pad_size), (0, 0)), constant_values=1.0)
-    temperature = jnp.pad(
-      temperature, ((0, pad_size), (0, 0)), constant_values=1.0
-    )
+  padded_batch, num_blocks, logits, k, p, temperature = canonicalize_args(
+    batch_size,
+    block_token,
+    (logits, 0.0),
+    (k, 1),
+    (p, 1.0),
+    (temperature, 1.0),
+  )
 
   random_u128_in_u32s = tuple(
     sample_random_u128_in_u32s(rng_key, (padded_batch, 1))
diff --git a/tallax/vllm/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index f91342e..7fc8c44 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -20,6 +20,7 @@
 from jax.sharding import NamedSharding, PartitionSpec as P
 
 from tallax.vllm.utils.high_precision_uint import modulo_u128_u64
+from tallax.vllm.utils.high_precision_uint import sample_random_u128_in_u32s
 from tallax.tax.bitonic.topk import bitonic_topk_arrays
 from tallax.tax.cumsum import cumsum_arrays
 from tallax.tax.gather import take_along_axis_arrays
@@ -28,6 +29,16 @@
 _SAMPLING_EPS = 1e-5
 
 
+def canonicalize_args(batch_size, *args):
+  result = []
+  for val in args:
+    val = jnp.atleast_1d(val)
+    if val.shape[0] == 1:
+      val = jnp.broadcast_to(val, (batch_size,))
+    result.append(val)
+  return result
+
+
 def broadcast_to(x, shape):
   if x.shape[1] == shape[1] and shape[0] % NUM_LANES == 0 and x.shape[0] == 1:
     return pltpu.repeat(
@@ -140,10 +151,12 @@ def top_p_and_sample_arrays(
   shape = topk_logits_transposed.shape
 
   # Greedy sample (before temperature scaling)
-  greedy_sampled = topk_idx[:, 0]
+  greedy_sampled = topk_idx[:, :1]
 
   # Temperature scaling
-  topk_logits_scaled = topk_logits_transposed / temperature[None, :].astype(topk_logits_transposed.dtype)
+  topk_logits_scaled = topk_logits_transposed / temperature[None, :].astype(
+    topk_logits_transposed.dtype
+  )
 
   # Top-p masking in i32 space
   unnorm_probs_i32_sorted = top_p_integer_mask(
@@ -152,7 +165,10 @@ def top_p_and_sample_arrays(
 
   # Re-sort back to original index order for bitwise-matching sampling
   inverted_idxs, unnorm_probs_i32_unsorted = bitonic_topk_arrays(
-    [-topk_idx_transposed, unnorm_probs_i32_sorted], k=topk_idx_transposed.shape[0], axis=0, num_keys=1
+    [-topk_idx_transposed, unnorm_probs_i32_sorted],
+    k=topk_idx_transposed.shape[0],
+    axis=0,
+    num_keys=1,
   )
   idxs = -inverted_idxs
 
@@ -167,23 +183,28 @@ def top_p_and_sample_arrays(
   cumsums = cumsum_arrays(unnorm_probs_i32_unsorted, axis=0)
   threshold_local_idx = sum((c < target_cumsum) for c in cumsums)
 
-  next_tokens = (idxs * (jax.lax.broadcasted_iota(jnp.int32, idxs.shape, 0) == threshold_local_idx)).sum(0)
+  next_tokens = (
+    idxs
+    * (
+      jax.lax.broadcasted_iota(jnp.int32, idxs.shape, 0) == threshold_local_idx
+    )
+  ).sum(0)
   result = jnp.where(temperature < _SAMPLING_EPS, greedy_sampled, next_tokens)
 
   if not debug:
     return result
 
   # Build debug output matching the reference format but for the k-slice
-  # Convert sampled CDF value to tuple of (high_u32, low_u32)
-  target_cumsum_i64 = target_cumsum.astype(jnp.int64)
-  random_unnorm_cdf_sampled_high = (target_cumsum_i64 >> 32).astype(jnp.uint32)
-  random_unnorm_cdf_sampled_low = (target_cumsum_i64 & 0xFFFFFFFF).astype(jnp.uint32)
+  random_unnorm_cdf_sampled_low = target_cumsum.T
 
   debug_results = {
     "greedy_sampled": greedy_sampled,
-    "topk_logits_unsorted": topk_logits,  # Original sorted logits in batch-first format
+    "topk_logits_unsorted": topk_logits.T,  # Original sorted logits in batch-first format
     "topk_topp_unnorm_probs_i32_unsorted": unnorm_probs_i32_unsorted.T,  # Transpose back to [batch, k]
-    "random_unnorm_cdf_sampled": (random_unnorm_cdf_sampled_high, random_unnorm_cdf_sampled_low),
+    "random_unnorm_cdf_sampled": (
+      jnp.zeros_like(random_unnorm_cdf_sampled_low),
+      random_unnorm_cdf_sampled_low,
+    ),
     "next_tokens": next_tokens,
   }
 
@@ -193,67 +214,82 @@ def top_p_and_sample_arrays(
 def top_p_and_sample_refs(
   topk_logits_ref,
   topk_idx_ref,
-  rng_key_ref,
+  random_u128_in_u32s_refs,
   top_p_ref,
   temperature_ref,
-  dim0_offset_ref,
   sampled_tokens_ref,
-  *,
-  vocab_size: int,
-  replace_val: float,
+  debug_arrays_ref=None,
 ):
   """Pallas kernel body for top-p filtering + sampling."""
-  sampled_tokens_ref[...] = top_p_and_sample_arrays(
+  result = top_p_and_sample_arrays(
     topk_logits=topk_logits_ref[...],
     topk_idx=topk_idx_ref[...],
-    random_u128_in_u32s=[rng_key_ref[0, i:i+1] for i in range(4)] if False else
-      # Generate u128 from the rng key
-      list(jax.random.bits(rng_key_ref[...].reshape(-1), (4, 1, topk_logits_ref.shape[0]), jnp.uint32)),
+    random_u128_in_u32s=[ref[...] for ref in random_u128_in_u32s_refs],
     top_p=top_p_ref[...],
     temperature=temperature_ref[...],
+    debug=debug_arrays_ref is not None,
   )
+  if debug_arrays_ref is None:
+    sampled_tokens_ref[...] = result
+    return
+
+  sampled_tokens, debug_results = result
+  sampled_tokens_ref[...] = sampled_tokens
+  for key, val in debug_results.items():
+    if isinstance(val, tuple):
+      for ref, v in zip(debug_arrays_ref[key], val, strict=True):
+        ref[...] = v
+    else:
+      debug_arrays_ref[key][...] = val
 
 
 def _top_p_and_sample(
-  topk_logits: jax.Array,
-  topk_idx: jax.Array,
-  random_u128_in_u32s: list,
-  top_p: jax.Array,
-  temperature: jax.Array,
-  *,
-  vocab_size: int,
-  replace_val: float,
-  interpret: bool = False,
-  dim0_offset: int = 0,
-  debug: bool = False,
-) -> jax.Array:
-  """Fused TPU kernel for sampling with top-p filtering and temperature scaling.
+  topk_logits, topk_idx, random_u128_in_u32s, top_p, temperature, *, debug=False
+):
+  batch_size, k = topk_logits.shape
+  top_p, temperature = canonicalize_args(batch_size, top_p, temperature)
+
+  output_shape = jax.ShapeDtypeStruct((batch_size,), jnp.int32)
+
+  if debug:
+    debug_out_shapes = {
+      "greedy_sampled": jax.ShapeDtypeStruct((batch_size, 1), jnp.int32),
+      "topk_logits_unsorted": jax.ShapeDtypeStruct(
+        (batch_size, k), jnp.float32
+      ),
+      "topk_topp_unnorm_probs_i32_unsorted": jax.ShapeDtypeStruct(
+        (batch_size, k), jnp.int32
+      ),
+      "random_unnorm_cdf_sampled": (
+        jax.ShapeDtypeStruct((batch_size, 1), jnp.uint32),
+      )
+      * 2,
+      "next_tokens": jax.ShapeDtypeStruct((batch_size, 1), jnp.int32),
+    }
+  else:
+    debug_out_shapes = None
+
+  results = pl.pallas_call(
+    top_p_and_sample_refs,
+    out_shape=(output_shape, debug_out_shapes),
+  )(
+    topk_logits,
+    topk_idx,
+    random_u128_in_u32s,
+    top_p,
+    temperature,
+  )
 
-  Args:
-    topk_logits: Sorted logits of shape (batch_size, k)
-    topk_idx: Indices corresponding to sorted logits (batch_size, k)
-    random_u128_in_u32s: List of 4 u32 arrays for random sampling
-    top_p: Top-p threshold values
-    temperature: Temperature values
-    vocab_size: Vocabulary size for sampling
-    replace_val: Value to replace filtered logits with
-    interpret: If True, run in CPU interpret mode
-    dim0_offset: Offset for dim0 (batch) axis, for sharding
-    debug: If True, return (tokens, debug_results) with intermediate values
+  def trim(x):
+    if x.shape[1] == 1:
+      x = x.squeeze(1)
+    return x[:batch_size]
 
-  Returns:
-    Sampled tokens of shape (batch_size,), or (tokens, debug_dict) if debug=True
-  """
-  # Run the arrays-based implementation directly (no inner pallas_call needed
-  # since the outer top_p_and_sample already wraps this)
-  return top_p_and_sample_arrays(
-    topk_logits=topk_logits,
-    topk_idx=topk_idx,
-    random_u128_in_u32s=random_u128_in_u32s,
-    top_p=top_p,
-    temperature=temperature,
-    debug=debug,
-  )
+  sampled_tokens, debug_aux = jax.tree.map(trim, results)
+
+  if debug:
+    return sampled_tokens, debug_aux
+  return sampled_tokens
 
 
 @functools.partial(
@@ -266,17 +302,13 @@ def _top_p_and_sample(
     "debug",
   ),
 )
-def bounded_topk_topp_and_sample(
+def topp_and_sample(
   topk_logits: jax.Array,
   topk_idx: jax.Array,
   rng_key: jax.Array,
   top_p: jax.Array,
   temperature: jax.Array,
   *,
-  vocab_size: int,
-  max_k: int,
-  replace_val: float,
-  interpret: bool = False,
   debug: bool = False,
 ) -> jax.Array:
   """Sharded wrapper for top-p sampling with custom partitioning.
@@ -299,62 +331,79 @@ def bounded_topk_topp_and_sample(
     Sampled tokens of shape (batch_size,), or (tokens, debug_dict) if debug=True.
   """
   # Generate random u128 outside the sharded function
-  from tallax.vllm.utils.high_precision_uint import sample_random_u128_in_u32s
   batch_size = topk_logits.shape[0]
-  random_u128_in_u32s = list(sample_random_u128_in_u32s(rng_key, (batch_size, 1)))
+  random_u128_in_u32s = list(
+    sample_random_u128_in_u32s(rng_key, (batch_size, 1))
+  )
+
+  out_shapes = jax.eval_shape(
+    _top_p_and_sample,
+    topk_logits,
+    topk_idx,
+    random_u128_in_u32s,
+    top_p,
+    temperature,
+    debug=debug,
+  )
 
   @custom_partitioning
   def sharded_top_p_and_sample(
     topk_logits, topk_idx, random_u128_in_u32s, top_p, temperature
   ):
-    return _top_p_and_sample(
-      topk_logits,
-      topk_idx,
-      random_u128_in_u32s,
-      top_p,
-      temperature,
-      vocab_size=vocab_size,
-      replace_val=replace_val,
-      interpret=interpret,
-      debug=debug,
+    return jax.tree.leaves(
+      _top_p_and_sample(
+        topk_logits,
+        topk_idx,
+        random_u128_in_u32s,
+        top_p,
+        temperature,
+        debug=debug,
+      )
     )
 
   def infer_sharding_from_operands(mesh, arg_shapes, result_shape):
     batch_spec = arg_shapes[0].sharding.spec[0]
-    return NamedSharding(mesh, P(batch_spec))
+    return [
+      NamedSharding(mesh, P(batch_spec))
+      for _ in range(len(jax.tree.leaves(out_shapes)))
+    ]
 
   def partition(mesh, arg_shapes, out_shapes):
     arg_shardings, out_shardings = jax.tree.map(
       lambda s: s.sharding, (arg_shapes, out_shapes)
     )
-    batch_axis_name = arg_shardings[0].spec[0]
 
-    def shmap_fn(topk_logits, topk_idx, random_u128_in_u32s, top_p, temperature):
-      dim0_offset = 0
-      if batch_axis_name is not None:
-        dim0_offset = jax.lax.axis_index(batch_axis_name) * topk_logits.shape[0]
-      return _top_p_and_sample(
-        topk_logits,
-        topk_idx,
-        random_u128_in_u32s,
-        top_p,
-        temperature,
-        vocab_size=vocab_size,
-        replace_val=replace_val,
-        interpret=interpret,
-        dim0_offset=dim0_offset,
-        debug=debug,
+    def shmap_fn(
+      topk_logits, topk_idx, random_u128_in_u32s, top_p, temperature
+    ):
+      return jax.tree.leaves(
+        _top_p_and_sample(
+          topk_logits,
+          topk_idx,
+          random_u128_in_u32s,
+          top_p,
+          temperature,
+          debug=debug,
+        )
       )
 
     return mesh, shmap_fn, out_shardings, arg_shardings
 
+  sharding_rule = "b k, b k, r r r r, b, b -> b"
+  if debug:
+    sharding_rule += ", b" * (len(jax.tree.leaves(out_shapes)) - 1)
   sharded_top_p_and_sample.def_partition(
     infer_sharding_from_operands=infer_sharding_from_operands,
     partition=partition,
-    sharding_rule="b k, b k, r r r r, b, b -> b",
+    sharding_rule=sharding_rule,
     need_replication_factors=("k", "r"),
   )
-
-  return sharded_top_p_and_sample(
+  flat_outs = sharded_top_p_and_sample(
     topk_logits, topk_idx, random_u128_in_u32s, top_p, temperature
   )
+  sampled_tokens, debug_aux = jax.tree.unflatten(
+    jax.tree.structure(out_shapes), flat_outs
+  )
+  if debug:
+    return sampled_tokens, debug_aux
+  return sampled_tokens

From 57587ff8c72a11db1bf1d7e9bebea6fbaede9348 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 19:45:28 +0000
Subject: [PATCH 09/33] fix static argnames

---
 tallax/vllm/bounded_k/top_p_and_sample.py | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/tallax/vllm/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index 7fc8c44..566d16e 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -294,13 +294,7 @@ def trim(x):
 
 @functools.partial(
   jit,
-  static_argnames=(
-    "vocab_size",
-    "max_k",
-    "replace_val",
-    "interpret",
-    "debug",
-  ),
+  static_argnames=("debug",),
 )
 def topp_and_sample(
   topk_logits: jax.Array,

From baabf0702dd544f073f691cd69c280a478336beb Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 19:46:23 +0000
Subject: [PATCH 10/33] remove legacy import

---
 tallax/vllm/bounded_k/__init__.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tallax/vllm/bounded_k/__init__.py b/tallax/vllm/bounded_k/__init__.py
index 4032cfb..0e07401 100644
--- a/tallax/vllm/bounded_k/__init__.py
+++ b/tallax/vllm/bounded_k/__init__.py
@@ -8,6 +8,6 @@
 enabled.
 """
 
-from tallax.vllm.bounded_k.top_p_and_sample import bounded_topk_topp_and_sample
+# from tallax.vllm.bounded_k.top_p_and_sample import bounded_topk_topp_and_sample
 
-__all__ = ["bounded_topk_topp_and_sample"]
+# __all__ = ["bounded_topk_topp_and_sample"]

From cf15b5e529eab96b440722bdfd6f6ae9d0509b8f Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 19:47:17 +0000
Subject: [PATCH 11/33] more legacy imports removed

---
 tallax/vllm/__init__.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tallax/vllm/__init__.py b/tallax/vllm/__init__.py
index 9ce2d21..27aae80 100644
--- a/tallax/vllm/__init__.py
+++ b/tallax/vllm/__init__.py
@@ -9,13 +9,14 @@
 """
 
 from tallax.vllm.sampling import topk_topp_and_sample
-from tallax.vllm.bounded_k import bounded_topk_topp_and_sample
+
+# from tallax.vllm.bounded_k 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",
-  "bounded_topk_topp_and_sample",
+  # "bounded_topk_topp_and_sample",
   "arbitrary_topk_topp_and_sample",
   "reference_topk_topp_and_sample",
 ]

From e3f735a9bcf049d425ce91fc067225a61f8bc664 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 21:08:20 +0000
Subject: [PATCH 12/33] Add pad value for find_boundary_idx to deal with pad
 issues

---
 tallax/vllm/arbitrary_k/kernel.py    |  1 +
 tallax/vllm/arbitrary_k/topk_mask.py | 12 +++++++++---
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/tallax/vllm/arbitrary_k/kernel.py b/tallax/vllm/arbitrary_k/kernel.py
index 5f434be..58703e0 100644
--- a/tallax/vllm/arbitrary_k/kernel.py
+++ b/tallax/vllm/arbitrary_k/kernel.py
@@ -176,6 +176,7 @@ def scoped_body(scoped_ref):
     unnorm_probs_i32,
     map_fn=lambda x: U48(x, max_val=2**24 - 1),
     target=random_unnorm_cdf_sampled,
+    pad_val=0,
   )
 
   sampled_tokens_ref[...] = jnp.where(
diff --git a/tallax/vllm/arbitrary_k/topk_mask.py b/tallax/vllm/arbitrary_k/topk_mask.py
index f7c5747..2b0bac0 100644
--- a/tallax/vllm/arbitrary_k/topk_mask.py
+++ b/tallax/vllm/arbitrary_k/topk_mask.py
@@ -26,6 +26,7 @@ def _find_boundary_chunk(
   map_fn,
   target,
   chunk_size: int,
+  pad_val,
   active_chunk: jax.Array | None = None,
   ref_offset: jax.Array | int = 0,
 ):
@@ -40,6 +41,7 @@ def _find_boundary_chunk(
     map_fn: Unary function mapping chunks to binary counters
     target: Target count (shape [batch, 1])
     chunk_size: Size of each chunk
+    pad_val: Fill value for out-of-bounds positions in the boundary slice
     active_chunk: Optional subset of ref to search in
     ref_offset: Offset into ref for indexing
 
@@ -85,12 +87,12 @@ def _find_boundary_chunk(
     boundary_slice = jnp.where(
       (ref_offset[:, :1] + iota1) < ref.shape[1],
       boundary_slice,
-      get_dtype_info(boundary_slice).min,
+      pad_val,
     )
   return ref_offset, boundary_slice, target
 
 
-def find_boundary_idx(ref_or_arr, map_fn, target):
+def find_boundary_idx(ref_or_arr, map_fn, target, pad_val):
   """Find the lowest idx where map_fn(ref[...]).cumsum(1) >= target.
 
   Uses a two-level chunk search: first finds the right chunk of chunks,
@@ -100,6 +102,7 @@ def find_boundary_idx(ref_or_arr, map_fn, target):
     ref_or_arr: Pallas ref or array of shape [batch, vocab_size]
     map_fn: Maps chunks to binary counts
     target: Target cumulative sum value
+    pad_val: Fill value for out-of-bounds positions in boundary slices
 
   Returns:
     Index array of shape [batch, NUM_LANES]
@@ -109,7 +112,7 @@ def find_boundary_idx(ref_or_arr, map_fn, target):
 
     def scoped_body(scoped_ref):
       scoped_ref[...] = arr
-      return find_boundary_idx(scoped_ref, map_fn, target)
+      return find_boundary_idx(scoped_ref, map_fn, target, pad_val)
 
     return pl.run_scoped(scoped_body, pltpu.VMEM(arr.shape, arr.dtype))
   ref = ref_or_arr
@@ -120,12 +123,14 @@ def scoped_body(scoped_ref):
     map_fn=map_fn,
     target=target,
     chunk_size=int(math.sqrt(ref.shape[1] // NUM_LANES)) * NUM_LANES,
+    pad_val=pad_val,
   )
   ref_offset, boundary_slice, target = _find_boundary_chunk(
     ref,
     map_fn=map_fn,
     target=target,
     chunk_size=NUM_LANES,
+    pad_val=pad_val,
     ref_offset=ref_offset,
     active_chunk=boundary_slice,
   )
@@ -200,6 +205,7 @@ def topk_mask(
         lambda chunk: (chunk > threshold).astype(jnp.int32),
         reduce_fn="sum",
       ),
+      pad_val=get_dtype_info(logits_ref).min,
     )
     threshold = pltpu.repeat(
       threshold,

From ed3cd0521a219428ee31e2b3c456b1572680cb93 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 21:10:35 +0000
Subject: [PATCH 13/33] fix find_boundary_idx

---
 tallax/vllm/arbitrary_k/kernel.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tallax/vllm/arbitrary_k/kernel.py b/tallax/vllm/arbitrary_k/kernel.py
index 58703e0..cb69cf3 100644
--- a/tallax/vllm/arbitrary_k/kernel.py
+++ b/tallax/vllm/arbitrary_k/kernel.py
@@ -27,7 +27,7 @@
 )
 from tallax.vllm.arbitrary_k.topp_mask import topp_mask
 from tallax.vllm.arbitrary_k.topk_mask import topk_mask, find_boundary_idx
-from tallax.tax.utils import NUM_LANES, map_reduce
+from tallax.tax.utils import NUM_LANES, map_reduce, get_dtype_info
 
 _SAMPLING_EPS = 1e-5
 
@@ -79,6 +79,7 @@ def sample_probs(unnorm_probs_i32, random_u128_in_u32s, max_val=2**24 - 1):
     unnorm_probs_i32,
     map_fn=lambda x: U48(x, max_val=max_val),
     target=target_u48,
+    pad_val=0,
   )
 
 
@@ -140,6 +141,7 @@ def scoped_body(scoped_ref):
       chunk == pltpu.repeat(logits_max_lanes, chunk.shape[1] // NUM_LANES, 1)
     ).astype(jnp.int32),
     target=jnp.broadcast_to(jnp.float32(1), logits_max_lanes.shape),
+    pad_val=get_dtype_info(logits_ref).min,
   )[:, :1]
 
   # Stage 2: Top-k masking

From a498538b1f703601d10ad091bf2f5f1dba56eeef Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 21:17:22 +0000
Subject: [PATCH 14/33] use sample_probs fn

---
 tallax/vllm/arbitrary_k/kernel.py | 25 ++++++++-----------------
 1 file changed, 8 insertions(+), 17 deletions(-)

diff --git a/tallax/vllm/arbitrary_k/kernel.py b/tallax/vllm/arbitrary_k/kernel.py
index cb69cf3..b1ecf45 100644
--- a/tallax/vllm/arbitrary_k/kernel.py
+++ b/tallax/vllm/arbitrary_k/kernel.py
@@ -74,18 +74,18 @@ def sample_probs(unnorm_probs_i32, random_u128_in_u32s, max_val=2**24 - 1):
     random_u128_in_u32s,
     total_sum_u48.to_u64_in_u32s(),
   )
-  target_u48 = U48.from_u64_in_u32s(sampled_u64_in_u32s)
+  random_unnorm_cdf_sampled = U48.from_u64_in_u32s(sampled_u64_in_u32s)
   return find_boundary_idx(
     unnorm_probs_i32,
     map_fn=lambda x: U48(x, max_val=max_val),
-    target=target_u48,
+    target=random_unnorm_cdf_sampled,
     pad_val=0,
-  )
+  ), random_unnorm_cdf_sampled
 
 
 def arbitrary_topk_topp_and_sample_kernel(
   logits_ref,
-  random_u128_in_u32s_ref,
+  random_u128_in_u32s_refs,
   k_ref,
   p_ref,
   temperature_ref,
@@ -116,7 +116,7 @@ def scoped_body(scoped_ref):
       scoped_ref[...] = logits_ref[...].astype(jnp.float32)
       return arbitrary_topk_topp_and_sample_kernel(
         scoped_ref,
-        random_u128_in_u32s_ref,
+        random_u128_in_u32s_refs,
         k_ref,
         p_ref,
         temperature_ref,
@@ -167,20 +167,11 @@ def scoped_body(scoped_ref):
   )
 
   # Stage 5: Sample
-  random_u128_parts = [ref[...] for ref in random_u128_in_u32s_ref]
-  total_sum_u48 = U48.map_reduce_sum(unnorm_probs_i32, max_val=2**24 - 1)
-  sampled_u64_in_u32s = modulo_u128_u64(
-    random_u128_parts,
-    total_sum_u48.to_u64_in_u32s(),
-  )
-  random_unnorm_cdf_sampled = U48.from_u64_in_u32s(sampled_u64_in_u32s)
-  next_tokens = find_boundary_idx(
+  next_tokens, random_unnorm_cdf_sampled = sample_probs(
     unnorm_probs_i32,
-    map_fn=lambda x: U48(x, max_val=2**24 - 1),
-    target=random_unnorm_cdf_sampled,
-    pad_val=0,
+    [ref[...] for ref in random_u128_in_u32s_refs],
+    max_val=2**24 - 1,
   )
-
   sampled_tokens_ref[...] = jnp.where(
     temperature < _SAMPLING_EPS, greedy_sampled, next_tokens
   )

From 11f7c16079591fecbbd4c16f50f7ad0970aef5e8 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 22:05:26 +0000
Subject: [PATCH 15/33] fix imports

---
 tallax/vllm/arbitrary_k/kernel.py         |  1 -
 tallax/vllm/arbitrary_k/topk_mask.py      | 50 -----------------------
 tallax/vllm/bounded_k/__init__.py         |  4 +-
 tallax/vllm/bounded_k/top_p_and_sample.py | 14 ++-----
 tallax/vllm/sampling.py                   | 11 ++---
 5 files changed, 9 insertions(+), 71 deletions(-)

diff --git a/tallax/vllm/arbitrary_k/kernel.py b/tallax/vllm/arbitrary_k/kernel.py
index b1ecf45..eac7685 100644
--- a/tallax/vllm/arbitrary_k/kernel.py
+++ b/tallax/vllm/arbitrary_k/kernel.py
@@ -130,7 +130,6 @@ def scoped_body(scoped_ref):
     return pl.run_scoped(scoped_body, pltpu.VMEM(logits_ref.shape, jnp.float32))
 
   batch_size = logits_ref.shape[0]
-  vocab_size = logits_ref.shape[1]
   logits_max = map_reduce(logits_ref, reduce_fn="max")
   logits_max_lanes = jnp.broadcast_to(logits_max, (batch_size, NUM_LANES))
 
diff --git a/tallax/vllm/arbitrary_k/topk_mask.py b/tallax/vllm/arbitrary_k/topk_mask.py
index 2b0bac0..743f9d6 100644
--- a/tallax/vllm/arbitrary_k/topk_mask.py
+++ b/tallax/vllm/arbitrary_k/topk_mask.py
@@ -222,53 +222,3 @@ def topk_mask(
       )
     )
   return jnp.where(mask, logits, replace_val).astype(logits_ref.dtype)
-
-
-def topk_mask_pallas_kernel(
-  logits_ref,
-  k_ref,
-  output_ref,
-  *,
-  replace_val: float,
-  stable: bool,
-):
-  output_ref[...] = topk_mask(
-    logits_ref, k_ref, replace_val=replace_val, stable=stable
-  )
-
-
-@functools.partial(
-  jax.jit, static_argnames=["replace_val", "stable", "interpret"]
-)
-def topk_mask_pallas(
-  x: jax.Array,
-  k: int,
-  replace_val: float = -1e12,
-  stable: bool = True,
-  interpret: bool = False,
-) -> jax.Array:
-  """Pallas-based topk mask with parallel chunk-based reduction.
-
-  Args:
-    x: Input array of shape [batch, vocab_size]
-    k: Number of top elements
-    replace_val: Value for masked elements
-    stable: Whether to use stable masking
-    interpret: Whether to use interpret mode
-
-  Returns:
-    Masked array
-  """
-  batch_size, _vocab_size = x.shape
-  k = jnp.broadcast_to(k, (batch_size, 1))
-  output_shape = jax.ShapeDtypeStruct(x.shape, x.dtype)
-  return pl.pallas_call(
-    functools.partial(
-      topk_mask_pallas_kernel,
-      replace_val=replace_val,
-      stable=stable,
-    ),
-    compiler_params=pltpu.CompilerParams(vmem_limit_bytes=int(0.9 * 2**27)),
-    out_shape=output_shape,
-    interpret=interpret,
-  )(x, k)
diff --git a/tallax/vllm/bounded_k/__init__.py b/tallax/vllm/bounded_k/__init__.py
index 0e07401..f042b46 100644
--- a/tallax/vllm/bounded_k/__init__.py
+++ b/tallax/vllm/bounded_k/__init__.py
@@ -8,6 +8,6 @@
 enabled.
 """
 
-# from tallax.vllm.bounded_k.top_p_and_sample import bounded_topk_topp_and_sample
+from tallax.vllm.bounded_k.top_p_and_sample import top_p_and_sample
 
-# __all__ = ["bounded_topk_topp_and_sample"]
+__all__ = ["top_p_and_sample"]
diff --git a/tallax/vllm/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index 566d16e..ad927de 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -29,16 +29,6 @@
 _SAMPLING_EPS = 1e-5
 
 
-def canonicalize_args(batch_size, *args):
-  result = []
-  for val in args:
-    val = jnp.atleast_1d(val)
-    if val.shape[0] == 1:
-      val = jnp.broadcast_to(val, (batch_size,))
-    result.append(val)
-  return result
-
-
 def broadcast_to(x, shape):
   if x.shape[1] == shape[1] and shape[0] % NUM_LANES == 0 and x.shape[0] == 1:
     return pltpu.repeat(
@@ -247,7 +237,9 @@ def _top_p_and_sample(
   topk_logits, topk_idx, random_u128_in_u32s, top_p, temperature, *, debug=False
 ):
   batch_size, k = topk_logits.shape
-  top_p, temperature = canonicalize_args(batch_size, top_p, temperature)
+  top_p, temperature = (
+    jnp.broadcast_to(v, (batch_size,)) for v in (top_p, temperature)
+  )
 
   output_shape = jax.ShapeDtypeStruct((batch_size,), jnp.int32)
 
diff --git a/tallax/vllm/sampling.py b/tallax/vllm/sampling.py
index 9809f5c..75185fb 100644
--- a/tallax/vllm/sampling.py
+++ b/tallax/vllm/sampling.py
@@ -15,7 +15,7 @@
 @functools.partial(
   jax.jit, static_argnames=("max_k", "num_bins", "bins_topm_schedule")
 )
-def topk_topp_and_sample(
+def bounded_topk_topp_and_sample(
   rng_key,
   logits,
   tpu_sampling_metadata,
@@ -38,7 +38,6 @@ def topk_topp_and_sample(
   Returns:
     Sampled token indices.
   """
-  vocab_size = logits.shape[1]
   topk_logits, topk_idxs = top_bounded_k(
     logits,
     k=tpu_sampling_metadata.top_k,
@@ -49,11 +48,9 @@ def topk_topp_and_sample(
     guarantee_convergence=True,
   )
   return top_p_and_sample(
-    topk_logits,
-    topk_idxs,
-    rng_key,
+    topk_logits=topk_logits,
+    topk_idx=topk_idxs,
+    rng_key=rng_key,
     top_p=tpu_sampling_metadata.top_p,
     temperature=tpu_sampling_metadata.temperature,
-    vocab_size=vocab_size,
-    replace_val=-1e12,
   )

From 088365a5dfa01f9ca7d3883ad67a18954a04ca5e Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 22:07:38 +0000
Subject: [PATCH 16/33] fix imports

---
 tallax/vllm/bounded_k/__init__.py | 4 ++--
 tallax/vllm/sampling.py           | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tallax/vllm/bounded_k/__init__.py b/tallax/vllm/bounded_k/__init__.py
index f042b46..90db535 100644
--- a/tallax/vllm/bounded_k/__init__.py
+++ b/tallax/vllm/bounded_k/__init__.py
@@ -8,6 +8,6 @@
 enabled.
 """
 
-from tallax.vllm.bounded_k.top_p_and_sample import top_p_and_sample
+from tallax.vllm.bounded_k.top_p_and_sample import topp_and_sample
 
-__all__ = ["top_p_and_sample"]
+__all__ = ["topp_and_sample"]
diff --git a/tallax/vllm/sampling.py b/tallax/vllm/sampling.py
index 75185fb..4aea4bb 100644
--- a/tallax/vllm/sampling.py
+++ b/tallax/vllm/sampling.py
@@ -8,7 +8,7 @@
 
 import functools
 import jax
-from tallax.vllm.bounded_k import top_p_and_sample
+from tallax.vllm.bounded_k import topp_and_sample
 from tallax.tax.divide_and_filter_topk.topk import top_bounded_k
 
 
@@ -47,7 +47,7 @@ def bounded_topk_topp_and_sample(
     bins_topm_schedule=bins_topm_schedule,
     guarantee_convergence=True,
   )
-  return top_p_and_sample(
+  return topp_and_sample(
     topk_logits=topk_logits,
     topk_idx=topk_idxs,
     rng_key=rng_key,

From 57b5d4abbf93a3680622d7ee89aed1a2ff51fc51 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 22:09:05 +0000
Subject: [PATCH 17/33] fix imports

---
 tallax/vllm/__init__.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/tallax/vllm/__init__.py b/tallax/vllm/__init__.py
index 27aae80..78a5e17 100644
--- a/tallax/vllm/__init__.py
+++ b/tallax/vllm/__init__.py
@@ -8,15 +8,12 @@
   - reference: Pure JAX reference implementation (no Pallas).
 """
 
-from tallax.vllm.sampling import topk_topp_and_sample
-
-# from tallax.vllm.bounded_k import bounded_topk_topp_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",
-  # "bounded_topk_topp_and_sample",
+  "bounded_topk_topp_and_sample",
   "arbitrary_topk_topp_and_sample",
   "reference_topk_topp_and_sample",
 ]

From 203c7a6ce376501aadb4d6ea494aa4e995852ffe Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 22:12:19 +0000
Subject: [PATCH 18/33] pass stable kwarg

---
 tallax/tax/divide_and_filter_topk/topk.py | 24 ++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/tallax/tax/divide_and_filter_topk/topk.py b/tallax/tax/divide_and_filter_topk/topk.py
index 0e481ff..87d8020 100644
--- a/tallax/tax/divide_and_filter_topk/topk.py
+++ b/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):
@@ -98,7 +98,9 @@ def bitonic_topk_arrays(operands, k, num_keys, val_dtype=None, max_index=None):
   pack = val_dtype == jnp.bfloat16 and max_index <= 2**16
   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 = _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))
@@ -178,6 +180,7 @@ 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
@@ -232,7 +235,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),
+    num_keys=1 + int(stable),
   )
   sorted_bin_indices = pad(sorted_bin_indices, (NUM_SUBLANES, NUM_LANES))
 
@@ -338,7 +341,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 +382,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,7 +443,7 @@ 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 = (
@@ -542,7 +545,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 +570,7 @@ def _():
     "guarantee_convergence",
     "replace_val",
     "interpret",
+    "stable",
   ),
 )
 def _top_bounded_k(
@@ -581,6 +585,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 +728,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)),

From fceb2a893f18c49eccb4f15c30357ee66a266646 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 22:22:57 +0000
Subject: [PATCH 19/33] fix stable topk and update README with stable sort
 runtimes

---
 README.md                                 |  4 +-
 tallax/tax/divide_and_filter_topk/topk.py | 64 ++++++++++++++++-------
 2 files changed, 46 insertions(+), 22 deletions(-)

diff --git a/README.md b/README.md
index 3c3ef06..b961124 100644
--- a/README.md
+++ b/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
diff --git a/tallax/tax/divide_and_filter_topk/topk.py b/tallax/tax/divide_and_filter_topk/topk.py
index 87d8020..9181121 100644
--- a/tallax/tax/divide_and_filter_topk/topk.py
+++ b/tallax/tax/divide_and_filter_topk/topk.py
@@ -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,21 +92,34 @@ 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")
+  if num_keys == 2:
+    # sort vals descending, indices ascending so we negate indices
+    operands[1] = -operands[1]
   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=(num_keys > 1))]
-  operands = _bitonic_topk_arrays(
-    operands, k=k, num_keys=min(len(operands), num_keys)
+  pack = (
+    val_dtype is not None
+    and max_index is not None
+    and (val_dtype == jnp.bfloat16 and max_index <= 2**16)
   )
   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 = [
+      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:
+    operands = (
+      list(unpack_bf16_u16_from_i32(operands[0], stable=stable)) + operands[1:]
+    )
+  if num_keys == 2:
+    # 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)]
 
 
@@ -184,7 +198,7 @@ def compute_idxs(i):
   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)
@@ -228,15 +242,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),
-  )
+  )[-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
@@ -447,7 +468,7 @@ def _():
       # 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)
       )
@@ -773,6 +794,7 @@ def _top_bounded_k(
     "guarantee_convergence",
     "replace_val",
     "interpret",
+    "stable",
   ),
 )
 @functools.wraps(_top_bounded_k)
@@ -788,6 +810,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(
@@ -802,6 +825,7 @@ def _closed_topk(logits: jax.Array, k: jax.Array):
       guarantee_convergence=guarantee_convergence,
       replace_val=replace_val,
       interpret=interpret,
+      stable=stable,
     )
 
   @custom_partitioning

From 37aa0ca74919d236911d9d63d4d048a25106c578 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Sun, 15 Feb 2026 22:25:39 +0000
Subject: [PATCH 20/33] closure over bool for jax.eval_shape

---
 tallax/vllm/bounded_k/top_p_and_sample.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tallax/vllm/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index ad927de..7736800 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -323,13 +323,15 @@ def topp_and_sample(
   )
 
   out_shapes = jax.eval_shape(
-    _top_p_and_sample,
+    functools.partial(
+      _top_p_and_sample,
+      debug=debug,
+    ),
     topk_logits,
     topk_idx,
     random_u128_in_u32s,
     top_p,
     temperature,
-    debug=debug,
   )
 
   @custom_partitioning

From ebe3923980285c80425b0b4ae9f29a483d7233f2 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Mon, 16 Feb 2026 08:02:45 +0000
Subject: [PATCH 21/33] fix shapes and cumsum logic

---
 tallax/vllm/bounded_k/top_p_and_sample.py | 30 +++++++++++------------
 1 file changed, 15 insertions(+), 15 deletions(-)

diff --git a/tallax/vllm/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index 7736800..58d718c 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -170,16 +170,18 @@ def top_p_and_sample_arrays(
       unnorm_probs_i32_unsorted.sum(0, keepdims=True).astype(jnp.uint32),
     ],
   )[1]
-  cumsums = cumsum_arrays(unnorm_probs_i32_unsorted, axis=0)
-  threshold_local_idx = sum((c < target_cumsum) for c in cumsums)
+  cumsum = cumsum_arrays(unnorm_probs_i32_unsorted, axis=0)
+  threshold_local_idx = (cumsum < target_cumsum).sum(0, keepdims=True)
 
   next_tokens = (
     idxs
     * (
       jax.lax.broadcasted_iota(jnp.int32, idxs.shape, 0) == threshold_local_idx
     )
-  ).sum(0)
-  result = jnp.where(temperature < _SAMPLING_EPS, greedy_sampled, next_tokens)
+  ).sum(0, keepdims=True)
+  result = jnp.where(
+    temperature[None, :] < _SAMPLING_EPS, greedy_sampled, next_tokens
+  )
 
   if not debug:
     return result
@@ -241,7 +243,7 @@ def _top_p_and_sample(
     jnp.broadcast_to(v, (batch_size,)) for v in (top_p, temperature)
   )
 
-  output_shape = jax.ShapeDtypeStruct((batch_size,), jnp.int32)
+  output_shape = jax.ShapeDtypeStruct((1, batch_size), jnp.int32)
 
   if debug:
     debug_out_shapes = {
@@ -261,7 +263,7 @@ def _top_p_and_sample(
   else:
     debug_out_shapes = None
 
-  results = pl.pallas_call(
+  sampled_tokens, debug_aux = pl.pallas_call(
     top_p_and_sample_refs,
     out_shape=(output_shape, debug_out_shapes),
   )(
@@ -272,16 +274,14 @@ def _top_p_and_sample(
     temperature,
   )
 
-  def trim(x):
-    if x.shape[1] == 1:
-      x = x.squeeze(1)
-    return x[:batch_size]
-
-  sampled_tokens, debug_aux = jax.tree.map(trim, results)
+  sampled_tokens = sampled_tokens.squeeze(0)
+  if not debug:
+    return sampled_tokens
 
-  if debug:
-    return sampled_tokens, debug_aux
-  return sampled_tokens
+  debug_aux = jax.tree.map(
+    lambda x: x.squeeze(1) if x.shape[1] == 1 else x, debug_aux
+  )
+  return sampled_tokens, debug_aux
 
 
 @functools.partial(

From c53e5127212fdd85513f6201d153980db429a725 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Mon, 16 Feb 2026 08:05:28 +0000
Subject: [PATCH 22/33] fix shapes

---
 tallax/vllm/bounded_k/top_p_and_sample.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tallax/vllm/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index 58d718c..165b5eb 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -141,7 +141,7 @@ def top_p_and_sample_arrays(
   shape = topk_logits_transposed.shape
 
   # Greedy sample (before temperature scaling)
-  greedy_sampled = topk_idx[:, :1]
+  greedy_sampled = topk_idx[:1, :]
 
   # Temperature scaling
   topk_logits_scaled = topk_logits_transposed / temperature[None, :].astype(

From fe675f9e575abfe832533aa8f9e1e6406689e4af Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Mon, 16 Feb 2026 08:07:55 +0000
Subject: [PATCH 23/33] fix shapes

---
 tallax/vllm/bounded_k/top_p_and_sample.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tallax/vllm/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index 165b5eb..f25c864 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -141,7 +141,7 @@ def top_p_and_sample_arrays(
   shape = topk_logits_transposed.shape
 
   # Greedy sample (before temperature scaling)
-  greedy_sampled = topk_idx[:1, :]
+  greedy_sampled = topk_idx_transposed[:1, :]
 
   # Temperature scaling
   topk_logits_scaled = topk_logits_transposed / temperature[None, :].astype(
@@ -190,14 +190,14 @@ def top_p_and_sample_arrays(
   random_unnorm_cdf_sampled_low = target_cumsum.T
 
   debug_results = {
-    "greedy_sampled": greedy_sampled,
-    "topk_logits_unsorted": topk_logits.T,  # Original sorted logits in batch-first format
+    "greedy_sampled": greedy_sampled.T,
+    "topk_logits_unsorted": topk_logits,  # Original sorted logits in batch-first format
     "topk_topp_unnorm_probs_i32_unsorted": unnorm_probs_i32_unsorted.T,  # Transpose back to [batch, k]
     "random_unnorm_cdf_sampled": (
       jnp.zeros_like(random_unnorm_cdf_sampled_low),
       random_unnorm_cdf_sampled_low,
     ),
-    "next_tokens": next_tokens,
+    "next_tokens": next_tokens.T,
   }
 
   return result, debug_results

From dff7d1eaf4a6734a6a32fa69be10d17c0d2d2d7f Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Mon, 16 Feb 2026 08:10:00 +0000
Subject: [PATCH 24/33] fix num operands logic

---
 tallax/vllm/bounded_k/top_p_and_sample.py | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/tallax/vllm/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index f25c864..8a2f69e 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -389,9 +389,4 @@ def shmap_fn(
   flat_outs = sharded_top_p_and_sample(
     topk_logits, topk_idx, random_u128_in_u32s, top_p, temperature
   )
-  sampled_tokens, debug_aux = jax.tree.unflatten(
-    jax.tree.structure(out_shapes), flat_outs
-  )
-  if debug:
-    return sampled_tokens, debug_aux
-  return sampled_tokens
+  return jax.tree.unflatten(jax.tree.structure(out_shapes), flat_outs)

From 3721165e4db6e6f608c8b398e74e22fe27df3ad2 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Mon, 16 Feb 2026 08:32:27 +0000
Subject: [PATCH 25/33] fix passing remainder through topk logic

---
 tallax/tax/divide_and_filter_topk/topk.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tallax/tax/divide_and_filter_topk/topk.py b/tallax/tax/divide_and_filter_topk/topk.py
index 9181121..16707d6 100644
--- a/tallax/tax/divide_and_filter_topk/topk.py
+++ b/tallax/tax/divide_and_filter_topk/topk.py
@@ -205,6 +205,7 @@ def loop_body(i, 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
@@ -215,7 +216,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,
   )
 

From 76dab20195baf4bcbbc877e265965463a0dc1723 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Mon, 16 Feb 2026 15:50:06 +0000
Subject: [PATCH 26/33] avoid inverting indices for bf16 packing

---
 tallax/tax/divide_and_filter_topk/topk.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tallax/tax/divide_and_filter_topk/topk.py b/tallax/tax/divide_and_filter_topk/topk.py
index 16707d6..73dec43 100644
--- a/tallax/tax/divide_and_filter_topk/topk.py
+++ b/tallax/tax/divide_and_filter_topk/topk.py
@@ -96,15 +96,15 @@ def bitonic_topk_arrays(operands, k, num_keys, val_dtype=None, max_index=None):
   stable = num_keys > 1
   if not jnp.issubdtype(operands[1].dtype, jnp.integer):
     raise ValueError("Expected second operand to be indices, so integer dtype")
-  if num_keys == 2:
-    # sort vals descending, indices ascending so we negate indices
-    operands[1] = -operands[1]
   dtypes = [x.dtype for x in operands]
   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[:2], stable=stable),
@@ -116,7 +116,7 @@ def bitonic_topk_arrays(operands, k, num_keys, val_dtype=None, max_index=None):
     operands = (
       list(unpack_bf16_u16_from_i32(operands[0], stable=stable)) + operands[1:]
     )
-  if num_keys == 2:
+  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)

From 6352ed1cb8873f4cc299fb99a1754ac39b22ecef Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Mon, 16 Feb 2026 15:57:26 +0000
Subject: [PATCH 27/33] rename idx->idxs to reflect plural

---
 tallax/vllm/bounded_k/top_p_and_sample.py | 27 ++++++++++-------------
 1 file changed, 12 insertions(+), 15 deletions(-)

diff --git a/tallax/vllm/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index 8a2f69e..a91b7ac 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -110,7 +110,7 @@ def top_p_mask(*, topk_logits, p, replace_val, axis):
 def top_p_and_sample_arrays(
   *,
   topk_logits,
-  topk_idx,
+  topk_idxs,
   random_u128_in_u32s,
   top_p,
   temperature,
@@ -131,17 +131,14 @@ def top_p_and_sample_arrays(
   """
   topk_logits = topk_logits.astype(jnp.float32)
 
-  # Store original shape for debug
-  batch_size, k = topk_logits.shape
-
   # Shift to dim 0 for sublane-based reductions
   topk_logits_transposed = topk_logits.T
-  topk_idx_transposed = topk_idx.T
+  topk_idxs_transposed = topk_idxs.T
   random_u128_in_u32s_transposed = [x.T for x in random_u128_in_u32s]
   shape = topk_logits_transposed.shape
 
   # Greedy sample (before temperature scaling)
-  greedy_sampled = topk_idx_transposed[:1, :]
+  greedy_sampled = topk_idxs_transposed[:1, :]
 
   # Temperature scaling
   topk_logits_scaled = topk_logits_transposed / temperature[None, :].astype(
@@ -155,8 +152,8 @@ def top_p_and_sample_arrays(
 
   # Re-sort back to original index order for bitwise-matching sampling
   inverted_idxs, unnorm_probs_i32_unsorted = bitonic_topk_arrays(
-    [-topk_idx_transposed, unnorm_probs_i32_sorted],
-    k=topk_idx_transposed.shape[0],
+    [-topk_idxs_transposed, unnorm_probs_i32_sorted],
+    k=topk_idxs_transposed.shape[0],
     axis=0,
     num_keys=1,
   )
@@ -191,8 +188,9 @@ def top_p_and_sample_arrays(
 
   debug_results = {
     "greedy_sampled": greedy_sampled.T,
-    "topk_logits_unsorted": topk_logits,  # Original sorted logits in batch-first format
-    "topk_topp_unnorm_probs_i32_unsorted": unnorm_probs_i32_unsorted.T,  # Transpose back to [batch, k]
+    "topk_logits": topk_logits,  # Top-k logits
+    "topk_idxs": topk_idxs,
+    "topk_topp_unnorm_probs_i32_topk_filtered_unsorted": unnorm_probs_i32_unsorted.T,  # Transpose back to [batch, k]
     "random_unnorm_cdf_sampled": (
       jnp.zeros_like(random_unnorm_cdf_sampled_low),
       random_unnorm_cdf_sampled_low,
@@ -215,7 +213,7 @@ def top_p_and_sample_refs(
   """Pallas kernel body for top-p filtering + sampling."""
   result = top_p_and_sample_arrays(
     topk_logits=topk_logits_ref[...],
-    topk_idx=topk_idx_ref[...],
+    topk_idxs=topk_idx_ref[...],
     random_u128_in_u32s=[ref[...] for ref in random_u128_in_u32s_refs],
     top_p=top_p_ref[...],
     temperature=temperature_ref[...],
@@ -248,10 +246,9 @@ def _top_p_and_sample(
   if debug:
     debug_out_shapes = {
       "greedy_sampled": jax.ShapeDtypeStruct((batch_size, 1), jnp.int32),
-      "topk_logits_unsorted": jax.ShapeDtypeStruct(
-        (batch_size, k), jnp.float32
-      ),
-      "topk_topp_unnorm_probs_i32_unsorted": jax.ShapeDtypeStruct(
+      "topk_logits": jax.ShapeDtypeStruct((batch_size, k), jnp.float32),
+      "topk_idxs": jax.ShapeDtypeStruct((batch_size, k), jnp.int32),
+      "topk_topp_unnorm_probs_i32_topk_filtered_unsorted": jax.ShapeDtypeStruct(
         (batch_size, k), jnp.int32
       ),
       "random_unnorm_cdf_sampled": (

From bcde5f99d4570c79291161338dca3e79e2af739a Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Mon, 16 Feb 2026 16:30:15 +0000
Subject: [PATCH 28/33] fix bf16 stable sort

---
 tallax/tax/divide_and_filter_topk/topk.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tallax/tax/divide_and_filter_topk/topk.py b/tallax/tax/divide_and_filter_topk/topk.py
index 73dec43..f377cf5 100644
--- a/tallax/tax/divide_and_filter_topk/topk.py
+++ b/tallax/tax/divide_and_filter_topk/topk.py
@@ -106,6 +106,7 @@ def bitonic_topk_arrays(operands, k, num_keys, val_dtype=None, max_index=None):
     # sort vals descending, indices ascending so we negate indices
     operands[1] = -operands[1]
   if pack:
+    operands[1] = 2**16 - 1 - operands[1]
     operands = [
       pack_bf16_u16_to_i32(*operands[:2], stable=stable),
       *operands[2:],
@@ -116,6 +117,7 @@ def bitonic_topk_arrays(operands, k, num_keys, val_dtype=None, max_index=None):
     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]

From 35c90c0aadabd9b622c788951c7e4941860e6038 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Mon, 16 Feb 2026 16:41:04 +0000
Subject: [PATCH 29/33] control topk_logits dtype

---
 tallax/vllm/arbitrary_k/kernel.py | 5 +++--
 tallax/vllm/reference.py          | 5 ++++-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/tallax/vllm/arbitrary_k/kernel.py b/tallax/vllm/arbitrary_k/kernel.py
index eac7685..91503fa 100644
--- a/tallax/vllm/arbitrary_k/kernel.py
+++ b/tallax/vllm/arbitrary_k/kernel.py
@@ -178,7 +178,8 @@ def scoped_body(scoped_ref):
   # Debug: write full intermediate arrays if debug refs are provided
   if debug_arrays_ref is not None:
     debug_arrays_ref["greedy_sampled"][...] = greedy_sampled
-    debug_arrays_ref["topk_logits_unsorted"][...] = topk_logits_pre_temperature
+    ref = debug_arrays_ref["topk_logits_unsorted"]
+    ref[...] = topk_logits_pre_temperature.astype(ref.dtype)
     debug_arrays_ref["topk_topp_unnorm_probs_i32_unsorted"][...] = (
       unnorm_probs_i32
     )
@@ -256,7 +257,7 @@ def arbitrary_topk_topp_and_sample(
     debug_out_shapes = {
       "greedy_sampled": jax.ShapeDtypeStruct((padded_batch, 1), jnp.int32),
       "topk_logits_unsorted": jax.ShapeDtypeStruct(
-        (padded_batch, vocab_size), jnp.float32
+        (padded_batch, vocab_size), logits.dtype
       ),
       "topk_topp_unnorm_probs_i32_unsorted": jax.ShapeDtypeStruct(
         (padded_batch, vocab_size), jnp.int32
diff --git a/tallax/vllm/reference.py b/tallax/vllm/reference.py
index 2b2909b..11c4bdd 100644
--- a/tallax/vllm/reference.py
+++ b/tallax/vllm/reference.py
@@ -73,6 +73,7 @@ def reference_topk_topp_and_sample(
     Sampled token indices [batch], or (tokens, debug_dict) if debug=True
   """
   assert stable
+  underlying_logits_dtype = logits.dtype
   with jax.enable_x64(True):
     shape = logits.shape
     scale = 2**_SCALE_BITS - 1
@@ -139,7 +140,9 @@ def reference_topk_topp_and_sample(
 
     debug_results = {
       "greedy_sampled": greedy_sampled,
-      "topk_logits_unsorted": topk_logits_unsorted,
+      "topk_logits_unsorted": topk_logits_unsorted.astype(
+        underlying_logits_dtype
+      ),
       "topk_topp_unnorm_probs_i32_unsorted": unnorm_probs_i32,
       "random_unnorm_cdf_sampled": (
         sampled_total_high.squeeze(1),

From ec1bf52f6b5d96f16032983009716ee9d492691f Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Mon, 16 Feb 2026 19:44:03 +0000
Subject: [PATCH 30/33] hardcode consts so they're tied between implementations
 strongly

---
 tallax/constants.py                           |  3 +++
 tallax/vllm/arbitrary_k/kernel.py             | 18 ++++--------------
 tallax/vllm/arbitrary_k/topk_mask.py          |  5 ++---
 tallax/vllm/arbitrary_k/topp_mask.py          |  7 +++----
 tallax/vllm/bounded_k/top_p_and_sample.py     | 18 ++++++------------
 tallax/vllm/reference.py                      | 12 ++++--------
 tallax/vllm/sampling.py                       |  3 ++-
 tests/vllm/arbitrary_k/topk_mask_test.py      |  8 ++++----
 tests/vllm/bounded_k/top_p_and_sample_test.py |  7 +++----
 tests/vllm/utils/high_precision_uint_test.py  |  3 ++-
 10 files changed, 33 insertions(+), 51 deletions(-)
 create mode 100644 tallax/constants.py

diff --git a/tallax/constants.py b/tallax/constants.py
new file mode 100644
index 0000000..c5032af
--- /dev/null
+++ b/tallax/constants.py
@@ -0,0 +1,3 @@
+REPLACE_VAL = -1e12
+SAMPLING_EPS = 1e-5
+SCALE_BITS = 24
diff --git a/tallax/vllm/arbitrary_k/kernel.py b/tallax/vllm/arbitrary_k/kernel.py
index 91503fa..bdf3916 100644
--- a/tallax/vllm/arbitrary_k/kernel.py
+++ b/tallax/vllm/arbitrary_k/kernel.py
@@ -25,12 +25,11 @@
   modulo_u128_u64,
   sample_random_u128_in_u32s,
 )
+from tallax.constants import SAMPLING_EPS, SCALE_BITS
 from tallax.vllm.arbitrary_k.topp_mask import topp_mask
 from tallax.vllm.arbitrary_k.topk_mask import topk_mask, find_boundary_idx
 from tallax.tax.utils import NUM_LANES, map_reduce, get_dtype_info
 
-_SAMPLING_EPS = 1e-5
-
 
 def canonicalize_args(batch_size, block_token, *args_with_pad):
   """Canonicalize per-sample args and pad batch to block_token multiple.
@@ -58,17 +57,17 @@ def canonicalize_args(batch_size, block_token, *args_with_pad):
   return (padded_batch, num_blocks, *result)
 
 
-def sample_probs(unnorm_probs_i32, random_u128_in_u32s, max_val=2**24 - 1):
+def sample_probs(unnorm_probs_i32, random_u128_in_u32s):
   """Sample from unnormalized i32 probabilities using U48 arithmetic.
 
   Args:
     unnorm_probs_i32: Unnormalized probabilities [batch, vocab_size] in i32
     random_u128_in_u32s: List of 4 u32 arrays for random sampling
-    max_val: Maximum bound of unnorm_probs_i32 values
 
   Returns:
     Sampled token indices [batch, NUM_LANES]
   """
+  max_val = 2**SCALE_BITS - 1
   total_sum_u48 = U48.map_reduce_sum(unnorm_probs_i32, max_val=max_val)
   sampled_u64_in_u32s = modulo_u128_u64(
     random_u128_in_u32s,
@@ -93,7 +92,6 @@ def arbitrary_topk_topp_and_sample_kernel(
   debug_arrays_ref=None,
   *,
   stable: bool,
-  replace_val: float,
   underlying_logits_dtype=None,
 ):
   """Pallas kernel body for combined topk + topp + sample.
@@ -107,7 +105,6 @@ def arbitrary_topk_topp_and_sample_kernel(
     sampled_tokens_ref: Output sampled tokens [block_token, 1]
     debug_arrays_ref: Optional dict of refs for full intermediate arrays
     stable: Whether to use stable masking
-    replace_val: Replacement value for masked elements
     underlying_logits_dtype: Original dtype if logits were cast
   """
   if logits_ref.dtype != jnp.float32:
@@ -123,7 +120,6 @@ def scoped_body(scoped_ref):
         sampled_tokens_ref,
         debug_arrays_ref=debug_arrays_ref,
         stable=stable,
-        replace_val=replace_val,
         underlying_logits_dtype=logits_ref.dtype,
       )
 
@@ -148,7 +144,6 @@ def scoped_body(scoped_ref):
   topk_logits = topk_mask(
     logits_ref,
     k_ref,
-    replace_val=replace_val,
     stable=stable,
     underlying_dtype=underlying_logits_dtype,
   )
@@ -169,10 +164,9 @@ def scoped_body(scoped_ref):
   next_tokens, random_unnorm_cdf_sampled = sample_probs(
     unnorm_probs_i32,
     [ref[...] for ref in random_u128_in_u32s_refs],
-    max_val=2**24 - 1,
   )
   sampled_tokens_ref[...] = jnp.where(
-    temperature < _SAMPLING_EPS, greedy_sampled, next_tokens
+    temperature < SAMPLING_EPS, greedy_sampled, next_tokens
   )
 
   # Debug: write full intermediate arrays if debug refs are provided
@@ -197,7 +191,6 @@ def scoped_body(scoped_ref):
   jax.jit,
   static_argnames=[
     "stable",
-    "replace_val",
     "block_token",
     "interpret",
     "debug",
@@ -211,7 +204,6 @@ def arbitrary_topk_topp_and_sample(
   temperature: jax.Array,
   *,
   stable: bool = True,
-  replace_val: float = -1e12,
   block_token: int = 8,
   interpret: bool = False,
   debug: bool = False,
@@ -228,7 +220,6 @@ def arbitrary_topk_topp_and_sample(
     p: Top-p values [batch] or scalar
     temperature: Temperature values [batch] or scalar
     stable: Whether to use stable masking
-    replace_val: Replacement value for masked elements
     block_token: Number of tokens per block
     interpret: Whether to use interpret mode
     debug: If True, return (sampled_tokens, debug_results) dict
@@ -290,7 +281,6 @@ def arbitrary_topk_topp_and_sample(
     functools.partial(
       arbitrary_topk_topp_and_sample_kernel,
       stable=stable,
-      replace_val=replace_val,
     ),
     out_shape=(output_shape, debug_out_shapes),
     grid=(num_blocks,),
diff --git a/tallax/vllm/arbitrary_k/topk_mask.py b/tallax/vllm/arbitrary_k/topk_mask.py
index 743f9d6..472ca4c 100644
--- a/tallax/vllm/arbitrary_k/topk_mask.py
+++ b/tallax/vllm/arbitrary_k/topk_mask.py
@@ -12,6 +12,7 @@
 from jax.experimental import pallas as pl
 from jax.experimental.pallas import tpu as pltpu
 
+from tallax.constants import REPLACE_VAL
 from tallax.vllm.utils.binary_search import binary_search
 from tallax.tax.utils import (
   NUM_LANES,
@@ -148,7 +149,6 @@ def topk_mask(
   logits_ref,
   k_ref,
   *,
-  replace_val: float,
   stable: bool,
   underlying_dtype=None,
 ):
@@ -157,7 +157,6 @@ def topk_mask(
   Args:
     logits_ref: Input logits reference [batch, vocab_size]
     k_ref: Number of top elements to keep [batch, 1]
-    replace_val: Replacement value for masked elements
     stable: Whether to use stable masking (exactly k elements kept)
     underlying_dtype: Original dtype if logits were cast (for bf16 search convergence)
 
@@ -221,4 +220,4 @@ def topk_mask(
         jax.lax.broadcasted_iota(jnp.int32, logits_ref.shape, 1) <= boundary_idx
       )
     )
-  return jnp.where(mask, logits, replace_val).astype(logits_ref.dtype)
+  return jnp.where(mask, logits, REPLACE_VAL).astype(logits_ref.dtype)
diff --git a/tallax/vllm/arbitrary_k/topp_mask.py b/tallax/vllm/arbitrary_k/topp_mask.py
index 8c6f249..c1d5802 100644
--- a/tallax/vllm/arbitrary_k/topp_mask.py
+++ b/tallax/vllm/arbitrary_k/topp_mask.py
@@ -13,6 +13,7 @@
 import jax.numpy as jnp
 from jax.experimental.pallas import tpu as pltpu
 
+from tallax.constants import SCALE_BITS
 from tallax.tax.utils import NUM_LANES, map_reduce
 from tallax.vllm.utils.high_precision_uint import U48
 from tallax.vllm.utils.binary_search import binary_search
@@ -30,7 +31,6 @@ def topp_mask(
   logits: jax.Array,
   top_p: jax.Array,
   *,
-  scale_bits: int = 24,
   logits_max: jax.Array = None,
 ) -> jax.Array:
   """Apply top-p mask using binary search over integer probability space.
@@ -38,7 +38,6 @@ def topp_mask(
   Args:
     logits: Input logits [batch, vocab_size], float32
     top_p: Top-p threshold [batch, 1] or broadcastable
-    scale_bits: Precision bits for probability scaling (default 24)
     logits_max: Pre-computed max logits [batch, 1] (optional, computed if None)
 
   Returns:
@@ -52,7 +51,7 @@ def topp_mask(
   if logits_max is None:
     logits_max = map_reduce(logits, reduce_fn="max")
 
-  scale = 2**scale_bits - 1
+  scale = 2**SCALE_BITS - 1
   unnorm_probs_i32 = map_chunks(
     logits,
     lambda logits: (jnp.exp(logits - logits_max) * scale).astype(jnp.int32),
@@ -80,7 +79,7 @@ def predicate_fn(threshold):
   threshold_i32, _, _ = binary_search(
     predicate_fn,
     *(jnp.full(bound_shape, v, jnp.int32) for v in (0, scale)),
-    num_iter=scale_bits,
+    num_iter=SCALE_BITS,
   )
 
   # Apply mask
diff --git a/tallax/vllm/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index a91b7ac..6a30208 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -22,12 +22,11 @@
 from tallax.vllm.utils.high_precision_uint import modulo_u128_u64
 from tallax.vllm.utils.high_precision_uint import sample_random_u128_in_u32s
 from tallax.tax.bitonic.topk import bitonic_topk_arrays
+from tallax.constants import REPLACE_VAL, SAMPLING_EPS, SCALE_BITS
 from tallax.tax.cumsum import cumsum_arrays
 from tallax.tax.gather import take_along_axis_arrays
 from tallax.tax.utils import NUM_SUBLANES, NUM_LANES
 
-_SAMPLING_EPS = 1e-5
-
 
 def broadcast_to(x, shape):
   if x.shape[1] == shape[1] and shape[0] % NUM_LANES == 0 and x.shape[0] == 1:
@@ -59,7 +58,7 @@ def top_p_integer_mask(*, topk_logits, p, axis):
   shape = topk_logits.shape
 
   exp_logits = jnp.exp(topk_logits - topk_logits[:1, :])
-  scale = 2**24 - 1
+  scale = 2**SCALE_BITS - 1
   unnorm_probs_i32 = (exp_logits * scale).astype(jnp.int32)
   if unnorm_probs_i32.shape[axis] > 2**7:
     raise NotImplementedError(
@@ -79,17 +78,16 @@ def top_p_integer_mask(*, topk_logits, p, axis):
   return jnp.where(unnorm_probs_i32 >= thresholds, unnorm_probs_i32, 0)
 
 
-def top_p_mask(*, topk_logits, p, replace_val, axis):
+def top_p_mask(*, topk_logits, p, axis):
   """Apply top-p filtering mask to sorted logits (float version for backwards compat).
 
   Args:
     topk_logits: Sorted logits (descending order)
     p: Top-p threshold(s)
-    replace_val: Value to replace filtered logits with
     axis: Axis along which to apply filtering (must be 0)
 
   Returns:
-    Masked logits with values outside top-p set to replace_val
+    Masked logits with values outside top-p set to REPLACE_VAL
   """
   if axis != 0:
     raise NotImplementedError("topp_mask only supports axis=0")
@@ -104,7 +102,7 @@ def top_p_mask(*, topk_logits, p, replace_val, axis):
   thresholds = take_along_axis_arrays(
     topk_logits, broadcast_to(threshold_idx, shape), axis=0
   )
-  return jnp.where(topk_logits >= thresholds, topk_logits, replace_val)
+  return jnp.where(topk_logits >= thresholds, topk_logits, REPLACE_VAL)
 
 
 def top_p_and_sample_arrays(
@@ -177,7 +175,7 @@ def top_p_and_sample_arrays(
     )
   ).sum(0, keepdims=True)
   result = jnp.where(
-    temperature[None, :] < _SAMPLING_EPS, greedy_sampled, next_tokens
+    temperature[None, :] < SAMPLING_EPS, greedy_sampled, next_tokens
   )
 
   if not debug:
@@ -304,10 +302,6 @@ def topp_and_sample(
     rng_key: RNG key for sampling.
     top_p: Top-p threshold values.
     temperature: Temperature values.
-    vocab_size: Total vocabulary size.
-    max_k: Maximum k value (bounded implementation supports k <= 128).
-    replace_val: Value to replace filtered logits with.
-    interpret: If True, run in CPU interpret mode.
     debug: If True, return (tokens, debug_results) with intermediate values.
 
   Returns:
diff --git a/tallax/vllm/reference.py b/tallax/vllm/reference.py
index 11c4bdd..51054fb 100644
--- a/tallax/vllm/reference.py
+++ b/tallax/vllm/reference.py
@@ -12,9 +12,7 @@
 import jax.numpy as jnp
 
 from tallax.vllm.utils.high_precision_uint import sample_random_u128_in_u32s
-
-_SAMPLING_EPS = 1e-5
-_SCALE_BITS = 24
+from tallax.constants import REPLACE_VAL, SAMPLING_EPS, SCALE_BITS
 
 
 def _python_u128_modulo_u64_impl(r0, r1, r2, r3, m0, m1):
@@ -52,7 +50,6 @@ def reference_topk_topp_and_sample(
   temperature: jax.Array,
   *,
   stable: bool = True,
-  replace_val: float = -1e12,
   debug: bool = False,
 ) -> jax.Array:
   """Reference implementation of topk + topp + sample in pure JAX.
@@ -66,7 +63,6 @@ def reference_topk_topp_and_sample(
     p: Top-p values [batch] or scalar
     temperature: Temperature values [batch] or scalar
     stable: Must be True (reference assumes stable)
-    replace_val: Replacement value for masked elements
     debug: If True, return (tokens, debug_results) with intermediate values
 
   Returns:
@@ -76,7 +72,7 @@ def reference_topk_topp_and_sample(
   underlying_logits_dtype = logits.dtype
   with jax.enable_x64(True):
     shape = logits.shape
-    scale = 2**_SCALE_BITS - 1
+    scale = 2**SCALE_BITS - 1
 
     logits = logits.astype(jnp.float32)
 
@@ -87,7 +83,7 @@ def reference_topk_topp_and_sample(
     sorted_indices = logits.argsort(axis=1, stable=True, descending=True)
     sorted_logits = jnp.take_along_axis(logits, sorted_indices, axis=1)
     sorted_logits = jnp.where(
-      jnp.arange(shape[1])[None, :] < k[:, None], sorted_logits, replace_val
+      jnp.arange(shape[1])[None, :] < k[:, None], sorted_logits, REPLACE_VAL
     )
     topk_logits_unsorted = jnp.take_along_axis(
       sorted_logits, sorted_indices.argsort(), axis=1
@@ -129,7 +125,7 @@ def reference_topk_topp_and_sample(
       ],
     )
     next_tokens = (unnorm_probs_i32.cumsum(axis=1) < sampled_total).sum(axis=1)
-    result = jnp.where(temperature < _SAMPLING_EPS, greedy_sampled, next_tokens)
+    result = jnp.where(temperature < SAMPLING_EPS, greedy_sampled, next_tokens)
 
     if not debug:
       return result
diff --git a/tallax/vllm/sampling.py b/tallax/vllm/sampling.py
index 4aea4bb..676a902 100644
--- a/tallax/vllm/sampling.py
+++ b/tallax/vllm/sampling.py
@@ -8,6 +8,7 @@
 
 import functools
 import jax
+from tallax.constants import REPLACE_VAL
 from tallax.vllm.bounded_k import topp_and_sample
 from tallax.tax.divide_and_filter_topk.topk import top_bounded_k
 
@@ -41,7 +42,7 @@ def bounded_topk_topp_and_sample(
   topk_logits, topk_idxs = top_bounded_k(
     logits,
     k=tpu_sampling_metadata.top_k,
-    replace_val=-1e12,
+    replace_val=REPLACE_VAL,
     max_k=max_k,
     num_bins=num_bins,
     bins_topm_schedule=bins_topm_schedule,
diff --git a/tests/vllm/arbitrary_k/topk_mask_test.py b/tests/vllm/arbitrary_k/topk_mask_test.py
index 6907990..193d0de 100644
--- a/tests/vllm/arbitrary_k/topk_mask_test.py
+++ b/tests/vllm/arbitrary_k/topk_mask_test.py
@@ -5,6 +5,7 @@
 import jax.numpy as jnp
 import numpy as np
 
+from tallax.constants import REPLACE_VAL
 from tallax.vllm.arbitrary_k.topk_mask import topk_mask_pallas
 from tallax.tax.utils import is_cpu_platform
 
@@ -17,16 +18,15 @@ def test_topk_mask(seed, k_val):
   key = jax.random.PRNGKey(seed)
   batch_size, vocab_size = 4, 256
   logits = jax.random.normal(key, (batch_size, vocab_size), dtype=jnp.float32)
-  replace_val = -1e12
   k = jnp.full((batch_size,), k_val, dtype=jnp.int32)
-  masked = topk_mask_pallas(logits, k, replace_val=replace_val, stable=True)
+  masked = topk_mask_pallas(logits, k, stable=True)
 
   # Check count
-  counts = (masked != replace_val).sum(axis=1)
+  counts = (masked != REPLACE_VAL).sum(axis=1)
   np.testing.assert_array_equal(counts, jnp.full((batch_size,), k_val))
 
   # Check values match jax.lax.top_k
   for b in range(batch_size):
     ref_vals, _ = jax.lax.top_k(logits[b], k_val)
-    actual_vals = jnp.sort(masked[b][masked[b] != replace_val])[::-1]
+    actual_vals = jnp.sort(masked[b][masked[b] != REPLACE_VAL])[::-1]
     np.testing.assert_allclose(actual_vals[:k_val], ref_vals, atol=1e-5)
diff --git a/tests/vllm/bounded_k/top_p_and_sample_test.py b/tests/vllm/bounded_k/top_p_and_sample_test.py
index 6133801..a67ac9c 100644
--- a/tests/vllm/bounded_k/top_p_and_sample_test.py
+++ b/tests/vllm/bounded_k/top_p_and_sample_test.py
@@ -8,6 +8,7 @@
 import jax.numpy as jnp
 import numpy as np
 
+from tallax.constants import REPLACE_VAL
 from tallax.vllm.bounded_k.top_p_and_sample import top_p_mask as pallas_top_p_mask
 from tallax.vllm.tpu_inference_sampling_as_standalone_file import (
   topp_mask as tpu_inference_top_p_mask,
@@ -29,13 +30,11 @@ def test_top_p_mask(shape, seed, p_threshold):
   else:
     p_array = jnp.full(shape[:1], p_threshold, dtype=jnp.float32)
 
-  replace_val = -1e12
-
   sort_indices = jnp.argsort(logits, axis=1, descending=True)
   sorted_logits = jnp.take_along_axis(logits, sort_indices, axis=1)
 
   result_sorted = pallas_top_p_mask(
-    topk_logits=sorted_logits.T, p=p_array, replace_val=replace_val, axis=0,
+    topk_logits=sorted_logits.T, p=p_array, axis=0,
   ).T
 
   inverse_sort_indices = jnp.argsort(sort_indices, axis=1)
@@ -44,7 +43,7 @@ def test_top_p_mask(shape, seed, p_threshold):
   result_ref = jnp.zeros_like(logits)
   for i in range(shape[0]):
     result_ref = result_ref.at[i].set(
-      tpu_inference_top_p_mask(logits[i:i+1], float(p_array[i]), replace_val)[0]
+      tpu_inference_top_p_mask(logits[i:i+1], float(p_array[i]), REPLACE_VAL)[0]
     )
 
   np.testing.assert_array_equal(result_original_order, result_ref)
diff --git a/tests/vllm/utils/high_precision_uint_test.py b/tests/vllm/utils/high_precision_uint_test.py
index 15e9521..744d8cb 100644
--- a/tests/vllm/utils/high_precision_uint_test.py
+++ b/tests/vllm/utils/high_precision_uint_test.py
@@ -5,6 +5,7 @@
 import jax.numpy as jnp
 import numpy as np
 
+from tallax.constants import SCALE_BITS
 from tallax.vllm.utils.high_precision_uint import U48, modulo_u128_u64
 
 
@@ -12,7 +13,7 @@
 def test_u48(seed):
   """U48.map_reduce_sum matches i64 sum and < operator is consistent."""
   key = jax.random.PRNGKey(seed)
-  scale = 2**24 - 1
+  scale = 2**SCALE_BITS - 1
   vals = jax.random.randint(key, (4, 512), 0, scale, dtype=jnp.int32)
   u48_sum = U48.map_reduce_sum(vals, max_val=scale)
 

From e4208216a0546497c8fc79d6f27d1556a4c10b9b Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Mon, 16 Feb 2026 19:49:54 +0000
Subject: [PATCH 31/33] add debug path

---
 tallax/vllm/sampling.py | 45 ++++++++++++++++++++++++++++++++++-------
 1 file changed, 38 insertions(+), 7 deletions(-)

diff --git a/tallax/vllm/sampling.py b/tallax/vllm/sampling.py
index 676a902..c4a4d83 100644
--- a/tallax/vllm/sampling.py
+++ b/tallax/vllm/sampling.py
@@ -8,21 +8,25 @@
 
 import functools
 import jax
+from jax import numpy as jnp
 from tallax.constants import REPLACE_VAL
 from tallax.vllm.bounded_k import topp_and_sample
 from tallax.tax.divide_and_filter_topk.topk import top_bounded_k
 
 
 @functools.partial(
-  jax.jit, static_argnames=("max_k", "num_bins", "bins_topm_schedule")
+  jax.jit, static_argnames=("max_k", "num_bins", "bins_topm_schedule", "debug")
 )
 def bounded_topk_topp_and_sample(
-  rng_key,
   logits,
-  tpu_sampling_metadata,
+  rng_key,
+  top_k,
+  top_p,
+  temperature,
   max_k: int,
   num_bins: int | None = None,
   bins_topm_schedule: int | None = None,
+  debug: bool = False,
 ):
   """Combined top-k, top-p filtering, and sampling for vLLM inference.
 
@@ -41,17 +45,44 @@ def bounded_topk_topp_and_sample(
   """
   topk_logits, topk_idxs = top_bounded_k(
     logits,
-    k=tpu_sampling_metadata.top_k,
+    k=top_k,
     replace_val=REPLACE_VAL,
     max_k=max_k,
     num_bins=num_bins,
     bins_topm_schedule=bins_topm_schedule,
     guarantee_convergence=True,
   )
-  return topp_and_sample(
+  outs = topp_and_sample(
     topk_logits=topk_logits,
     topk_idx=topk_idxs,
     rng_key=rng_key,
-    top_p=tpu_sampling_metadata.top_p,
-    temperature=tpu_sampling_metadata.temperature,
+    top_p=top_p,
+    temperature=temperature,
+    debug=debug,
+  )
+  if not debug:
+    return outs
+  sampled, debug_vals = outs
+  # Rebuild the unreduced shape arrays
+  debug_vals["topk_logits_unsorted"] = jax.vmap(
+    lambda ind, updates: jnp.full_like(
+      updates, REPLACE_VAL, shape=logits.shape[1:]
+    )
+    .at[ind]
+    .set(updates)
+  )(debug_vals["topk_idxs"], debug_vals["topk_logits"])
+  debug_vals["topk_topp_unnorm_probs_i32_unsorted"] = jax.vmap(
+    lambda ind, updates: jnp.zeros_like(
+      updates,
+      shape=logits.shape[1:],
+    )
+    .at[ind]
+    .set(updates)
+  )(
+    debug_vals["topk_idxs"].sort(1),
+    debug_vals["topk_topp_unnorm_probs_i32_topk_filtered_unsorted"],
   )
+  del debug_vals["topk_idxs"]
+  del debug_vals["topk_logits"]
+  del debug_vals["topk_topp_unnorm_probs_i32_topk_filtered_unsorted"]
+  return sampled, debug_vals

From 5cdc082ac090eb7d6266780ad30c0551cf25fa64 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Mon, 16 Feb 2026 19:59:29 +0000
Subject: [PATCH 32/33] add stable control for topk

---
 tallax/vllm/sampling.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/tallax/vllm/sampling.py b/tallax/vllm/sampling.py
index c4a4d83..7943924 100644
--- a/tallax/vllm/sampling.py
+++ b/tallax/vllm/sampling.py
@@ -15,7 +15,14 @@
 
 
 @functools.partial(
-  jax.jit, static_argnames=("max_k", "num_bins", "bins_topm_schedule", "debug")
+  jax.jit,
+  static_argnames=(
+    "max_k",
+    "num_bins",
+    "bins_topm_schedule",
+    "debug",
+    "stable",
+  ),
 )
 def bounded_topk_topp_and_sample(
   logits,
@@ -25,6 +32,7 @@ def bounded_topk_topp_and_sample(
   temperature,
   max_k: int,
   num_bins: int | None = None,
+  stable: bool = True,
   bins_topm_schedule: int | None = None,
   debug: bool = False,
 ):
@@ -51,6 +59,7 @@ def bounded_topk_topp_and_sample(
     num_bins=num_bins,
     bins_topm_schedule=bins_topm_schedule,
     guarantee_convergence=True,
+    stable=stable,
   )
   outs = topp_and_sample(
     topk_logits=topk_logits,

From 91d0acff1e16b30428f4cf6ae2da68b4fa703384 Mon Sep 17 00:00:00 2001
From: Oliver Dutton <oliver@peptone.io>
Date: Mon, 16 Feb 2026 20:08:56 +0000
Subject: [PATCH 33/33] fix topk_logits dtype

---
 tallax/vllm/sampling.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tallax/vllm/sampling.py b/tallax/vllm/sampling.py
index 7943924..6260788 100644
--- a/tallax/vllm/sampling.py
+++ b/tallax/vllm/sampling.py
@@ -79,7 +79,7 @@ def bounded_topk_topp_and_sample(
     )
     .at[ind]
     .set(updates)
-  )(debug_vals["topk_idxs"], debug_vals["topk_logits"])
+  )(debug_vals["topk_idxs"], debug_vals["topk_logits"]).astype(logits.dtype)
   debug_vals["topk_topp_unnorm_probs_i32_unsorted"] = jax.vmap(
     lambda ind, updates: jnp.zeros_like(
       updates,