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/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/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/tax/divide_and_filter_topk/topk.py b/tallax/tax/divide_and_filter_topk/topk.py
index 0e481ff..f377cf5 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):
@@ -80,7 +80,8 @@ def to_comparison_dtype(x):
 
 
 def bitonic_topk_arrays(operands, k, num_keys, val_dtype=None, max_index=None):
-  """Top-k of arrays, packing bf16 and indices into single 32-bit dtype if possible.
+  """Top-k of arrays, first operand of value, second of indices, possibly with further operands.
+  Packing bf16 and indices into single 32-bit dtype if possible.
 
   Args:
     operands: List of arrays to find top-k from.
@@ -91,19 +92,36 @@ def bitonic_topk_arrays(operands, k, num_keys, val_dtype=None, max_index=None):
   Returns:
     List of top-k arrays.
   """
-  if val_dtype is None or max_index is None:
-    return _bitonic_topk_arrays(operands, k=k, num_keys=num_keys)
-  assert len(operands) == 2 and type(operands) == list
+  assert type(operands) == list
+  stable = num_keys > 1
+  if not jnp.issubdtype(operands[1].dtype, jnp.integer):
+    raise ValueError("Expected second operand to be indices, so integer dtype")
   dtypes = [x.dtype for x in operands]
-  pack = val_dtype == jnp.bfloat16 and max_index <= 2**16
+  pack = (
+    val_dtype is not None
+    and max_index is not None
+    and (val_dtype == jnp.bfloat16 and max_index <= 2**16)
+  )
+  if num_keys == 2 and not pack:
+    # sort vals descending, indices ascending so we negate indices
+    operands[1] = -operands[1]
   if pack:
-    operands = [pack_bf16_u16_to_i32(*operands, stable=(num_keys > 1))]
-  operands = _bitonic_topk_arrays(operands, k=k, num_keys=min(len(operands), num_keys))
+    operands[1] = 2**16 - 1 - operands[1]
+    operands = [
+      pack_bf16_u16_to_i32(*operands[:2], stable=stable),
+      *operands[2:],
+    ]
+    num_keys = max(1, num_keys - 1)
+  operands = _bitonic_topk_arrays(operands, k=k, num_keys=num_keys)
   if pack:
-    assert len(operands) == 1
-    operands = list(unpack_bf16_u16_from_i32(operands[0], stable=False))
-    assert len(operands) == 2
-  # convert back dtypes (incase theyve changed)
+    operands = (
+      list(unpack_bf16_u16_from_i32(operands[0], stable=stable)) + operands[1:]
+    )
+    operands[1] = 2**16 - 1 - operands[1]
+  if num_keys == 2 and not pack:
+    # invert back so sort vals descending, indices ascending
+    operands[1] = -operands[1]
+  # convert back dtypes (incase they've changed)
   return [x.astype(dtype) for x, dtype in zip(operands, dtypes, strict=True)]
 
 
@@ -178,16 +196,18 @@ def compute_idxs(i):
     )
 
   num_full_slices = vocab_size // num_bins
+
   def loop_body(i, bins_topk_outs):
     # Process in reverse order, due to >= swap
     # this leads to stable index ordering
-    i = num_full_bins - 1 - i
+    i = num_full_slices - 1 - i
     vals = logits[..., pl.dslice(num_bins * i, num_bins)]
     idxs = compute_idxs(i)
     return update_bins_topk(vals, idxs, *bins_topk_outs)
 
   # Handle remaining elements if vocab_size doesn't divide num_bins
   # We do this first as we loop the array backwards, due to >= swap this leads to stable index ordering and cases where the top-k includes -inf
+  bins_topk_outs = (bins_topk_vals, bins_topk_idxs)
   rem_vals = _extract_remainder_slice(logits, num_bins)
   if rem_vals is not None:
     # Create idxs for the final segment
@@ -198,7 +218,7 @@ def loop_body(i, bins_topk_outs):
   bins_topk_outs = unrolled_fori_loop(
     num_full_slices,
     loop_body,
-    (bins_topk_vals, bins_topk_idxs),
+    bins_topk_outs,
     unroll=unroll,
   )
 
@@ -225,15 +245,22 @@ def _merge_unconverged_bins_topk(
   num_packed_bins = pl.cdiv(max_k, m) - 1
   if num_packed_bins > NUM_LANES or num_packed_bins > num_bins:
     raise NotImplementedError
-  bin_vals = bins_topm_vals_ref[:, pl.dslice((m - 1) * num_bins, num_bins)]
+  bin_vals, bin_idxs = (
+    ref[:, pl.dslice((m - 1) * num_bins, num_bins)]
+    for ref in (bins_topm_vals_ref, bins_topm_idxs_ref)
+  )
   # Use bitonic_topk_arrays descending to get bin indices ordered by contribution count
   bin_indices = jax.lax.broadcasted_iota(jnp.int32, (block_token, num_bins), 1)
-  # Sort descending by num_gt_k to get top NUM_LANES bin indices
-  _, sorted_bin_indices = bitonic_topk_arrays(
-    [bin_vals, bin_indices],
+  bin_operands = [bin_vals, bin_indices]
+  if stable:
+    # stable sort comparison using indices
+    bin_operands.insert(1, bin_idxs)
+  # Sort descending by m'th largest value of each bin to get bins which may contribute to top-k
+  sorted_bin_indices = bitonic_topk_arrays(
+    bin_operands,
     k=num_packed_bins,
-    num_keys=1+int(stable),
-  )
+    num_keys=1 + int(stable),
+  )[-1]
   sorted_bin_indices = pad(sorted_bin_indices, (NUM_SUBLANES, NUM_LANES))
 
   # Repeat first num_packed_bins values across NUM_LANES positions to create packing permutation
@@ -338,7 +365,7 @@ def _merge_unconverged_bins_topk(
       k=max_k,
       val_dtype=logits_ref.dtype,
       max_index=logits_ref.shape[1],
-      num_keys=1+int(stable),
+      num_keys=1 + int(stable),
     )
   )
 
@@ -379,9 +406,9 @@ def dynamic_topk_refs(
   block_token = logits_ref.shape[0]
   shape = (block_token, bins_topm_vals_ref.shape[1])
   block_topk = bins_topm_vals_ref.shape[0]
-  assert block_topk % block_token == 0, (
-    "block_topk must be a multiple of block_token"
-  )
+  assert (
+    block_topk % block_token == 0
+  ), "block_topk must be a multiple of block_token"
 
   pid = pl.program_id(0)
 
@@ -440,11 +467,11 @@ def _():
       # the largest m-th largest value, then top-k is guaranteed to be in bins
       # top-(m-1) collated
       pivot = bins_topm_vals[m - 1].max(-1, keepdims=True)
-      # if stable sort, we must include values at the boundary in to the bitonic stable sort 
+      # if stable sort, we must include values at the boundary in to the bitonic stable sort
       # if not we just need k values greater than or equal to
       _comp = operator.gt if stable else operator.ge
       num_larger = (
-        sum(_comp(v, pivot) for v in bins_topm_vals[: m - 1])
+        sum((_comp(v, pivot) for v in bins_topm_vals[: m - 1]))
         .astype(jnp.float32)
         .sum(-1)
       )
@@ -542,7 +569,7 @@ def _():
           k=max_k,
           val_dtype=logits_ref.dtype,
           max_index=logits_ref.shape[1],
-          num_keys=1+int(stable),
+          num_keys=1 + int(stable),
         )
         topk_vals_ref[...], topk_idxs_ref[...] = (
           vals.astype(topk_vals_ref.dtype),
@@ -567,6 +594,7 @@ def _():
     "guarantee_convergence",
     "replace_val",
     "interpret",
+    "stable",
   ),
 )
 def _top_bounded_k(
@@ -581,6 +609,7 @@ def _top_bounded_k(
   guarantee_convergence: bool = True,
   replace_val: float | int | None = None,
   interpret: bool = False,
+  stable: bool = True,
 ):
   """
   High-level interface for adaptive binned top-k computation on TPU.
@@ -723,6 +752,7 @@ def _top_bounded_k(
       bins_topm_schedule=bins_topm_schedule,
       guarantee_convergence=guarantee_convergence,
       replace_val=replace_val,
+      stable=stable,
     ),
     in_specs=(
       pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0)),
@@ -767,6 +797,7 @@ def _top_bounded_k(
     "guarantee_convergence",
     "replace_val",
     "interpret",
+    "stable",
   ),
 )
 @functools.wraps(_top_bounded_k)
@@ -782,6 +813,7 @@ def top_bounded_k(
   guarantee_convergence: bool = False,
   replace_val: float | int | None = None,
   interpret: bool = False,
+  stable: bool = True,
 ):
   def _closed_topk(logits: jax.Array, k: jax.Array):
     return _top_bounded_k(
@@ -796,6 +828,7 @@ def _closed_topk(logits: jax.Array, k: jax.Array):
       guarantee_convergence=guarantee_convergence,
       replace_val=replace_val,
       interpret=interpret,
+      stable=stable,
     )
 
   @custom_partitioning
diff --git a/tallax/vllm/__init__.py b/tallax/vllm/__init__.py
index 8f15360..78a5e17 100644
--- a/tallax/vllm/__init__.py
+++ b/tallax/vllm/__init__.py
@@ -8,14 +8,12 @@
   - reference: Pure JAX reference implementation (no Pallas).
 """
 
-from tallax.vllm.sampling import topk_topp_and_sample
-from tallax.vllm.bounded_k import top_p_and_sample
-from tallax.vllm.arbitrary_k import topk_topp_mask_and_sample
-from tallax.vllm.reference import reference_topk_topp_mask_and_sample
+from tallax.vllm.sampling import bounded_topk_topp_and_sample
+from tallax.vllm.arbitrary_k import arbitrary_topk_topp_and_sample
+from tallax.vllm.reference import reference_topk_topp_and_sample
 
 __all__ = [
-  "topk_topp_and_sample",
-  "top_p_and_sample",
-  "topk_topp_mask_and_sample",
-  "reference_topk_topp_mask_and_sample",
+  "bounded_topk_topp_and_sample",
+  "arbitrary_topk_topp_and_sample",
+  "reference_topk_topp_and_sample",
 ]
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 236c81f..bdf3916 100644
--- a/tallax/vllm/arbitrary_k/kernel.py
+++ b/tallax/vllm/arbitrary_k/kernel.py
@@ -25,48 +25,73 @@
   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
+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.
 
-def sample_probs(unnorm_probs_i32, random_u128_in_u32s, max_val=2**24 - 1):
+  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):
   """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,
     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 topk_topp_mask_and_sample_kernel(
+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,
   sampled_tokens_ref,
-  debug_results_ref=None,
+  debug_arrays_ref=None,
   *,
   stable: bool,
-  replace_val: float,
   underlying_logits_dtype=None,
 ):
   """Pallas kernel body for combined topk + topp + sample.
@@ -78,24 +103,23 @@ 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
   """
   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,
+        random_u128_in_u32s_refs,
         k_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,
       )
 
@@ -112,6 +136,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
@@ -119,12 +144,12 @@ def scoped_body(scoped_ref):
   topk_logits = topk_mask(
     logits_ref,
     k_ref,
-    replace_val=replace_val,
     stable=stable,
     underlying_dtype=underlying_logits_dtype,
   )
 
   # Stage 3: Temperature scaling
+  topk_logits_pre_temperature = topk_logits
   topk_logits /= temperature
   logits_max /= temperature
 
@@ -136,64 +161,42 @@ def scoped_body(scoped_ref):
   )
 
   # Stage 5: Sample
-  next_tokens = sample_probs(
-    unnorm_probs_i32, [ref[...] for ref in random_u128_in_u32s_ref]
+  next_tokens, random_unnorm_cdf_sampled = sample_probs(
+    unnorm_probs_i32,
+    [ref[...] for ref in random_u128_in_u32s_refs],
   )
-
   sampled_tokens_ref[...] = jnp.where(
-    temperature < _SAMPLING_EPS, greedy_sampled, next_tokens
+    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,
-    )
-
-
-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]
+  # Debug: write full intermediate arrays if debug refs are provided
+  if debug_arrays_ref is not None:
+    debug_arrays_ref["greedy_sampled"][...] = greedy_sampled
+    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
     )
-  if "next_tokens" in debug_refs:
-    debug_refs["next_tokens"][...] = next_tokens[:1, :1].astype(jnp.int32)
+    # Store as tuple of (high_u32, low_u32)
+    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
 
 
 @functools.partial(
   jax.jit,
   static_argnames=[
     "stable",
-    "replace_val",
     "block_token",
     "interpret",
     "debug",
   ],
 )
-def topk_topp_mask_and_sample(
+def arbitrary_topk_topp_and_sample(
   logits: jax.Array,
   rng_key: jax.Array,
   k: jax.Array,
@@ -201,7 +204,6 @@ def topk_topp_mask_and_sample(
   temperature: jax.Array,
   *,
   stable: bool = True,
-  replace_val: float = -1e12,
   block_token: int = 8,
   interpret: bool = False,
   debug: bool = False,
@@ -218,7 +220,6 @@ def topk_topp_mask_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
@@ -228,103 +229,86 @@ def topk_topp_mask_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))
   )
 
-  # 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:
-    result = pl.pallas_call(
-      functools.partial(
-        topk_topp_mask_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), logits.dtype
       ),
-      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 via SMEM
-    out_shapes = (output_shape, debug_shapes)
-    result, debug_results = pl.pallas_call(
-      functools.partial(
-        topk_topp_mask_and_sample_kernel,
-        stable=stable,
-        replace_val=replace_val,
+      "topk_topp_unnorm_probs_i32_unsorted": jax.ShapeDtypeStruct(
+        (padded_batch, vocab_size), jnp.int32
       ),
-      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": (
+        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)
+      ),
+      "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)),
-        jax.tree.map(
-          lambda _: pl.BlockSpec(memory_space=pltpu.SMEM),
-          debug_shapes,
-        ),
+      )
+      * 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,
+    ),
+    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,
       ),
-      compiler_params=pltpu.CompilerParams(vmem_limit_bytes=int(0.9 * 2**27)),
-      interpret=interpret,
-    )(logits, random_u128_in_u32s, k, p, temperature)
+      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)
 
-    return result[:batch_size, 0], debug_results
+  if debug:
+    return sampled_tokens, debug_aux
+  return sampled_tokens
diff --git a/tallax/vllm/arbitrary_k/topk_mask.py b/tallax/vllm/arbitrary_k/topk_mask.py
index f7c5747..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,
@@ -26,6 +27,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 +42,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 +88,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 +103,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 +113,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 +124,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,
   )
@@ -143,7 +149,6 @@ def topk_mask(
   logits_ref,
   k_ref,
   *,
-  replace_val: float,
   stable: bool,
   underlying_dtype=None,
 ):
@@ -152,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)
 
@@ -200,6 +204,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,
@@ -215,54 +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)
-
-
-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)
+  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/__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/bounded_k/top_p_and_sample.py b/tallax/vllm/bounded_k/top_p_and_sample.py
index e3abd89..6a30208 100644
--- a/tallax/vllm/bounded_k/top_p_and_sample.py
+++ b/tallax/vllm/bounded_k/top_p_and_sample.py
@@ -20,13 +20,13 @@
 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.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:
@@ -58,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(
@@ -78,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")
@@ -103,16 +102,17 @@ 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(
   *,
   topk_logits,
-  topk_idx,
+  topk_idxs,
   random_u128_in_u32s,
   top_p,
   temperature,
+  debug=False,
 ):
   """Fused top-p filtering + sampling on pre-sorted top-k logits.
 
@@ -122,131 +122,175 @@ 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)
 
   # 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_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_idxs_transposed[:1, :]
 
-  topk_logits = topk_logits / temperature[None, :].astype(topk_logits.dtype)
+  # Temperature scaling
+  topk_logits_scaled = topk_logits_transposed / temperature[None, :].astype(
+    topk_logits_transposed.dtype
+  )
 
-  unnorm_probs_i32 = top_p_integer_mask(
-    topk_logits=topk_logits, p=top_p, axis=0
+  # 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_idxs_transposed, unnorm_probs_i32_sorted],
+    k=topk_idxs_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)
-  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, keepdims=True)
+  result = jnp.where(
+    temperature[None, :] < SAMPLING_EPS, greedy_sampled, next_tokens
+  )
+
+  if not debug:
+    return result
+
+  # Build debug output matching the reference format but for the k-slice
+  random_unnorm_cdf_sampled_low = target_cumsum.T
+
+  debug_results = {
+    "greedy_sampled": greedy_sampled.T,
+    "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,
+    ),
+    "next_tokens": next_tokens.T,
+  }
 
-  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)
+  return result, debug_results
 
 
 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)),
+    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[...],
+    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,
-  rng_key: jax.Array,
-  top_p: jax.Array,
-  temperature: jax.Array,
-  *,
-  vocab_size: int,
-  replace_val: float,
-  interpret: bool = False,
-  dim0_offset: int = 0,
-) -> 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 = (
+    jnp.broadcast_to(v, (batch_size,)) for v in (top_p, temperature)
+  )
 
-  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,)
-    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
+  output_shape = jax.ShapeDtypeStruct((1, batch_size), jnp.int32)
+
+  if debug:
+    debug_out_shapes = {
+      "greedy_sampled": jax.ShapeDtypeStruct((batch_size, 1), jnp.int32),
+      "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": (
+        jax.ShapeDtypeStruct((batch_size, 1), jnp.uint32),
+      )
+      * 2,
+      "next_tokens": jax.ShapeDtypeStruct((batch_size, 1), jnp.int32),
+    }
+  else:
+    debug_out_shapes = None
+
+  sampled_tokens, debug_aux = 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,
+  )
 
-  Returns:
-    Sampled tokens of shape (batch_size,)
-  """
-  batch_size = topk_logits.shape[0]
+  sampled_tokens = sampled_tokens.squeeze(0)
+  if not debug:
+    return sampled_tokens
 
-  # 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(
-    topk_logits=topk_logits,
-    topk_idx=topk_idx,
-    random_u128_in_u32s=random_u128_in_u32s,
-    top_p=top_p,
-    temperature=temperature,
+  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(
   jit,
-  static_argnames=(
-    "vocab_size",
-    "replace_val",
-    "interpret",
-  ),
+  static_argnames=("debug",),
 )
-def top_p_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,
-  replace_val: float,
-  interpret: bool = False,
+  debug: bool = False,
 ) -> jax.Array:
   """Sharded wrapper for top-p sampling with custom partitioning.
 
@@ -258,64 +302,82 @@ def top_p_and_sample(
     rng_key: RNG key for sampling.
     top_p: Top-p threshold values.
     temperature: Temperature values.
-    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.
   """
+  # Generate random u128 outside the sharded function
+  batch_size = topk_logits.shape[0]
+  random_u128_in_u32s = list(
+    sample_random_u128_in_u32s(rng_key, (batch_size, 1))
+  )
+
+  out_shapes = jax.eval_shape(
+    functools.partial(
+      _top_p_and_sample,
+      debug=debug,
+    ),
+    topk_logits,
+    topk_idx,
+    random_u128_in_u32s,
+    top_p,
+    temperature,
+  )
 
   @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,
-      top_p,
-      temperature,
-      vocab_size=vocab_size,
-      replace_val=replace_val,
-      interpret=interpret,
+    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, rng_key, 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,
-        top_p,
-        temperature,
-        vocab_size=vocab_size,
-        replace_val=replace_val,
-        interpret=interpret,
-        dim0_offset=dim0_offset,
+    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, b, b -> b",
+    sharding_rule=sharding_rule,
     need_replication_factors=("k", "r"),
   )
-
-  return sharded_top_p_and_sample(
-    topk_logits, topk_idx, rng_key, top_p, temperature
+  flat_outs = sharded_top_p_and_sample(
+    topk_logits, topk_idx, random_u128_in_u32s, top_p, temperature
   )
+  return jax.tree.unflatten(jax.tree.structure(out_shapes), flat_outs)
diff --git a/tallax/vllm/reference.py b/tallax/vllm/reference.py
index 4120266..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):
@@ -39,11 +37,12 @@ def u128_modulo_u64_pure_callback(r_parts, m_parts):
     (jax.ShapeDtypeStruct(r_parts[0].shape, jnp.uint32),) * 2,
     *r_parts,
     *m_parts,
+    vmap_method="legacy_vectorized",
   )
   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,
@@ -51,7 +50,6 @@ def reference_topk_topp_mask_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.
@@ -65,16 +63,16 @@ def reference_topk_topp_mask_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:
     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
+    scale = 2**SCALE_BITS - 1
 
     logits = logits.astype(jnp.float32)
 
@@ -85,7 +83,7 @@ def reference_topk_topp_mask_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
@@ -127,17 +125,25 @@ def reference_topk_topp_mask_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
 
+    # 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 = {
       "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,
+      "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),
+        sampled_total_low.squeeze(1),
+      ),
       "next_tokens": next_tokens,
     }
     return result, debug_results
diff --git a/tallax/vllm/sampling.py b/tallax/vllm/sampling.py
index 9809f5c..6260788 100644
--- a/tallax/vllm/sampling.py
+++ b/tallax/vllm/sampling.py
@@ -8,20 +8,33 @@
 
 import functools
 import jax
-from tallax.vllm.bounded_k import top_p_and_sample
+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",
+    "stable",
+  ),
 )
-def topk_topp_and_sample(
-  rng_key,
+def bounded_topk_topp_and_sample(
   logits,
-  tpu_sampling_metadata,
+  rng_key,
+  top_k,
+  top_p,
+  temperature,
   max_k: int,
   num_bins: int | None = None,
+  stable: bool = True,
   bins_topm_schedule: int | None = None,
+  debug: bool = False,
 ):
   """Combined top-k, top-p filtering, and sampling for vLLM inference.
 
@@ -38,22 +51,47 @@ 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,
-    replace_val=-1e12,
+    k=top_k,
+    replace_val=REPLACE_VAL,
     max_k=max_k,
     num_bins=num_bins,
     bins_topm_schedule=bins_topm_schedule,
     guarantee_convergence=True,
+    stable=stable,
+  )
+  outs = topp_and_sample(
+    topk_logits=topk_logits,
+    topk_idx=topk_idxs,
+    rng_key=rng_key,
+    top_p=top_p,
+    temperature=temperature,
+    debug=debug,
   )
-  return top_p_and_sample(
-    topk_logits,
-    topk_idxs,
-    rng_key,
-    top_p=tpu_sampling_metadata.top_p,
-    temperature=tpu_sampling_metadata.temperature,
-    vocab_size=vocab_size,
-    replace_val=-1e12,
+  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"]).astype(logits.dtype)
+  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
diff --git a/tests/vllm/arbitrary_k/kernel_test.py b/tests/vllm/arbitrary_k/kernel_test.py
index f59f166..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
   )
 
@@ -43,18 +43,35 @@ 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 (tuple of high and low u32)
+  np.testing.assert_array_equal(
+    kernel_debug["random_unnorm_cdf_sampled"][0][0],
+    ref_debug["random_unnorm_cdf_sampled"][0][0],
+  )
   np.testing.assert_array_equal(
-    kernel_debug["topp_nonzero_count"][0, 0],
-    ref_debug["topp_nonzero_count"][0],
+    kernel_debug["random_unnorm_cdf_sampled"][1][0],
+    ref_debug["random_unnorm_cdf_sampled"][1][0],
   )
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)