Fix import path and shape handling in topk sampling kernel

ANALYSIS_block_token_crash.py

@@ -0,0 +1,153 @@
+"""Analysis: Why block_token != 8 causes TPU runtime crashes
+
+ROOT CAUSE:
+===========
+
+In topk_mask.py, _find_boundary_chunk() creates a compile-time unrolled loop
+over batch_size, which is set to block_token when called from the kernel.
+
+PROBLEMATIC CODE (topk_mask.py, lines 89-92):
+=============================================
+
+```python
+batch_size = ref.shape[0]  # This is block_token in the kernel!
+
+# Creates batch_size separate pl.dslice operations
+boundary_slices = [
+    ref[:, pl.dslice(pl.multiple_of(ref_offset[i, 0], chunk_size), chunk_size)]
+    for i in range(batch_size)
+]
+
+# Merges them with batch_size-1 jnp.where operations
+boundary_slice = boundary_slices[0]
+for i in range(1, batch_size):
+    boundary_slice = jnp.where(iota0 == i, boundary_slices[i], boundary_slice)
+```
+
+EXECUTION FLOW:
+==============
+
+1. topk_topp_mask_and_sample_kernel receives logits_ref[block_token, vocab_size]
+2. Calls topk_mask_ref_inputs(logits_ref, ...) when stable=True
+3. topk_mask_ref_inputs calls find_boundary_idx(logits_ref, ...)
+4. find_boundary_idx calls _find_boundary_chunk(ref, ...) where ref.shape[0] = block_token
+5. _find_boundary_chunk sets batch_size = ref.shape[0] = block_token
+6. **Creates block_token slices and block_token-1 merge operations**
+
+WHY block_token=8 WORKS BUT OTHERS CRASH:
+=========================================
+
+1. **Compile-time loop unrolling explosion**
+   - With block_token=8: 8 pl.dslice ops + 7 jnp.where ops
+   - With block_token=16: 16 pl.dslice ops + 15 jnp.where ops
+   - With block_token=32: 32 pl.dslice ops + 31 jnp.where ops
+   - With block_token=64: 64 pl.dslice ops + 63 jnp.where ops
+
+   The compiler must analyze and optimize all these operations statically.
+   Larger values cause exponential growth in compilation complexity.
+
+2. **pl.dslice constraint violations**
+   - Each pl.dslice requires compile-time proof that:
+     * offset is a multiple of chunk_size (via pl.multiple_of)
+     * slice doesn't violate memory bounds
+   - With more slices, the compiler's symbolic analysis becomes harder
+   - May fail to prove constraints for non-power-of-2 or large block_token
+
+3. **TPU VMEM pressure**
+   - vmem_limit_bytes = 0.9 * 2^27 ≈ 120MB
+   - Each boundary_slice array: [block_token, chunk_size] in f32
+     = block_token * 256 * 4 bytes (for vocab_size=1024)
+   - With block_token=8: 8 * 256 * 4 = 8KB per slice
+   - With block_token=64: 64 * 256 * 4 = 64KB per slice
+   - Total for all slices + intermediates can exceed VMEM limit
+
+4. **TPU hardware alignment**
+   - TPU has natural block sizes (often powers of 2, especially 8)
+   - block_token=8 aligns with:
+     * 128-lane architecture (8 is factor of 128)
+     * VMEM tile sizes
+     * DMA transfer boundaries
+   - Other values may cause misalignment → undefined behavior
+
+5. **Nested iterations amplify the problem**
+   - find_boundary_idx calls _find_boundary_chunk **twice**:
+     * First with chunk_size = sqrt(vocab/128) * 128 = 256 (for vocab=1024)
+     * Second with chunk_size = 128
+   - Each call creates block_token slices
+   - **Total: 2 * block_token slices + 2 * (block_token-1) merges**
+
+SPECIFIC CRASH SCENARIOS:
+========================
+
+block_token=16:
+- 32 pl.dslice operations
+- 30 jnp.where merges
+- Compilation time increases significantly
+- May exceed compiler resource limits
+
+block_token=32:
+- 64 pl.dslice operations
+- 62 jnp.where merges
+- High probability of VMEM overflow
+- Symbolic analysis likely fails
+
+block_token=64:
+- 128 pl.dslice operations
+- 126 jnp.where merges
+- Almost certain to crash during compilation or execution
+- Violates fundamental TPU constraints
+
+WHY CPU THREADS CRASH:
+=====================
+
+When TPU compiler fails:
+1. May corrupt internal state
+2. Can trigger segfaults in compiler backend
+3. Errors propagate to JAX Python runtime
+4. Crashes Python interpreter threads
+
+SOLUTIONS:
+==========
+
+Option 1: Restrict block_token to 8 (current approach)
+- Safe, tested, works
+- Limits flexibility
+
+Option 2: Rewrite _find_boundary_chunk to avoid per-batch slicing
+- Use vectorized operations instead of loop
+- More complex but supports arbitrary block_token
+- Would require significant refactoring
+
+Option 3: Use different algorithm for large block_token
+- Keep current impl for block_token <= 8
+- Fall back to simpler (slower) algorithm for larger values
+
+Option 4: Make block_token a compile-time constant
+- Set via static_argnames
+- Allow compiler to specialize per block_token
+- May help with larger values but still risky
+
+RECOMMENDATION:
+==============
+
+Keep block_token=8 as the default and only supported value.
+Add assertion to enforce this:
+
+```python
+assert block_token == 8, (
+    "Only block_token=8 is currently supported due to TPU compiler "
+    "constraints. See topk_mask.py _find_boundary_chunk for details."
+)
+```
+
+The performance impact of block_token=8 is minimal since batch sizes
+are typically padded to multiples of 8 anyway.
+"""
+
+import sys
+
+def main():
+    print(__doc__)
+
+if __name__ == "__main__":
+    main()

tallax/vllm/high_precision_uint.py

@@ -51,10 +51,10 @@ def to_f32(self) -> jax.Array:
       f32 array with same shape as each part
     """
     result = self.parts[0].astype(jnp.float32)
-    scale = 2.0 ** 16
+    scale = jnp.float32(2**16)
     for part in self.parts[1:]:
       result += part.astype(jnp.float32) * scale
-      scale *= 2.0 ** 16
+      scale *= jnp.float32(2**16)
     return result
 
   @classmethod
@@ -71,9 +71,9 @@ def from_f32(cls, x: jax.Array, num_parts: int) -> 'HighPrecisionUInt':
     parts = []
     remainder = x
     for _ in range(num_parts):
-      part_f32 = jnp.fmod(remainder, 2.0 ** 16)
+      part_f32 = jnp.fmod(remainder, jnp.float32(2**16))
       parts.append(part_f32.astype(jnp.int32))
-      remainder = jnp.floor(remainder / (2.0 ** 16))
+      remainder = jnp.floor(remainder / jnp.float32(2**16))
     return cls(parts)
 
   def harmonize(self) -> 'HighPrecisionUInt':
@@ -197,20 +197,17 @@ def compare_ge(self, other: 'HighPrecisionUInt') -> jax.Array:
     other_parts = other.parts + [0] * (max_len - len(other.parts))
 
     # Compare from MSB to LSB
-    result = jnp.ones_like(self.parts[0], dtype=bool)
+    result = jnp.zeros_like(self.parts[0], dtype=bool)
     still_equal = jnp.ones_like(self.parts[0], dtype=bool)
 
     for i in range(max_len - 1, -1, -1):
       part_greater = self_parts[i] > other_parts[i]
-      part_less = self_parts[i] < other_parts[i]
       part_equal = self_parts[i] == other_parts[i]
 
       # Update result: if still equal and current part is greater, set true
       # if still equal and current part is less, set false
-      result = jnp.where(still_equal & part_greater, True, result)
-      result = jnp.where(still_equal & part_less, False, result)
-
+      result |= (still_equal & part_greater)
       # Update still_equal
       still_equal = still_equal & part_equal
 
-    return result
+    return result | still_equal

tallax/vllm/topk_mask.py

@@ -20,7 +20,7 @@
 from jax.experimental import pallas as pl
 from jax.experimental.pallas import tpu as pltpu
 
-from tallax.tax.binary_search import binary_search
+from tallax.vllm.binary_search import binary_search
 from tallax.tax.utils import NUM_LANES, get_dtype_info
 
 
@@ -39,26 +39,27 @@ def _find_boundary_chunk(
   the k'th matching element.
 
   Args:
-    logits_ref: Logits reference of shape [batch, vocab_size]
+    ref: Reference array of shape [batch, vocab_size]
     target: Target value to match (shape [batch, 1])
     k: Target count (shape [batch, 1])
     chunk_size: Size of each chunk
+    active_chunk: Optional subset of ref to search in
+    ref_offset: Offset into ref for indexing
 
   Returns:
-    Tuple of (chunk_index, cumsum_before_chunk) where:
-      - chunk_index: Index of chunk containing k'th match (shape [batch, 1])
-      - cumsum_before_chunk: Cumulative matches before this chunk (shape [batch, 1])
+    Tuple of (ref_offset, boundary_slice, k) where:
+      - ref_offset: Updated offset to boundary chunk start
+      - boundary_slice: Extracted boundary chunk [batch, chunk_size]
+      - k: Updated k count after subtracting earlier chunks
   """
-  active_chunk = ref if active_chunk is None else active_chunk
+  arr = ref if active_chunk is None else active_chunk
   # Calculate number of chunks using ceiling division
-  num_chunks = pl.cdiv(active_chunk.shape[1], chunk_size)
-  # If chunk size doesn't evenly divide array, we make the first chunk smaller. This choice avoids OOB indexing to extract chunk_size length slices later.
-  chunk_offsets = [0] + sorted(filter(lambda x: x>0, [active_chunk.shape[1] - i * chunk_size for i in range(num_chunks)]))
-  # Split into chunks (no padding needed, variable sizes OK)
-  chunks = [active_chunk[:, lb:ub] for lb, ub in zip(chunk_offsets[:-1], chunk_offsets[1:], strict=True)]
-  assert sum([c.shape[1] for c in chunks]) == active_chunk.shape[1]
-  first_chunk_size = chunks[0].shape[1]
-  assert len(set([c.shape[1] for c in chunks[1:]])) <= 1
+  num_chunks = pl.cdiv(arr.shape[1], chunk_size)
+  # Split into chunks with multiples of chunk_size (may be OOB for last chunk)
+  chunks = [
+    arr[:, i * chunk_size : (i + 1) * chunk_size]
+    for i in range(num_chunks)
+  ]
 
   # Count matches in each chunk, keeping (batch, 1) shape
   num_matches = [
@@ -70,20 +71,33 @@ def _find_boundary_chunk(
   cumsums = [num_matches[0]]
   for i in range(1, len(num_matches)):
     cumsums.append(cumsums[i - 1] + num_matches[i])
+
   boundary_idx = sum((c < k) for c in cumsums)
+  # Subtract counts from all chunks BEFORE the boundary chunk
   k -= sum((i == (boundary_idx - 1)) * c for i, c in enumerate(cumsums))
 
   # We'll do batch_size separate dslices into arr
   batch_size = ref.shape[0]
   iota0, iota1 = (jax.lax.broadcasted_iota(jnp.int32, (batch_size, chunk_size), dim) for dim in (0, 1))
-
-  # first chunk may not be chunk_size in length if arr doesn't evenly divide
-  ref_offset += jnp.maximum(first_chunk_size + (boundary_idx - 1) * chunk_size, 0)
-  # We indead into ref instead of ref_slice as dynamic_slice not supported on arrays
-  boundary_slices = [ref[:, pl.dslice(ref_offset[i, 0], chunk_size)] for i in range(batch_size)]
+
+  # Update offset by multiples of chunk_size
+  ref_offset += boundary_idx * chunk_size
+  # Assure compiler offset is a multiple of chunk_size
+  # This is us guaranteeing when using multiple iterations of find_boundary_chunk that current chunk_size evenly divides all previous chunk_sizes
+  # Index into ref (not ref_slice) as dynamic_slice not supported on arrays
+  # These dslices may be OOB, which is fine - we mask them out later
+  boundary_slices = [ref[:, pl.dslice(pl.multiple_of(ref_offset[i, 0], chunk_size), chunk_size)] for i in range(batch_size)]
   boundary_slice = boundary_slices[0]
   for i in range(1, batch_size):
     boundary_slice = jnp.where(iota0 == i, boundary_slices[i], boundary_slice)
+
+  # Mask OOB indices to dtype min to ensure they don't interfere with comparisons
+  if num_chunks * chunk_size != arr.shape[1]:
+    boundary_slice = jnp.where(
+      (ref_offset + iota1) < ref.shape[1],
+      boundary_slice,
+      get_dtype_info(ref).min
+    )
   return ref_offset, boundary_slice, k
 
 def find_boundary_idx(ref, k, threshold):
@@ -117,12 +131,19 @@ def find_boundary_idx(ref, k, threshold):
     cumsums.append(cumsums[i - 1] + num_matches[i])
   return (ref_offset + sum((c < k) for c in cumsums))  
 
+def alu_gt(lhs, rhs):
+   # equiv to (logits > pivot).astype(jnp.int32), but avoids masks
+   # Only valid if no NaN and no inf/-inf in values.
+  assert lhs.dtype == jnp.float32
+  return ((rhs - lhs).view(jnp.int32) >> 31)
+
 def topk_mask_ref_inputs(
   logits_ref,
   k_ref,
   *,
   replace_val: float,
   stable: bool,
+  use_alu: bool,
 ):
   """Pallas kernel for topk masking with parallel chunk-based reduction.
 
@@ -139,7 +160,10 @@ def topk_mask_ref_inputs(
   logits = logits_ref[...]
   k = k_ref[...]
   # next value after the largest value where less than k gt it.
-  predicate_fn = lambda pivot: (logits > pivot).sum(-1, keepdims=True) < k
+  if use_alu:
+    predicate_fn = lambda pivot: alu_gt(logits, pivot).sum(-1, keepdims=True) < k # Use ALU as faster
+  else:
+    predicate_fn = lambda pivot: (logits > pivot).sum(-1, keepdims=True) < k
   bound_shape = (logits.shape[0], 1)
   _, threshold = binary_search(
     predicate_fn,
@@ -148,12 +172,11 @@ def topk_mask_ref_inputs(
 
   if not stable:
     # Simple threshold masking
-    output_ref[...] = jnp.where(
+    return jnp.where(
       logits >= threshold,
       logits,
       replace_val
     )
-    return
 
   # Step 2: Find exact boundary for stable masking
   boundary_idx = find_boundary_idx(
@@ -174,20 +197,22 @@ def topk_mask_pallas_kernel(
   *,
   replace_val: float,
   stable: bool,
+  use_alu: bool,
 ):
-  output_ref[...] = topk_mask_ref_inputs(logits_ref, k_ref, replace_val=replace_val, stable=stable)
+  output_ref[...] = topk_mask_ref_inputs(logits_ref, k_ref, replace_val=replace_val, stable=stable, use_alu=use_alu)
 
 
 @functools.partial(
   jax.jit,
-  static_argnames=["replace_val", "stable", "interpret"]
+  static_argnames=["replace_val", "stable", "interpret", "use_alu"]
 )
 def topk_mask_pallas(
   x: jax.Array,
   k: int,
   replace_val: float = -1e12,
   stable: bool = True,
   interpret: bool = False,
+  use_alu: bool = False
 ) -> jax.Array:
   """Pallas-based topk mask with parallel chunk-based reduction.
 
@@ -210,7 +235,9 @@ def topk_mask_pallas(
       topk_mask_pallas_kernel,
       replace_val=replace_val,
       stable=stable,
+      use_alu=use_alu,
     ),
+    compiler_params=pltpu.CompilerParams(vmem_limit_bytes=int(0.9 * 2**27)),
     out_shape=output_shape,
     interpret=interpret,
   )(x, k)