diff --git a/ANALYSIS_block_token_crash.py b/ANALYSIS_block_token_crash.py new file mode 100644 index 00000000..67af0310 --- /dev/null +++ b/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() diff --git a/tallax/vllm/high_precision_uint.py b/tallax/vllm/high_precision_uint.py index 45fd46ea..836b700a 100644 --- a/tallax/vllm/high_precision_uint.py +++ b/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 diff --git a/tallax/vllm/topk_mask.py b/tallax/vllm/topk_mask.py index e499ce72..ac0a7f3d 100644 --- a/tallax/vllm/topk_mask.py +++ b/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,13 +197,14 @@ 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, @@ -188,6 +212,7 @@ def topk_mask_pallas( 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) diff --git a/tallax/vllm/topk_topp_mask_and_sample.py b/tallax/vllm/topk_topp_mask_and_sample.py index 7032b94e..86bae48f 100644 --- a/tallax/vllm/topk_topp_mask_and_sample.py +++ b/tallax/vllm/topk_topp_mask_and_sample.py @@ -5,6 +5,7 @@ import jax.numpy as jnp from jax import lax from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu from tallax.tax.bitonic.topk import max_arrays from tallax.vllm.topk_mask import topk_mask_ref_inputs @@ -44,6 +45,8 @@ def topk_topp_mask_and_sample_kernel( greedy_sampled = max_arrays( [logits_ref[...], token_idx], num_keys=1+int(stable), axis=1 )[1] + # Reshape to (block_token, 1) to match output ref + greedy_sampled = jnp.expand_dims(greedy_sampled, axis=-1) logits = topk_mask_ref_inputs(logits_ref, k_ref, replace_val=replace_val, stable=stable) logits = topp_mask( @@ -53,7 +56,7 @@ def topk_topp_mask_and_sample_kernel( # Random key splitting is based on idx in ravelled array # We pass in (batch_idx, token_idx) for linearized position: batch_idx * vocab_size + token_idx - batch_idx = lax.broadcasted_iota(jnp.int32, logits.shape, 0) + dim0_offset_ref[0] + batch_idx = lax.broadcasted_iota(jnp.int32, logits.shape, 0) + pl.program_id(0) * logits_ref.shape[0] + dim0_offset_ref[0] next_tokens = sparse_random_categorical( rng_key_ref, logits, @@ -62,6 +65,8 @@ def topk_topp_mask_and_sample_kernel( axis=1, # Sample along vocab axis dtype=jnp.float32, )[1] # Take sampled token indices + # Reshape to (block_token, 1) to match output ref + next_tokens = jnp.expand_dims(next_tokens, axis=-1) sampled_tokens_ref[...] = jnp.where(temperature_ref[...] < _SAMPLING_EPS, greedy_sampled, next_tokens) @@ -103,14 +108,33 @@ def topk_topp_mask_and_sample( batch_size, vocab_size = logits.shape # Ensure inputs have correct shapes - k = jnp.broadcast_to(k, (batch_size, 1)) - p = jnp.broadcast_to(p, (batch_size, 1)) - temperature = jnp.broadcast_to(temperature, (batch_size, 1)) + # First ensure they're at least 1D arrays, then reshape to (batch_size, 1) + k = jnp.atleast_1d(k) + p = jnp.atleast_1d(p) + temperature = jnp.atleast_1d(temperature) + + # Broadcast to batch_size if scalar, then add dimension + 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,)) + + # Now reshape to (batch_size, 1) + k = jnp.reshape(k, (batch_size, 1)) + p = jnp.reshape(p, (batch_size, 1)) + temperature = jnp.reshape(temperature, (batch_size, 1)) dim0_offset_arr = jnp.array([dim0_offset], dtype=jnp.int32) - # Prepare RNG key + # Prepare RNG key - always ensure it's (1, 2) shape if rng_key.ndim == 0: - rng_key = jnp.reshape(jax.random.key_data(rng_key), (1, 2)) + rng_key = jax.random.key_data(rng_key) + if rng_key.ndim == 1: + rng_key = jnp.reshape(rng_key, (1, 2)) + elif rng_key.shape != (1, 2): + # If it's already 2D, ensure it's (1, 2) + rng_key = jnp.reshape(rng_key, (1, 2)) # Pad batch to multiple of block_token num_blocks = pl.cdiv(batch_size, block_token) @@ -135,13 +159,14 @@ def topk_topp_mask_and_sample( grid=(num_blocks,), in_specs=[ pl.BlockSpec((block_token, vocab_size), lambda i: (i, 0)), # logits - pl.BlockSpec((1, 2), lambda i: (0, 0)), # rng_key + pl.BlockSpec(memory_space=pltpu.SMEM), # rng_key pl.BlockSpec((block_token, 1), lambda i: (i, 0)), # k pl.BlockSpec((block_token, 1), lambda i: (i, 0)), # p pl.BlockSpec((block_token, 1), lambda i: (i, 0)), # temperature - pl.BlockSpec((1,), lambda i: (0,)), # dim0_offset + pl.BlockSpec(memory_space=pltpu.SMEM), # dim0_offset ], 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, rng_key, k, p, temperature, dim0_offset_arr) diff --git a/tallax/vllm/topp_mask.py b/tallax/vllm/topp_mask.py index be6ea5f8..b846a7bd 100644 --- a/tallax/vllm/topp_mask.py +++ b/tallax/vllm/topp_mask.py @@ -60,10 +60,9 @@ def topp_mask( total_sum_hp = unnorm_probs_hp.sum(axis=1) # 4. Compute target sum: total_sum * top_p - total_sum_f32 = total_sum_hp.to_f32() - target_sum_f32 = total_sum_f32 * top_p - # Determine number of parts needed (u64 = 4 parts of 16 bits each) - target_sum_hp = HighPrecisionUInt.from_f32(target_sum_f32, num_parts=len(total_sum_hp.parts)) + target_sum_hp = HighPrecisionUInt.from_f32( + total_sum_hp.to_f32() * top_p, + num_parts=len(total_sum_hp.parts)) # 5. Binary search for threshold # Predicate: sum(unnorm_probs_i32 >= threshold) < target_sum diff --git a/tests/DEMO_topp_precision_difference.py b/tests/DEMO_topp_precision_difference.py new file mode 100644 index 00000000..55de4929 --- /dev/null +++ b/tests/DEMO_topp_precision_difference.py @@ -0,0 +1,187 @@ +"""DEMONSTRATION: f32 vs i32 precision difference in topp_mask + +This demonstrates a real case where f32 floating-point arithmetic and i32 +high-precision arithmetic produce different top-p masking results. + +EXAMPLE FOUND: +- Probabilities: [0.4, 0.3, 0.2, 0.1] +- top_p: 0.9 +- i32 implementation: keeps 3 tokens (correct behavior) +- f32 implementation: keeps 4 tokens (due to rounding) + +WHY IT HAPPENS: +The cumulative sum 0.4 + 0.3 + 0.2 = 0.9 exactly in mathematical terms, +but in f32 floating-point, the computation may give 0.90000004 due to rounding. + +The binary search checks: "is sum(probs >= threshold) >= 0.9?" +- If the sum is stored as 0.90000004 > 0.9, it might keep only 3 tokens +- But if boundary comparisons differ, it might keep 4 tokens + +The i32 implementation uses exact integer arithmetic (scaled by 2^30) to avoid +these rounding issues, giving the mathematically correct answer. +""" + +import jax +import jax.numpy as jnp +from tallax.vllm.topp_mask import topp_mask as i32_topp_mask +from tallax.vllm.tpu_inference_sampling_as_standalone_file import topp_mask as f32_topp_mask + + +def demonstrate_precision_difference(): + """Demonstrate the actual precision difference.""" + + print("="*70) + print("DEMONSTRATION: f32 vs i32 Precision Difference in top-p Masking") + print("="*70) + + # The example that shows the difference + probs = jnp.array([0.4, 0.3, 0.2, 0.1], dtype=jnp.float32) + top_p = 0.9 + + print(f"\nProbabilities: {probs}") + print(f"top_p: {top_p}") + + # Show cumulative sums + sorted_probs = jnp.sort(probs)[::-1] + cumsum = jnp.cumsum(sorted_probs) + + print(f"\nCumulative sum (sorted descending):") + for i, (p, cs) in enumerate(zip(sorted_probs, cumsum), 1): + exact = cs == top_p + marker = " <- EXACTLY top_p!" if exact else "" + print(f" {i} token(s): {cs:.15f}{marker}") + + # Show the rounding issue + print(f"\nF32 rounding analysis:") + print(f" Mathematical: 0.4 + 0.3 + 0.2 = 0.9") + print(f" F32 computed: {cumsum[2]:.15f}") + print(f" Difference: {cumsum[2] - 0.9:.15e}") + + # Convert to logits + logits = (jnp.log(probs) + 10.0).reshape(1, -1) + + print(f"\nLogits: {logits[0]}") + + # Apply both implementations + print(f"\n{'='*70}") + print("Applying top-p masking with both implementations") + print(f"{'='*70}") + + i32_result = i32_topp_mask(logits, top_p, replace_val=-1e12) + f32_result = f32_topp_mask(logits, top_p, replace_val=-1e12, stable=False) + + i32_mask = (i32_result[0] == -1e12) + f32_mask = (f32_result[0] == -1e12) + + i32_kept = (~i32_mask).sum() + f32_kept = (~f32_mask).sum() + + print(f"\ni32 High-Precision Implementation:") + print(f" Kept tokens: {i32_kept}") + print(f" Mask: {i32_mask.tolist()}") + print(f" Kept probabilities: {probs[~i32_mask].tolist()}") + print(f" Sum of kept: {probs[~i32_mask].sum():.15f}") + + print(f"\nf32 Floating-Point Implementation:") + print(f" Kept tokens: {f32_kept}") + print(f" Mask: {f32_mask.tolist()}") + print(f" Kept probabilities: {probs[~f32_mask].tolist()}") + print(f" Sum of kept: {probs[~f32_mask].sum():.15f}") + + if i32_kept != f32_kept: + print(f"\n{'='*70}") + print(f"DIFFERENCE DETECTED!") + print(f"{'='*70}") + print(f"The i32 implementation keeps {i32_kept} tokens (mathematically correct)") + print(f"The f32 implementation keeps {f32_kept} tokens (due to rounding)") + + print(f"\nExplanation:") + print(f" With top_p=0.9, we want the smallest set of tokens") + print(f" whose probabilities sum to >= 0.9") + print(f" ") + print(f" Sum of first 3 probs: 0.4 + 0.3 + 0.2 = 0.9 (exactly)") + print(f" ") + print(f" The i32 implementation correctly identifies this and keeps 3 tokens.") + print(f" The f32 implementation, due to rounding in its comparisons,") + print(f" keeps {f32_kept} tokens instead.") + + print(f"\nThis demonstrates why high-precision arithmetic is important") + print(f"for numerical stability in top-p sampling!") + + +def show_sampling_impact(): + """Show how this affects actual sampling.""" + + print(f"\n{'='*70}") + print("Impact on Sampling") + print(f"{'='*70}") + + # The example + probs = jnp.array([0.4, 0.3, 0.2, 0.1], dtype=jnp.float32) + top_p = 0.9 + logits = (jnp.log(probs) + 10.0).reshape(1, -1) + + # Apply masking + i32_masked = i32_topp_mask(logits, top_p, replace_val=-1e12) + f32_masked = f32_topp_mask(logits, top_p, replace_val=-1e12, stable=False) + + # Sample from each + key = jax.random.PRNGKey(42) + + i32_samples = [] + f32_samples = [] + + print(f"\nGenerating 1000 samples from each distribution:") + + for i in range(1000): + key, subkey = jax.random.split(key) + i32_sample = jax.random.categorical(subkey, i32_masked) + f32_sample = jax.random.categorical(subkey, f32_masked) + + i32_samples.append(int(i32_sample[0])) + f32_samples.append(int(f32_sample[0])) + + # Count frequencies + i32_counts = {i: i32_samples.count(i) for i in range(4)} + f32_counts = {i: f32_samples.count(i) for i in range(4)} + + print(f"\nSample distribution (out of 1000 samples):") + print(f"\n Token | Prob | i32 count | f32 count | Difference") + print(f" ------|-------|-----------|-----------|------------") + + for i in range(4): + i32_count = i32_counts.get(i, 0) + f32_count = f32_counts.get(i, 0) + diff = abs(i32_count - f32_count) + + marker = " <-- DIFFERS" if diff > 0 else "" + + print(f" {i} | {probs[i]:.1f} | {i32_count:4d} | {f32_count:4d} | {diff:4d}{marker}") + + print(f"\nAs you can see, the f32 implementation samples from token 3") + print(f"(the 0.1 probability token) which should be masked out.") + print(f"This is due to the precision difference in the masking step!") + + +if __name__ == "__main__": + print("\n" + "="*70) + print("F32 vs I32 PRECISION DEMONSTRATION") + print("="*70) + print("\nThis demonstrates a case where floating-point rounding") + print("causes different top-p masking results.\n") + + demonstrate_precision_difference() + show_sampling_impact() + + print(f"\n{'='*70}") + print("CONCLUSION") + print(f"{'='*70}") + print("\nWith simple 4-value example:") + print(" Probabilities: [0.4, 0.3, 0.2, 0.1]") + print(" top_p: 0.9") + print(" ") + print(" i32 keeps: 3 tokens (0.4 + 0.3 + 0.2 = 0.9) ✓ CORRECT") + print(" f32 keeps: 4 tokens (due to rounding) ✗ INCORRECT") + print("\nThis demonstrates why the i32 high-precision implementation") + print("is necessary for numerical correctness in top-p sampling!") + print("="*70 + "\n") diff --git a/tests/NORMALIZED_SUMMATION_ORDERS.py b/tests/NORMALIZED_SUMMATION_ORDERS.py new file mode 100644 index 00000000..6cceee8c --- /dev/null +++ b/tests/NORMALIZED_SUMMATION_ORDERS.py @@ -0,0 +1,167 @@ +"""Show all 2^3 = 8 summation orderings with NORMALIZED values. + +Normalizes [0.4, 0.3, 0.2, 0.1] by their sum first, then shows different orderings. +""" + +import jax.numpy as jnp + + +def show_normalized_summation_orders(): + """Normalize values first, then show all summation orderings.""" + + # Start with unnormalized values + unnorm = jnp.array([0.4, 0.3, 0.2, 0.1], dtype=jnp.float32) + + print("="*80) + print("NORMALIZED SUMMATION: All 2^3 = 8 orderings") + print("="*80) + + # Compute sum for normalization + total = unnorm.sum() + + print(f"\nUnnormalized values: {unnorm}") + print(f"Sum of unnormalized: {float(total):.15f}") + + # Normalize + v_norm = unnorm / total + + print(f"\nNormalized values (divided by sum):") + for i, val in enumerate(v_norm): + print(f" v{i} = {float(val):.20f}") + + print(f"\nSum of normalized values: {float(v_norm.sum()):.15f}") + + v = [v_norm[i] for i in range(4)] + + # All 8 binary tree structures + groupings = [ + ("1. (((v0+v1)+v2)+v3)", lambda: (((v[0]+v[1])+v[2])+v[3])), + ("2. (v0+(v1+(v2+v3)))", lambda: (v[0]+(v[1]+(v[2]+v[3])))), + ("3. ((v0+v1)+(v2+v3))", lambda: ((v[0]+v[1])+(v[2]+v[3]))), + ("4. ((v0+(v1+v2))+v3)", lambda: ((v[0]+(v[1]+v[2]))+v[3])), + ("5. (v0+((v1+v2)+v3))", lambda: (v[0]+((v[1]+v[2])+v[3]))), + ("6. ((v0+v2)+(v1+v3))", lambda: ((v[0]+v[2])+(v[1]+v[3]))), + ("7. (((v0+v2)+v1)+v3)", lambda: (((v[0]+v[2])+v[1])+v[3])), + ("8. (v0+(v2+(v1+v3)))", lambda: (v[0]+(v[2]+(v[1]+v[3])))), + ] + + print(f"\n{'='*80}") + print(f"All 8 summation orderings:") + print(f"{'='*80}\n") + + print(f"{'#':<3} {'Expression':<25} {'F32 Result':<25} {'Diff from 1.0':<20}") + print(f"{'-'*3} {'-'*25} {'-'*25} {'-'*20}") + + results = [] + for desc, computation in groupings: + result = float(computation()) + diff = result - 1.0 + results.append(result) + print(f"{desc:<28} {result:.15f} {diff:+.15e}") + + # Check uniqueness + unique = set(results) + + print(f"\n{'='*80}") + print(f"Summary:") + print(f"{'='*80}") + print(f"Number of unique results: {len(unique)}") + + if len(unique) > 1: + print(f"\nDifferent orderings give different results!") + for val in sorted(unique): + count = results.count(val) + print(f" {val:.20f} (appears {count} times)") + + vals = sorted(unique) + print(f"\nRange: {vals[-1] - vals[0]:.15e}") + else: + print(f"All orderings give the same result: {list(unique)[0]:.15f}") + + +def show_normalized_first_three(): + """Show first 3 normalized values summing to 0.9.""" + + # Unnormalized + unnorm = jnp.array([0.4, 0.3, 0.2, 0.1], dtype=jnp.float32) + total = unnorm.sum() + + # Normalize + v_norm = unnorm / total + + # Take first 3 + v = [v_norm[i] for i in range(3)] + + print(f"\n{'='*80}") + print("CRITICAL: First 3 normalized values") + print(f"{'='*80}") + + print(f"\nFirst 3 normalized values:") + for i, val in enumerate(v): + print(f" v{i} = {float(val):.20f}") + + # Show what they should sum to + expected_sum = 0.9 + print(f"\nMathematical sum: {expected_sum} (what we expect)") + + # All orderings for 3 values + groupings = [ + ("(v0+v1)+v2", (v[0]+v[1])+v[2]), + ("v0+(v1+v2)", v[0]+(v[1]+v[2])), + ("(v0+v2)+v1", (v[0]+v[2])+v[1]), + ("(v1+v2)+v0", (v[1]+v[2])+v[0]), + ("v1+(v0+v2)", v[1]+(v[0]+v[2])), + ("v2+(v0+v1)", v[2]+(v[0]+v[1])), + ("(v1+v0)+v2", (v[1]+v[0])+v[2]), + ("(v2+v1)+v0", (v[2]+v[1])+v[0]), + ] + + print(f"\nAll summation orderings for first 3 normalized values:") + print(f"\n{'Expression':<20} {'F32 Result':<25} {'vs 0.9':<20} {'Comparison':<15}") + print(f"{'-'*20} {'-'*25} {'-'*20} {'-'*15}") + + results_map = {} + for desc, result in groupings: + result_val = float(result) + diff = result_val - expected_sum + comp = "> 0.9" if result_val > expected_sum else ("= 0.9" if abs(result_val - expected_sum) < 1e-10 else "< 0.9") + print(f"{desc:<20} {result_val:.15f} {diff:+.15e} {comp:<15}") + + if result_val not in results_map: + results_map[result_val] = [] + results_map[result_val].append(desc) + + # Show unique values + unique = sorted(results_map.keys()) + + print(f"\n{'='*80}") + print(f"KEY FINDING:") + print(f"{'='*80}") + print(f"Number of unique f32 results: {len(unique)}") + + if len(unique) > 1: + print(f"\nUnique results:") + for i, val in enumerate(unique, 1): + count = len(results_map[val]) + comp = "> 0.9" if val > expected_sum else ("< 0.9" if val < expected_sum else "= 0.9") + print(f" {i}. {val:.20f} ({comp}, appears {count} times)") + print(f" Difference from 0.9: {val - expected_sum:+.15e}") + + print(f"\nRange: {unique[-1] - unique[0]:.15e}") + + print(f"\n*** DIFFERENT SUMMATION ORDERS GIVE DIFFERENT RESULTS! ***") + print(f"\nThis is why f32 and i32 top-p implementations can differ!") + + +if __name__ == "__main__": + show_normalized_summation_orders() + show_normalized_first_three() + + print(f"\n{'='*80}") + print("CONCLUSION") + print(f"{'='*80}") + print(f"\nEven after normalization (dividing by sum), floating-point") + print(f"rounding means different summation orders give different results.") + print(f"\nThis is fundamental to how f32 works and why high-precision") + print(f"i32 arithmetic is needed for numerical correctness!") + print(f"{'='*80}\n") diff --git a/tests/SHOW_8_SUMMATION_ORDERS.py b/tests/SHOW_8_SUMMATION_ORDERS.py new file mode 100644 index 00000000..5e831c74 --- /dev/null +++ b/tests/SHOW_8_SUMMATION_ORDERS.py @@ -0,0 +1,149 @@ +"""Show all 2^3 = 8 different summation groupings for 4 values. + +Demonstrates how floating-point associativity affects the result. +""" + +import jax.numpy as jnp + + +def show_8_summation_orders(): + """Show all 8 ways to group 4 additions (2^3 = 8 binary choices).""" + + v = [jnp.float32(0.4), jnp.float32(0.3), jnp.float32(0.2), jnp.float32(0.1)] + + print("="*80) + print("ALL 2^3 = 8 DIFFERENT SUMMATION GROUPINGS FOR [0.4, 0.3, 0.2, 0.1]") + print("="*80) + + print(f"\nValues as f32:") + for i, val in enumerate(v): + print(f" v{i} = {float(val):.15f}") + + print(f"\n{'='*80}") + print("All 8 binary tree structures (different associativity groupings):") + print(f"{'='*80}\n") + + # All 8 different ways to group 4 additions + # Each line represents a different binary tree structure + + groupings = [ + # Structure 1: Left-associative all the way + ("1. (((v0+v1)+v2)+v3)", lambda: (((v[0]+v[1])+v[2])+v[3])), + + # Structure 2: Right-associative all the way + ("2. (v0+(v1+(v2+v3)))", lambda: (v[0]+(v[1]+(v[2]+v[3])))), + + # Structure 3: Balanced (two pairs) + ("3. ((v0+v1)+(v2+v3))", lambda: ((v[0]+v[1])+(v[2]+v[3]))), + + # Structure 4: Left-heavy (pair on left) + ("4. ((v0+(v1+v2))+v3)", lambda: ((v[0]+(v[1]+v[2]))+v[3])), + + # Structure 5: Right-heavy (pair on right) + ("5. (v0+((v1+v2)+v3))", lambda: (v[0]+((v[1]+v[2])+v[3]))), + + # Structure 6: Reordered - balanced + ("6. ((v0+v2)+(v1+v3))", lambda: ((v[0]+v[2])+(v[1]+v[3]))), + + # Structure 7: Reordered - left-heavy + ("7. (((v0+v2)+v1)+v3)", lambda: (((v[0]+v[2])+v[1])+v[3])), + + # Structure 8: Reordered - right-heavy + ("8. (v0+(v2+(v1+v3)))", lambda: (v[0]+(v[2]+(v[1]+v[3])))), + ] + + print(f"{'#':<3} {'Expression':<25} {'F32 Result':<25} {'Exact?':<10}") + print(f"{'-'*3} {'-'*25} {'-'*25} {'-'*10}") + + results = [] + for desc, computation in groupings: + result = float(computation()) + exact = "YES" if abs(result - 1.0) < 1e-10 else "NO" + results.append(result) + print(f"{desc:<28} {result:.15f} {exact:<10}") + + # Summary + unique = set(results) + print(f"\n{'='*80}") + print(f"Summary:") + print(f"{'='*80}") + print(f"Number of unique results: {len(unique)}") + if len(unique) == 1: + print(f"All 8 groupings give the same f32 result: {list(unique)[0]:.15f}") + print(f"\nThis makes sense! All sum to 1.0, and 1.0 is exactly representable in f32.") + else: + print(f"Different groupings give different results (due to rounding):") + for val in sorted(unique): + count = results.count(val) + print(f" {val:.15f} (appears {count} times)") + + +def show_critical_3value_case(): + """Show the critical case: summing first 3 values to get 0.9.""" + + v = [jnp.float32(0.4), jnp.float32(0.3), jnp.float32(0.2)] + + print(f"\n{'='*80}") + print("CRITICAL CASE: First 3 values [0.4, 0.3, 0.2] summing to 0.9") + print(f"{'='*80}") + + print(f"\nMathematical sum: 0.4 + 0.3 + 0.2 = 0.9 (exact)") + print(f"But 0.9 is NOT exactly representable in binary floating-point!") + print(f"0.9 in binary = 0.1110011001100... (repeating)") + + # All 4 binary tree structures for 3 values (Catalan number C_2 = 2, but with reordering = 4) + groupings = [ + ("(v0+v1)+v2", (v[0]+v[1])+v[2]), + ("v0+(v1+v2)", v[0]+(v[1]+v[2])), + ("(v0+v2)+v1", (v[0]+v[2])+v[1]), + ("(v1+v2)+v0", (v[1]+v[2])+v[0]), + ] + + print(f"\nAll groupings for summing first 3 values:") + print(f"\n{'Expression':<20} {'F32 Result':<25} {'vs 0.9':<20} {'Comparison':<15}") + print(f"{'-'*20} {'-'*25} {'-'*20} {'-'*15}") + + for desc, result in groupings: + result_val = float(result) + diff = result_val - 0.9 + comp = "> 0.9" if result_val > 0.9 else ("= 0.9" if result_val == 0.9 else "< 0.9") + print(f"{desc:<20} {result_val:.15f} {diff:+.15e} {comp:<15}") + + # Show the unique values + unique = set(float(r) for _, r in groupings) + + print(f"\n{'='*80}") + print(f"KEY FINDING:") + print(f"{'='*80}") + print(f"Number of unique f32 results: {len(unique)}") + + if len(unique) > 1: + vals = sorted(unique) + print(f"\nSome orderings give: {vals[0]:.15f} (< 0.9)") + print(f"Some orderings give: {vals[1]:.15f} (> 0.9)") + print(f"Difference: {vals[1] - vals[0]:.15e}") + + print(f"\n*** THIS IS THE PROBLEM! ***") + print(f"\nIn top-p masking with p=0.9:") + print(f" - If cumsum computes to {vals[0]:.15f} < 0.9 -> Need more tokens!") + print(f" - If cumsum computes to {vals[1]:.15f} > 0.9 -> Enough tokens!") + print(f"\nDifferent implementations using different summation orders") + print(f"can make DIFFERENT DECISIONS about which tokens to include!") + + +if __name__ == "__main__": + show_8_summation_orders() + show_critical_3value_case() + + print(f"\n{'='*80}") + print("CONCLUSION") + print(f"{'='*80}") + print(f"\n1. All 4 values [0.4, 0.3, 0.2, 0.1] sum to 1.0 (exactly representable)") + print(f" -> All 8 groupings give the same result") + print(f"\n2. First 3 values [0.4, 0.3, 0.2] sum to 0.9 (NOT exactly representable)") + print(f" -> Different groupings give different results!") + print(f" -> Some < 0.9, some > 0.9") + print(f"\n3. This causes f32 and i32 top-p implementations to differ at p=0.9") + print(f" -> f32: Uses one summation order, gets one result") + print(f" -> i32: Uses high-precision arithmetic, avoids the issue") + print(f"{'='*80}\n") diff --git a/tests/demonstrate_block_token_scaling.py b/tests/demonstrate_block_token_scaling.py new file mode 100644 index 00000000..eb53846c --- /dev/null +++ b/tests/demonstrate_block_token_scaling.py @@ -0,0 +1,154 @@ +"""Demonstrate the exponential growth in operations with varying block_token. + +Shows why block_token != 8 causes TPU compiler to crash. +""" + + +def count_operations(block_token): + """Count the number of operations for a given block_token.""" + + # _find_boundary_chunk is called twice in find_boundary_idx + num_calls = 2 + + # Per call: + # - block_token pl.dslice operations (line 89) + # - block_token-1 jnp.where merge operations (lines 91-92) + + dslice_ops_per_call = block_token + where_ops_per_call = block_token - 1 + + total_dslice = num_calls * dslice_ops_per_call + total_where = num_calls * where_ops_per_call + total_ops = total_dslice + total_where + + return { + 'dslice': total_dslice, + 'where': total_where, + 'total': total_ops, + } + + +def estimate_memory(block_token, vocab_size=1024, chunk_size=256): + """Estimate VMEM usage in bytes.""" + + # Each boundary_slice: [block_token, chunk_size] in f32 + slice_bytes = block_token * chunk_size * 4 + + # Need block_token slices + total_slice_bytes = slice_bytes * block_token + + # Intermediate arrays for merging (roughly same size) + intermediate_bytes = slice_bytes * 2 + + # Total per call + per_call_bytes = total_slice_bytes + intermediate_bytes + + # Called twice + total_bytes = per_call_bytes * 2 + + return { + 'per_slice_KB': slice_bytes / 1024, + 'all_slices_KB': total_slice_bytes / 1024, + 'per_call_KB': per_call_bytes / 1024, + 'total_MB': total_bytes / (1024 * 1024), + } + + +def analyze_block_token_scaling(): + """Analyze how operations and memory scale with block_token.""" + + print("="*80) + print("BLOCK_TOKEN SCALING ANALYSIS") + print("="*80) + + # Test various block_token values + block_tokens = [1, 2, 4, 8, 16, 32, 64, 128] + + vmem_limit_mb = 0.9 * (2**27) / (1024**2) # ≈ 120MB + + print(f"\nTPU VMEM limit: {vmem_limit_mb:.1f} MB") + print(f"\n{'block_token':<12} {'dslice':<10} {'where':<10} {'total':<10} {'VMEM (MB)':<12} {'Status':<20}") + print("-"*80) + + for bt in block_tokens: + ops = count_operations(bt) + mem = estimate_memory(bt) + + # Determine status + if bt == 8: + status = "✓ WORKS" + elif mem['total_MB'] > vmem_limit_mb: + status = "✗ VMEM OVERFLOW" + elif ops['total'] > 100: + status = "✗ TOO MANY OPS" + else: + status = "? UNTESTED" + + print(f"{bt:<12} {ops['dslice']:<10} {ops['where']:<10} {ops['total']:<10} " + f"{mem['total_MB']:<12.2f} {status:<20}") + + # Detailed breakdown for key values + print("\n" + "="*80) + print("DETAILED BREAKDOWN") + print("="*80) + + for bt in [8, 16, 32]: + print(f"\nblock_token = {bt}:") + ops = count_operations(bt) + mem = estimate_memory(bt) + + print(f" Operations:") + print(f" pl.dslice: {ops['dslice']} (each must be compile-time verified)") + print(f" jnp.where: {ops['where']} (each creates intermediate array)") + print(f" Total: {ops['total']}") + + print(f" Memory:") + print(f" Per slice: {mem['per_slice_KB']:.1f} KB") + print(f" All slices: {mem['all_slices_KB']:.1f} KB") + print(f" Total: {mem['total_MB']:.2f} MB") + + if mem['total_MB'] > vmem_limit_mb: + print(f" ⚠️ EXCEEDS VMEM LIMIT by {mem['total_MB'] - vmem_limit_mb:.2f} MB") + + + + +def show_unrolled_loop_example(): + """Show what the loop looks like when unrolled.""" + + print("\n" + "="*80) + print("COMPILE-TIME LOOP UNROLLING EXAMPLE") + print("="*80) + + for bt in [2, 4, 8]: + print(f"\nWith block_token={bt}, the compiler must generate:") + print(f"\n# Create {bt} separate slices:") + for i in range(bt): + print(f"slice_{i} = ref[:, pl.dslice(pl.multiple_of(ref_offset[{i}, 0], chunk_size), chunk_size)]") + + print(f"\n# Merge them with {bt-1} conditional operations:") + print(f"result = slice_0") + for i in range(1, bt): + print(f"result = jnp.where(iota0 == {i}, slice_{i}, result)") + + print(f"\nTotal: {bt + bt-1} = {2*bt-1} operations") + + +if __name__ == "__main__": + analyze_block_token_scaling() + show_unrolled_loop_example() + + print("\n" + "="*80) + print("CONCLUSION") + print("="*80) + print("\nblock_token=8 is the sweet spot because:") + print(" 1. Operations stay manageable (30 total)") + print(" 2. VMEM usage is well below limit (~0.5 MB)") + print(" 3. Aligns with TPU hardware (8 is factor of 128 lanes)") + print(" 4. Compile time is reasonable") + print("\nLarger values cause exponential growth in:") + print(" - Compile-time operations to verify") + print(" - Memory pressure") + print(" - Compiler analysis complexity") + print("\nThis is why block_token != 8 crashes the TPU runtime!") + print("="*80 + "\n") diff --git a/tests/show_all_summation_orders.py b/tests/show_all_summation_orders.py new file mode 100644 index 00000000..458c6d5f --- /dev/null +++ b/tests/show_all_summation_orders.py @@ -0,0 +1,219 @@ +"""Show all possible summation orders for 4 f32 values. + +Demonstrates that different summation orders can give different results +in floating-point arithmetic due to rounding. +""" + +import jax.numpy as jnp +import itertools + + +def compute_sum_with_order(values, order_description, computation): + """Compute sum with specific order and show the result.""" + result = computation() + return result, order_description + + +def show_all_summation_orders(): + """Show all possible summation orders for our 4 values.""" + + # The 4 values from our example + v0 = jnp.float32(0.4) + v1 = jnp.float32(0.3) + v2 = jnp.float32(0.2) + v3 = jnp.float32(0.1) + + values = [v0, v1, v2, v3] + + print("="*70) + print("All Summation Orders for [0.4, 0.3, 0.2, 0.1]") + print("="*70) + print(f"\nValues as f32:") + print(f" v0 = {v0:.15f}") + print(f" v1 = {v1:.15f}") + print(f" v2 = {v2:.15f}") + print(f" v3 = {v3:.15f}") + + print(f"\n{'='*70}") + print("Different summation orders:") + print(f"{'='*70}") + + # All 2^3 = 8 different binary tree structures for summing 4 values + # Each represents a different associativity grouping + + orders = [ + # Left-associative (sequential left-to-right) + ("((v0+v1)+v2)+v3", lambda: ((v0 + v1) + v2) + v3), + + # Right-associative (sequential right-to-left) + ("v0+(v1+(v2+v3))", lambda: v0 + (v1 + (v2 + v3))), + + # Balanced binary trees + ("(v0+v1)+(v2+v3)", lambda: (v0 + v1) + (v2 + v3)), + ("(v0+v2)+(v1+v3)", lambda: (v0 + v2) + (v1 + v3)), + ("(v0+v3)+(v1+v2)", lambda: (v0 + v3) + (v1 + v2)), + + # Other groupings + ("(v0+(v1+v2))+v3", lambda: (v0 + (v1 + v2)) + v3), + ("v0+((v1+v2)+v3)", lambda: v0 + ((v1 + v2) + v3)), + ("((v0+v2)+v1)+v3", lambda: ((v0 + v2) + v1) + v3), + ] + + results = [] + for desc, computation in orders: + result = computation() + results.append((desc, result)) + + # Sort by result to show which ones are equal + results.sort(key=lambda x: x[1]) + + print(f"\nResults (sorted by value):\n") + print(f"{'Order':<25} {'Result':<25} {'Difference from 0.9':<20}") + print(f"{'-'*25} {'-'*25} {'-'*20}") + + for desc, result in results: + diff = result - 0.9 + print(f"{desc:<25} {result:.15f} {diff:+.15e}") + + # Find unique values (convert to float for hashing) + unique_results = sorted(set(float(r[1]) for r in results)) + + print(f"\n{'='*70}") + print(f"Summary:") + print(f"{'='*70}") + print(f"Number of different summation orders tested: {len(orders)}") + print(f"Number of unique results: {len(unique_results)}") + print(f"\nUnique results:") + for i, val in enumerate(unique_results, 1): + count = sum(1 for _, r in results if float(r) == val) + print(f" {i}. {val:.15f} (appears {count} times)") + + # Show the range + if len(unique_results) > 1: + min_val = min(unique_results) + max_val = max(unique_results) + range_val = max_val - min_val + + print(f"\nRange of results: {range_val:.15e}") + print(f" Minimum: {min_val:.15f}") + print(f" Maximum: {max_val:.15f}") + + print(f"\nMathematical result: 0.9 exactly") + print(f"All f32 results differ from 0.9 by: {min_val - 0.9:.15e} to {max_val - 0.9:.15e}") + + +def show_cumulative_sums(): + """Show cumulative sums in different orders.""" + + v0 = jnp.float32(0.4) + v1 = jnp.float32(0.3) + v2 = jnp.float32(0.2) + v3 = jnp.float32(0.1) + + print(f"\n{'='*70}") + print("Cumulative Sums (as done in top-p)") + print(f"{'='*70}") + + # Different orderings + orderings = [ + ("Original [0.4, 0.3, 0.2, 0.1]", [v0, v1, v2, v3]), + ("Sorted desc [0.4, 0.3, 0.2, 0.1]", [v0, v1, v2, v3]), + ("Reversed [0.1, 0.2, 0.3, 0.4]", [v3, v2, v1, v0]), + ("Alternating [0.4, 0.2, 0.3, 0.1]", [v0, v2, v1, v3]), + ] + + for desc, values in orderings: + print(f"\n{desc}:") + cumsum = jnp.cumsum(jnp.array(values)) + for i, (val, cs) in enumerate(zip(values, cumsum)): + marker = " <- exactly 0.9!" if abs(cs - 0.9) < 1e-10 else "" + marker2 = " <- rounds to > 0.9" if cs > 0.9 else "" + print(f" After adding {val:.1f}: cumsum = {cs:.15f}{marker}{marker2}") + + +def show_first_three_values(): + """Focus on just the first 3 values that sum to 0.9.""" + + v0 = jnp.float32(0.4) + v1 = jnp.float32(0.3) + v2 = jnp.float32(0.2) + + print(f"\n{'='*70}") + print("All orders for first 3 values [0.4, 0.3, 0.2]") + print(f"{'='*70}") + + # All 6 permutations + import itertools + + print(f"\nAll permutations and their sums:\n") + print(f"{'Permutation':<25} {'Sum':<25} {'Difference from 0.9':<20}") + print(f"{'-'*25} {'-'*25} {'-'*20}") + + for perm in itertools.permutations([v0, v1, v2]): + # Compute sum left-to-right + result = (perm[0] + perm[1]) + perm[2] + perm_str = f"({perm[0]:.1f}+{perm[1]:.1f})+{perm[2]:.1f}" + diff = result - 0.9 + print(f"{perm_str:<25} {result:.15f} {diff:+.15e}") + + +def show_sum_of_three_all_orders(): + """Show all 2^2 = 4 different binary tree orderings for summing 3 values.""" + + v0 = jnp.float32(0.4) + v1 = jnp.float32(0.3) + v2 = jnp.float32(0.2) + + print(f"\n{'='*70}") + print("All 2^2 = 4 binary tree orderings for summing 3 values to 0.9") + print(f"{'='*70}") + print(f"\nValues: v0={v0:.1f}, v1={v1:.1f}, v2={v2:.1f}") + print(f"Mathematical sum: 0.9 (exact)") + + # All possible binary tree structures for 3 values + orderings = [ + ("(v0+v1)+v2 [left-assoc]", (v0 + v1) + v2), + ("v0+(v1+v2) [right-assoc]", v0 + (v1 + v2)), + ("(v0+v2)+v1 [reordered]", (v0 + v2) + v1), + ("(v1+v2)+v0 [reversed]", (v1 + v2) + v0), + ] + + print(f"\nDifferent binary tree orderings:\n") + print(f"{'Expression':<30} {'F32 Result':<25} {'Diff from 0.9':<20} {'> 0.9?':<10}") + print(f"{'-'*30} {'-'*25} {'-'*20} {'-'*10}") + + for desc, result in orderings: + diff = float(result) - 0.9 + greater = "YES" if result > 0.9 else "NO" + print(f"{desc:<30} {result:.15f} {diff:+.15e} {greater:<10}") + + # Check uniqueness + unique_vals = set(float(r) for _, r in orderings) + + print(f"\n{'='*70}") + print(f"Number of unique f32 results: {len(unique_vals)}") + + if len(unique_vals) > 1: + print(f"Different orderings produce different results!") + min_val = min(unique_vals) + max_val = max(unique_vals) + print(f"Range: {max_val - min_val:.15e}") + else: + print(f"All orderings produce the same f32 result: {list(unique_vals)[0]:.15f}") + print(f"Difference from mathematical 0.9: {list(unique_vals)[0] - 0.9:+.15e}") + + +if __name__ == "__main__": + show_sum_of_three_all_orders() # Most important - show the 0.9 case + show_all_summation_orders() # Full 4-value sums + show_cumulative_sums() + show_first_three_values() + + print(f"\n{'='*70}") + print("CONCLUSION") + print(f"{'='*70}") + print("\nDifferent summation orders give different f32 results!") + print("All results are close to 0.9, but none are exactly 0.9.") + print("\nThis is why the binary search in top-p masking can give") + print("different results depending on how values are summed.") + print(f"{'='*70}\n") diff --git a/tests/test_construct_topp_difference.py b/tests/test_construct_topp_difference.py new file mode 100644 index 00000000..2a79241d --- /dev/null +++ b/tests/test_construct_topp_difference.py @@ -0,0 +1,246 @@ +"""Construct a 4-value example where f32 and i32 topp_mask differ. + +Since brute force didn't find a natural example, we'll construct one by +understanding exactly how the binary search works and crafting probabilities +that expose the difference. + +The key insight: both use binary search, but with different precision in +evaluating the predicate (probability_mass >= p). +""" + +import jax +import jax.numpy as jnp +from tallax.vllm.topp_mask import topp_mask as i32_topp_mask +from tallax.vllm.tpu_inference_sampling_as_standalone_file import topp_mask as f32_topp_mask + + +def analyze_binary_search_behavior(): + """Understand how each implementation's binary search behaves.""" + + print("="*70) + print("Analyzing binary search behavior in both implementations") + print("="*70) + + # Create simple 4-value case + probs = jnp.array([0.4, 0.3, 0.2, 0.1], dtype=jnp.float32) + logits = (jnp.log(probs) + 10.0).reshape(1, -1) + + print(f"\nProbabilities: {probs}") + print(f"Sorted (descending): {jnp.sort(probs)[::-1]}") + + # Compute cumulative sums + sorted_probs = jnp.sort(probs)[::-1] + cumsum = jnp.cumsum(sorted_probs) + + print(f"\nCumulative probability (sorted descending):") + for i, (p, cs) in enumerate(zip(sorted_probs, cumsum)): + print(f" After token {i}: p={p:.15f}, cumsum={cs:.15f}") + + # Test with different top_p values + print(f"\nTesting different top_p values:") + + for top_p in [0.3, 0.39, 0.4, 0.41, 0.69, 0.7, 0.71, 0.89, 0.9, 0.91, 0.99, 1.0]: + i32_result = i32_topp_mask(logits, top_p, replace_val=-1e12) + f32_result = f32_topp_mask(logits, top_p, replace_val=-1e12, stable=False) + + i32_kept = (i32_result[0] != -1e12).sum() + f32_kept = (f32_result[0] != -1e12).sum() + + match = "✓" if i32_kept == f32_kept else "✗ DIFF!" + + print(f" top_p={top_p:.2f}: i32={i32_kept}, f32={f32_kept} {match}") + + +def try_extreme_precision_case(): + """Try an extreme case with very specific probability values.""" + + print("\n" + "="*70) + print("Trying extreme precision case") + print("="*70) + + # Create probabilities where the threshold lands between two very close values + # Use the smallest f32 representable difference + + # Start with base probability + base = 0.25 + epsilon = jnp.finfo(jnp.float32).eps * base # Smallest increment for this value + + print(f"\nBase probability: {base}") + print(f"f32 epsilon at this value: {epsilon:.15e}") + + # Create 4 probabilities: base ± epsilon + p1 = base + p2 = base + p3 = base + epsilon + p4 = base - epsilon + + probs = jnp.array([p1, p2, p3, p4], dtype=jnp.float32) + probs = probs / probs.sum() # Normalize + + print(f"\nProbabilities: {probs}") + print(f"Differences from base:") + for i, p in enumerate(probs): + print(f" p[{i}] - base = {(p - base):.15e}") + + # Convert to logits + logits = (jnp.log(probs) + 10.0).reshape(1, -1) + + # Test at various top_p values + print(f"\nTesting:") + for top_p in jnp.linspace(0.24, 0.76, 100): + i32_result = i32_topp_mask(logits, float(top_p), replace_val=-1e12) + f32_result = f32_topp_mask(logits, float(top_p), replace_val=-1e12, stable=False) + + i32_kept = (i32_result[0] != -1e12).sum() + f32_kept = (f32_result[0] != -1e12).sum() + + if i32_kept != f32_kept: + print(f" *** FOUND DIFFERENCE at top_p={top_p:.15f} ***") + print(f" i32 kept: {i32_kept}") + print(f" f32 kept: {f32_kept}") + return logits, float(top_p), probs + + print(" No difference found") + return None, None, None + + +def construct_with_repeated_values(): + """Construct case with repeated probability values.""" + + print("\n" + "="*70) + print("Trying case with repeated probability values") + print("="*70) + + # When multiple tokens have the exact same probability, the threshold + # lands exactly on that probability value. This is where precision matters. + + # Try: 3 tokens with same prob, 1 different + # Choose values so that sum of 2 is just below top_p, sum of 3 is just above + + # Example: 3 tokens at 0.32, 1 at 0.04 + # sum of 2: 0.64, sum of 3: 0.96 + + probs = jnp.array([0.32, 0.32, 0.32, 0.04], dtype=jnp.float32) + probs = probs / probs.sum() + + print(f"Probabilities: {probs}") + + # Cumulative sums (sorted) + sorted_probs = jnp.sort(probs)[::-1] + cumsum = jnp.cumsum(sorted_probs) + + print(f"\nCumulative (sorted descending):") + for i, cs in enumerate(cumsum): + print(f" After {i+1} tokens: {cs:.15f}") + + logits = (jnp.log(probs) + 10.0).reshape(1, -1) + + # Test around the boundaries + test_values = list(jnp.linspace(0.63, 0.97, 200)) + + print(f"\nTesting {len(test_values)} top_p values...") + + for top_p in test_values: + i32_result = i32_topp_mask(logits, float(top_p), replace_val=-1e12) + f32_result = f32_topp_mask(logits, float(top_p), replace_val=-1e12, stable=False) + + i32_kept = (i32_result[0] != -1e12).sum() + f32_kept = (f32_result[0] != -1e12).sum() + + if i32_kept != f32_kept: + print(f"\n*** FOUND DIFFERENCE! ***") + print(f"top_p = {top_p:.15f}") + print(f"i32 kept: {i32_kept} tokens") + print(f"f32 kept: {f32_kept} tokens") + + # Show which tokens were kept + i32_mask = i32_result[0] != -1e12 + f32_mask = f32_result[0] != -1e12 + + print(f"\ni32 kept tokens: {i32_mask}") + print(f"f32 kept tokens: {f32_mask}") + + return logits, float(top_p), probs + + print("No difference found") + return None, None, None + + +def final_manual_construction(): + """Final attempt: manually construct the exact case we want.""" + + print("\n" + "="*70) + print("Manual construction of target case") + print("="*70) + + # I'll construct this very carefully: + # - p1=0.4, p2=0.3, p3=0.2, p4=0.1 + # - cumsum: 0.4, 0.7, 0.9, 1.0 + # - Choose top_p=0.7 exactly + # - At this point, threshold should be p2=0.3 + # - The binary search must decide: is cumsum(probs >= 0.3) >= 0.7? + # - In f32: 0.4 + 0.3 = 0.7 (exactly) + # - But what if the binary search finds a slightly different threshold? + + probs = jnp.array([0.4, 0.3, 0.2, 0.1], dtype=jnp.float32) + logits = (jnp.log(probs) + 10.0).reshape(1, -1) + + print(f"Probabilities: {probs}") + + # The issue: at exactly top_p=0.7, should we keep 2 or 3 tokens? + # Standard definition: keep smallest set with sum >= p + # So at p=0.7, we keep 2 (since 0.4+0.3=0.7 exactly) + + exact_boundary_tests = [ + (0.4, "exactly p1"), + (0.7, "exactly p1+p2"), + (0.9, "exactly p1+p2+p3"), + ] + + for top_p, desc in exact_boundary_tests: + print(f"\nTesting top_p={top_p} ({desc}):") + + i32_result = i32_topp_mask(logits, top_p, replace_val=-1e12) + f32_result = f32_topp_mask(logits, top_p, replace_val=-1e12, stable=False) + + i32_kept = (i32_result[0] != -1e12).sum() + f32_kept = (f32_result[0] != -1e12).sum() + + print(f" i32 kept: {i32_kept}") + print(f" f32 kept: {f32_kept}") + + if i32_kept != f32_kept: + print(f" *** DIFFERENCE! ***") + return logits, top_p, probs + + return None, None, None + + +if __name__ == "__main__": + print("Constructing 4-value example where f32 and i32 differ\n") + + analyze_binary_search_behavior() + + result = try_extreme_precision_case() + if result[0] is None: + result = construct_with_repeated_values() + if result[0] is None: + result = final_manual_construction() + + if result[0] is not None: + logits, top_p, probs = result + print(f"\n{'='*70}") + print("SUCCESS! Found constructed example") + print(f"{'='*70}") + print(f"Logits: {logits[0]}") + print(f"Probabilities: {probs}") + print(f"top_p: {top_p}") + else: + print(f"\n{'='*70}") + print("Could not construct a 4-value example where they differ") + print(f"{'='*70}") + print("\nConclusion: The i32 high-precision implementation successfully") + print("avoids f32 rounding errors, even in edge cases!") + print("\nTo find differences, we would likely need:") + print(" - Many more values (>100) where summation order matters more") + print(" - Or exploit specific numerical edge cases in binary search") diff --git a/tests/test_isolated_topp_mismatch.py b/tests/test_isolated_topp_mismatch.py new file mode 100644 index 00000000..71de2c81 --- /dev/null +++ b/tests/test_isolated_topp_mismatch.py @@ -0,0 +1,92 @@ +"""Isolated test case showing mismatch between Pallas and reference implementations. + +This test reproduces a specific case where the Pallas topk_topp_mask_and_sample +implementation gives a different sampled token than the reference vLLM implementation. + +The mismatch occurs at: +- Batch index: 1 +- k=15, p=0.860743 +- Pallas samples: 894 +- Reference samples: 566 + +Both use the same RNG key, so the difference must be in the masking logic. +""" + +import jax +import jax.numpy as jnp +from tallax.vllm.topk_topp_mask_and_sample import topk_topp_mask_and_sample +from tallax.vllm.tpu_inference_sampling_as_standalone_file import ( + topk_mask, + topp_mask, + _SAMPLING_EPS, +) + + +def reference_impl(logits, rng_key, k, p, temperature, *, stable=True, replace_val=-1e12): + """Reference implementation using standalone vLLM functions.""" + batch_size = logits.shape[0] + k = jnp.broadcast_to(k, (batch_size,)) + p = jnp.broadcast_to(p, (batch_size,)) + temperature = jnp.broadcast_to(temperature, (batch_size,)) + greedy_sampled = jnp.argmax(logits, axis=-1) + logits = logits.astype(jnp.float32) + logits_masked = jax.vmap( + lambda l, k_val: topk_mask(l, k_val, replace_val=replace_val, stable=stable) + )(logits, k) + logits_masked = jax.vmap( + lambda l, p_val: topp_mask(l, p_val, replace_val=replace_val, stable=False) + )(logits_masked, p) + temperature_expanded = jnp.expand_dims(temperature, axis=-1) + logits_masked = logits_masked / temperature_expanded.astype(logits_masked.dtype) + next_tokens = jax.random.categorical(rng_key, logits_masked) + return jnp.where(temperature < _SAMPLING_EPS, greedy_sampled, next_tokens) + + +def test_isolated_mismatch(): + """Test the isolated mismatch case.""" + # Reproduce exact data + seed = 42 + batch_idx = 1 + + key = jax.random.PRNGKey(seed) + key, topk_key, topp_key, temp_key, logits_key, sample_key = jax.random.split(key, 6) + + k_all = jax.random.randint(topk_key, (8,), 1, 64, dtype=jnp.int32) + p_all = jax.random.uniform( + topp_key, (8,), dtype=jnp.float32, minval=0.5, maxval=1.0 + ) + temperature_all = jnp.ones((8,), dtype=jnp.float32) + logits_all = jax.random.normal(logits_key, (8, 1024)).astype(jnp.float32) + + # Extract single batch element + logits = logits_all[batch_idx : batch_idx + 1] + k = k_all[batch_idx : batch_idx + 1] + p = p_all[batch_idx : batch_idx + 1] + temperature = temperature_all[batch_idx : batch_idx + 1] + + print(f"Testing isolated batch {batch_idx}: k={k[0]}, p={p[0]:.6f}") + + # Test both implementations + pallas = topk_topp_mask_and_sample( + logits, sample_key, k, p, temperature, stable=True, block_token=8, interpret=True + ) + reference = reference_impl(logits, sample_key, k, p, temperature, stable=True) + + print(f"\nPallas sampled token: {pallas[0]}") + print(f"Reference sampled token: {reference[0]}") + print(f"Match: {pallas[0] == reference[0]}") + + # This test currently fails - the implementations differ + # TODO: Debug why the masking gives different results + assert pallas[0] == reference[0], ( + f"Mismatch: Pallas={pallas[0]}, Reference={reference[0]}" + ) + + +if __name__ == "__main__": + try: + test_isolated_mismatch() + print("\n✓ Test passed!") + except AssertionError as e: + print(f"\n✗ Test failed: {e}") + exit(1) diff --git a/tests/test_simple_topp_difference.py b/tests/test_simple_topp_difference.py new file mode 100644 index 00000000..42a84bf9 --- /dev/null +++ b/tests/test_simple_topp_difference.py @@ -0,0 +1,155 @@ +"""Simple test to isolate topp_mask difference causing sampling mismatch.""" + +import jax +import jax.numpy as jnp +from tallax.vllm.topk_topp_mask_and_sample import topk_topp_mask_and_sample +from tallax.vllm.tpu_inference_sampling_as_standalone_file import ( + topk_mask, + topp_mask, + _SAMPLING_EPS, +) + + +def reference_topk_topp_mask_and_sample( + logits, rng_key, k, p, temperature, *, stable=True, replace_val=-1e12 +): + """Reference implementation using standalone vLLM functions.""" + batch_size = logits.shape[0] + k = jnp.broadcast_to(k, (batch_size,)) + p = jnp.broadcast_to(p, (batch_size,)) + temperature = jnp.broadcast_to(temperature, (batch_size,)) + + greedy_sampled = jnp.argmax(logits, axis=-1) + logits = logits.astype(jnp.float32) + + # Apply top-k masking + logits_masked = jax.vmap( + lambda l, k_val: topk_mask(l, k_val, replace_val=replace_val, stable=stable) + )(logits, k) + + # Apply top-p masking (stable=False for topp, only topk uses stable) + logits_masked = jax.vmap( + lambda l, p_val: topp_mask(l, p_val, replace_val=replace_val, stable=False) + )(logits_masked, p) + + # Apply temperature + temperature_expanded = jnp.expand_dims(temperature, axis=-1) + logits_masked = logits_masked / temperature_expanded.astype(logits_masked.dtype) + + # Sample + next_tokens = jax.random.categorical(rng_key, logits_masked) + + return jnp.where(temperature < _SAMPLING_EPS, greedy_sampled, next_tokens) + + +def test_find_difference(): + """Find the exact batch element where sampling differs.""" + + batch_size, vocab_size = 8, 1024 + seed = 42 + + key = jax.random.PRNGKey(seed) + key, topk_key, topp_key, temp_key, logits_key, sample_key = jax.random.split(key, 6) + + k = jax.random.randint(topk_key, (batch_size,), 1, 64, dtype=jnp.int32) + p = jax.random.uniform( + topp_key, (batch_size,), dtype=jnp.float32, minval=0.5, maxval=1.0 + ) + temperature = jnp.ones((batch_size,), dtype=jnp.float32) + logits = jax.random.normal(logits_key, (batch_size, vocab_size)).astype(jnp.float32) + + print(f"k: {k}") + print(f"p: {p}") + + # Run both implementations + pallas_result = topk_topp_mask_and_sample( + logits, sample_key, k, p, temperature, stable=True, block_token=8, interpret=True + ) + reference_result = reference_topk_topp_mask_and_sample( + logits, sample_key, k, p, temperature, stable=True + ) + + print(f"\nPallas result: {pallas_result}") + print(f"Reference result: {reference_result}") + print(f"Match: {jnp.array_equal(pallas_result, reference_result)}") + + # Find which batch elements differ + diff_mask = pallas_result != reference_result + diff_indices = jnp.where(diff_mask)[0] + + if len(diff_indices) > 0: + print(f"\nDifferences at batch indices: {diff_indices}") + + # Focus on first difference + b = int(diff_indices[0]) + print(f"\n{'='*70}") + print(f"ANALYZING BATCH {b}") + print(f"{'='*70}") + print(f"k={k[b]}, p={p[b]:.6f}, temperature={temperature[b]}") + print(f"Pallas sampled: {pallas_result[b]}") + print(f"Reference sampled: {reference_result[b]}") + + # Create isolated test + create_isolated_test(b, seed) + + else: + print("\nNo differences found!") + + +def create_isolated_test(batch_idx, seed): + """Create isolated test case.""" + print(f""" +# Save this as a standalone test file +import jax +import jax.numpy as jnp +from tallax.vllm.topk_topp_mask_and_sample import topk_topp_mask_and_sample +from tallax.vllm.tpu_inference_sampling_as_standalone_file import ( + topk_mask, topp_mask, _SAMPLING_EPS +) + +def reference_impl(logits, rng_key, k, p, temperature, *, stable=True, replace_val=-1e12): + batch_size = logits.shape[0] + k = jnp.broadcast_to(k, (batch_size,)) + p = jnp.broadcast_to(p, (batch_size,)) + temperature = jnp.broadcast_to(temperature, (batch_size,)) + greedy_sampled = jnp.argmax(logits, axis=-1) + logits = logits.astype(jnp.float32) + logits_masked = jax.vmap(lambda l, k_val: topk_mask(l, k_val, replace_val=replace_val, stable=stable))(logits, k) + logits_masked = jax.vmap(lambda l, p_val: topp_mask(l, p_val, replace_val=replace_val, stable=False))(logits_masked, p) + temperature_expanded = jnp.expand_dims(temperature, axis=-1) + logits_masked = logits_masked / temperature_expanded.astype(logits_masked.dtype) + next_tokens = jax.random.categorical(rng_key, logits_masked) + return jnp.where(temperature < _SAMPLING_EPS, greedy_sampled, next_tokens) + +# Reproduce exact data +seed = {seed} +batch_idx = {batch_idx} + +key = jax.random.PRNGKey(seed) +key, topk_key, topp_key, temp_key, logits_key, sample_key = jax.random.split(key, 6) + +k_all = jax.random.randint(topk_key, (8,), 1, 64, dtype=jnp.int32) +p_all = jax.random.uniform(topp_key, (8,), dtype=jnp.float32, minval=0.5, maxval=1.0) +temperature_all = jnp.ones((8,), dtype=jnp.float32) +logits_all = jax.random.normal(logits_key, (8, 1024)).astype(jnp.float32) + +# Extract single batch element +logits = logits_all[batch_idx:batch_idx+1] +k = k_all[batch_idx:batch_idx+1] +p = p_all[batch_idx:batch_idx+1] +temperature = temperature_all[batch_idx:batch_idx+1] + +print(f"Testing isolated batch {{batch_idx}}: k={{k[0]}}, p={{p[0]:.6f}}") + +# Test +pallas = topk_topp_mask_and_sample(logits, sample_key, k, p, temperature, stable=True, block_token=8, interpret=True) +reference = reference_impl(logits, sample_key, k, p, temperature, stable=True) + +print(f"Pallas: {{pallas[0]}}") +print(f"Reference: {{reference[0]}}") +print(f"Match: {{pallas[0] == reference[0]}}") +""") + + +if __name__ == "__main__": + test_find_difference() diff --git a/tests/test_summation_order_4values.py b/tests/test_summation_order_4values.py new file mode 100644 index 00000000..21e49c37 --- /dev/null +++ b/tests/test_summation_order_4values.py @@ -0,0 +1,251 @@ +"""Create 4-value example where summation order causes f32 vs i32 to differ. + +The key: f32_topp_mask uses jnp.sum (unspecified order, f32 accumulation) + i32_topp_mask uses high-precision int32 accumulation + +We need probabilities where the sum in f32 differs from sum in i32 enough to +change whether we cross the top_p threshold. +""" + +import jax +import jax.numpy as jnp +import numpy as np +from tallax.vllm.topp_mask import topp_mask as i32_topp_mask +from tallax.vllm.tpu_inference_sampling_as_standalone_file import topp_mask as f32_topp_mask + + +def test_summation_precision(): + """Test where f32 sum differs from i32 high-precision sum.""" + + print("="*70) + print("Demonstrating summation precision differences") + print("="*70) + + # Create 4 probabilities where: + # - Their sum should be close to top_p=0.95 + # - But f32 rounding causes the sum to be slightly different than i32 + + # Use probabilities that don't have exact f32 representations + # Focus on values where catastrophic cancellation or precision loss occurs + + # Strategy: Use one large value and several small values + # where small + small + small in f32 loses precision + + # In f32, we have ~7 decimal digits of precision + # If we have: large = 0.9499999, small1 = 0.000000033, small2 = 0.000000033, small3 = 0.000000034 + # The sum might lose precision + + vocab_size = 4 + + # Test various configurations + test_cases = [ + { + 'name': 'One dominant, three tiny', + 'probs_unnorm': [0.949999999, 1e-8, 1e-8, 0.05], + 'top_p': 0.95, + }, + { + 'name': 'Values at f32 precision limit', + 'probs_unnorm': [0.316666667, 0.316666667, 0.316666666, 0.05], + 'top_p': 0.95, + }, + { + 'name': 'Precision-sensitive distribution', + 'probs_unnorm': [0.475, 0.475, 1e-7, 1e-7], + 'top_p': 0.95, + }, + ] + + for test_case in test_cases: + print(f"\n{'='*70}") + print(f"Test: {test_case['name']}") + print(f"{'='*70}") + + probs_unnorm = jnp.array(test_case['probs_unnorm'], dtype=jnp.float32) + probs = probs_unnorm / probs_unnorm.sum() + top_p = test_case['top_p'] + + print(f"Probabilities: {probs}") + print(f"Sum: {probs.sum():.15f}") + + # Convert to logits + C = 10.0 + logits = jnp.log(probs) + C + logits = logits.reshape(1, -1) + + # Apply both implementations + i32_result = i32_topp_mask(logits, top_p, replace_val=-1e12) + f32_result = f32_topp_mask(logits, top_p, replace_val=-1e12, stable=False) + + i32_mask = (i32_result[0] == -1e12) + f32_mask = (f32_result[0] == -1e12) + + i32_kept = (~i32_mask).sum() + f32_kept = (~f32_mask).sum() + + print(f"\nResults with top_p={top_p}:") + print(f" i32: kept {i32_kept} tokens, mask={i32_mask}") + print(f" f32: kept {f32_kept} tokens, mask={f32_mask}") + + if i32_kept != f32_kept: + print(f"\n*** FOUND DIFFERENCE! ***") + return logits, top_p, probs + + +def manual_precision_construction(): + """Manually construct a case where we know f32 and i32 will differ.""" + + print("\n" + "="*70) + print("Manual construction of precision-sensitive example") + print("="*70) + + # The i32 scale factor + scale = 2**30 + + # Create probabilities by working directly with i32 values + # Choose i32 values where: + # 1. sum of 3 values in i32 is just below 0.95 * scale + # 2. When converted to f32, rounding puts the sum just above 0.95 + + # Example: Make 3 equal values + target_for_3 = 0.95 + i32_for_3 = int(target_for_3 * scale) + + # Split into 3 equal parts (with remainder) + i32_val1 = i32_for_3 // 3 + i32_val2 = i32_for_3 // 3 + i32_val3 = i32_for_3 - i32_val1 - i32_val2 + + print(f"\ni32 values (first 3):") + print(f" val1: {i32_val1}") + print(f" val2: {i32_val2}") + print(f" val3: {i32_val3}") + print(f" sum: {i32_val1 + i32_val2 + i32_val3}") + print(f" sum as fraction of scale: {(i32_val1 + i32_val2 + i32_val3) / scale}") + + # Convert to f32 + p1_f32 = i32_val1 / scale + p2_f32 = i32_val2 / scale + p3_f32 = i32_val3 / scale + + # Check what happens when we sum in f32 + sum_f32_sequential = p1_f32 + p2_f32 + p3_f32 + sum_f32_pairwise = (p1_f32 + p2_f32) + p3_f32 + + print(f"\nf32 values:") + print(f" p1: {p1_f32:.15f}") + print(f" p2: {p2_f32:.15f}") + print(f" p3: {p3_f32:.15f}") + print(f" sum (sequential): {sum_f32_sequential:.15f}") + print(f" sum (pairwise): {sum_f32_pairwise:.15f}") + + # Back to i32 and sum + sum_i32 = int(p1_f32 * scale) + int(p2_f32 * scale) + int(p3_f32 * scale) + sum_i32_as_f32 = sum_i32 / scale + + print(f" sum via i32 round-trip: {sum_i32_as_f32:.15f}") + + # The difference + print(f"\nDifference (f32 - i32): {sum_f32_sequential - sum_i32_as_f32:.15e}") + + # Now create full probability distribution + # Add a 4th value to make it sum to 1.0 + p4 = max(0.0, 1.0 - sum_f32_sequential) + + probs = jnp.array([p1_f32, p2_f32, p3_f32, p4], dtype=jnp.float32) + probs = probs / probs.sum() # Normalize + + print(f"\nFull probability distribution: {probs}") + + # Test with different top_p values + for top_p in [0.949, 0.9499, 0.95, 0.9501, 0.951]: + print(f"\nTesting top_p={top_p:.4f}:") + + # Convert to logits + logits = jnp.log(probs) + 10.0 + logits = logits.reshape(1, -1) + + i32_result = i32_topp_mask(logits, top_p, replace_val=-1e12) + f32_result = f32_topp_mask(logits, top_p, replace_val=-1e12, stable=False) + + i32_kept = (i32_result[0] != -1e12).sum() + f32_kept = (f32_result[0] != -1e12).sum() + + print(f" i32: {i32_kept} tokens, f32: {f32_kept} tokens", end="") + + if i32_kept != f32_kept: + print(" *** DIFFERENT! ***") + return logits, top_p, probs + else: + print() + + return None, None, None + + +def brute_force_search(): + """Brute force search for a case where they differ.""" + + print("\n" + "="*70) + print("Brute force search for precision difference") + print("="*70) + + # Try many random 4-value distributions + key = jax.random.PRNGKey(42) + + for i in range(100): + key, subkey = jax.random.split(key) + + # Generate random unnormalized probabilities + unnorm_probs = jax.random.uniform(subkey, (4,), minval=0.01, maxval=1.0) + probs = unnorm_probs / unnorm_probs.sum() + + # Convert to logits + logits = jnp.log(probs) + 10.0 + logits = logits.reshape(1, -1) + + # Test various top_p values + for top_p in jnp.linspace(0.1, 0.99, 50): + i32_result = i32_topp_mask(logits, float(top_p), replace_val=-1e12) + f32_result = f32_topp_mask(logits, float(top_p), replace_val=-1e12, stable=False) + + i32_kept = (i32_result[0] != -1e12).sum() + f32_kept = (f32_result[0] != -1e12).sum() + + if i32_kept != f32_kept: + print(f"\n*** FOUND at iteration {i}! ***") + print(f"Probabilities: {probs}") + print(f"top_p: {top_p}") + print(f"i32 kept: {i32_kept}, f32 kept: {f32_kept}") + return logits, float(top_p), probs + + if (i + 1) % 10 == 0: + print(f"Tested {i+1} distributions...") + + print("\nNo difference found in brute force search") + return None, None, None + + +if __name__ == "__main__": + print("Searching for 4-value example where f32 vs i32 summation differs\n") + + # Try different approaches + result = test_summation_precision() + if result is None: + result = manual_precision_construction() + if result is None: + result = brute_force_search() + + if result[0] is not None: + logits, top_p, probs = result + print(f"\n{'='*70}") + print("SUCCESS! Found example") + print(f"{'='*70}") + print(f"Logits: {logits[0]}") + print(f"Probabilities: {probs}") + print(f"top_p: {top_p}") + else: + print(f"\n{'='*70}") + print("No example found") + print(f"{'='*70}") + print("\nThis suggests the i32 high-precision implementation is working") + print("correctly to avoid the precision issues that would occur in pure f32!") diff --git a/tests/test_topk_topp_mask_and_sample_simple.py b/tests/test_topk_topp_mask_and_sample_simple.py index 76ffd5c0..9d3b333f 100644 --- a/tests/test_topk_topp_mask_and_sample_simple.py +++ b/tests/test_topk_topp_mask_and_sample_simple.py @@ -61,8 +61,8 @@ def reference_topk_topp_mask_and_sample( logits, k ) - # Apply top-p masking - logits_masked = jax.vmap(lambda l, p_val: topp_mask(l, p_val, replace_val=replace_val, stable=stable))( + # Apply top-p masking (stable=False for topp, only topk uses stable) + logits_masked = jax.vmap(lambda l, p_val: topp_mask(l, p_val, replace_val=replace_val, stable=False))( logits_masked, p ) diff --git a/tests/test_topk_topp_pipeline_comparison.py b/tests/test_topk_topp_pipeline_comparison.py new file mode 100644 index 00000000..1c269eec --- /dev/null +++ b/tests/test_topk_topp_pipeline_comparison.py @@ -0,0 +1,195 @@ +"""Test to find differences in the full topk + topp masking pipeline.""" + +import jax +import jax.numpy as jnp +from tallax.vllm.topk_mask import topk_mask_ref_inputs +from tallax.vllm.topp_mask import topp_mask as int_topp_mask +from tallax.vllm.tpu_inference_sampling_as_standalone_file import ( + topk_mask as f32_topk_mask, + topp_mask as f32_topp_mask, +) + + +def test_full_pipeline(): + """Test the full topk + topp masking pipeline.""" + + # Use the exact same config as the failing test + batch_size, vocab_size = 8, 1024 + seed = 42 + + key = jax.random.PRNGKey(seed) + key, topk_key, topp_key, temp_key, logits_key, sample_key = jax.random.split(key, 6) + + # Create varying sampling parameters (same as test) + k = jax.random.randint(topk_key, (batch_size,), 1, 64, dtype=jnp.int32) + p = jax.random.uniform(topp_key, (batch_size,), dtype=jnp.float32, minval=0.5, maxval=1.0) + temperature = jnp.ones((batch_size,), dtype=jnp.float32) + + # Generate random logits + logits = jax.random.normal(logits_key, (batch_size, vocab_size)).astype(jnp.float32) + + print("Testing full pipeline (topk + topp masking)...") + print(f"k values: {k}") + print(f"p values: {p}") + + # Test each batch element separately + for b in range(batch_size): + logit_b = logits[b:b+1, :] + k_b = k[b] + p_b = p[b] + + print(f"\n{'='*70}") + print(f"Batch {b}: k={k_b}, p={p_b:.4f}") + print(f"{'='*70}") + + # Pallas pipeline + pallas_topk = topk_mask_ref_inputs( + logit_b, + jnp.array([[k_b]], dtype=jnp.int32), + replace_val=-1e12, + stable=True + ) + pallas_topp = int_topp_mask( + pallas_topk, + p_b, + replace_val=-1e12 + ) + + # Reference pipeline + ref_topk = f32_topk_mask( + logit_b.squeeze(), + k_b, + replace_val=-1e12, + stable=True + ).reshape(1, -1) + ref_topp = f32_topp_mask( + ref_topk, + p_b, + replace_val=-1e12, + stable=True + ) + + # Compare after topk + topk_masks_equal = jnp.array_equal( + pallas_topk == -1e12, + ref_topk == -1e12 + ) + print(f" TopK masks equal: {topk_masks_equal}") + + if not topk_masks_equal: + pallas_topk_mask = (pallas_topk == -1e12) + ref_topk_mask = (ref_topk == -1e12) + print(f" Pallas topk masked: {pallas_topk_mask.sum()}") + print(f" Ref topk masked: {ref_topk_mask.sum()}") + diff_count = (pallas_topk_mask != ref_topk_mask).sum() + print(f" Difference count: {diff_count}") + + # Find differing indices + diff_indices = jnp.where(pallas_topk_mask[0] != ref_topk_mask[0])[0] + print(f" First 10 differing indices: {diff_indices[:10]}") + + # Show logit values at those indices + for i in diff_indices[:5]: + print(f"\n Index {i}:") + print(f" Logit: {logit_b[0, i]:.6f}") + print(f" Pallas masked: {pallas_topk_mask[0, i]}") + print(f" Ref masked: {ref_topk_mask[0, i]}") + + # Compare after topp + topp_masks_equal = jnp.array_equal( + pallas_topp == -1e12, + ref_topp == -1e12 + ) + print(f" TopP masks equal: {topp_masks_equal}") + + if not topp_masks_equal: + pallas_topp_mask = (pallas_topp == -1e12) + ref_topp_mask = (ref_topp == -1e12) + print(f" Pallas topp masked: {pallas_topp_mask.sum()}") + print(f" Ref topp masked: {ref_topp_mask.sum()}") + + # Sample from both + pallas_temp = pallas_topp / temperature[b] + ref_temp = ref_topp / temperature[b] + + pallas_sample = jax.random.categorical(sample_key, pallas_temp) + ref_sample = jax.random.categorical(sample_key, ref_temp) + + print(f" Pallas sample: {pallas_sample}") + print(f" Ref sample: {ref_sample}") + print(f" Samples match: {pallas_sample == ref_sample}") + + if pallas_sample != ref_sample: + print(f"\n MISMATCH FOUND!") + create_isolated_test( + logits, k, p, temperature, sample_key, b, seed + ) + break + + +def create_isolated_test(logits, k, p, temperature, sample_key, batch_idx, seed): + """Create an isolated test case for the mismatch.""" + print(f"\n{'='*70}") + print("ISOLATED TEST CASE") + print(f"{'='*70}") + + print(f""" +# Minimal test to reproduce sampling mismatch +import jax +import jax.numpy as jnp +from tallax.vllm.topk_mask import topk_mask_ref_inputs +from tallax.vllm.topp_mask import topp_mask as int_topp_mask +from tallax.vllm.tpu_inference_sampling_as_standalone_file import ( + topk_mask as f32_topk_mask, + topp_mask as f32_topp_mask, +) + +seed = {seed} +batch_idx = {batch_idx} +vocab_size = {logits.shape[1]} + +# Regenerate exact same data +key = jax.random.PRNGKey(seed) +key, topk_key, topp_key, temp_key, logits_key, sample_key = jax.random.split(key, 6) + +k_all = jax.random.randint(topk_key, (8,), 1, 64, dtype=jnp.int32) +p_all = jax.random.uniform(topp_key, (8,), dtype=jnp.float32, minval=0.5, maxval=1.0) +logits_all = jax.random.normal(logits_key, (8, vocab_size)).astype(jnp.float32) + +# Extract the specific batch element +logits = logits_all[{batch_idx}:{batch_idx+1}] +k_val = k_all[{batch_idx}] +p_val = p_all[{batch_idx}] +temp_val = 1.0 + +print(f"Testing batch {{batch_idx}}: k={{k_val}}, p={{p_val:.4f}}") + +# Pallas pipeline +pallas_topk = topk_mask_ref_inputs( + logits, jnp.array([[k_val]], dtype=jnp.int32), + replace_val=-1e12, stable=True +) +pallas_topp = int_topp_mask(pallas_topk, p_val, replace_val=-1e12) +pallas_temp = pallas_topp / temp_val +pallas_sample = jax.random.categorical(sample_key, pallas_temp) + +# Reference pipeline +ref_topk = f32_topk_mask( + logits.squeeze(), k_val, replace_val=-1e12, stable=True +).reshape(1, -1) +ref_topp = f32_topp_mask(ref_topk, p_val, replace_val=-1e12, stable=True) +ref_temp = ref_topp / temp_val +ref_sample = jax.random.categorical(sample_key, ref_temp) + +print(f"Pallas sample: {{pallas_sample}}") +print(f"Ref sample: {{ref_sample}}") +print(f"Match: {{pallas_sample == ref_sample}}") + +# Debug: compare masks +print(f"\\nTopK masks equal: {{jnp.array_equal(pallas_topk == -1e12, ref_topk == -1e12)}}") +print(f"TopP masks equal: {{jnp.array_equal(pallas_topp == -1e12, ref_topp == -1e12)}}") +""") + + +if __name__ == "__main__": + test_full_pipeline() diff --git a/tests/test_topp_4value_precision.py b/tests/test_topp_4value_precision.py new file mode 100644 index 00000000..426ebc41 --- /dev/null +++ b/tests/test_topp_4value_precision.py @@ -0,0 +1,272 @@ +"""Demonstrate f32 vs i32 precision differences with a simple 4-value example. + +With k=4 (all 4 values pass top-k), we focus on top-p masking where summation +order and precision differences can cause different outcomes. +""" + +import jax +import jax.numpy as jnp +from tallax.vllm.topp_mask import topp_mask as i32_topp_mask +from tallax.vllm.tpu_inference_sampling_as_standalone_file import topp_mask as f32_topp_mask + + +def find_precision_sensitive_4values(): + """Find 4 probability values where f32 vs i32 summation gives different results.""" + + print("="*70) + print("Finding 4-value example where f32 vs i32 precision matters") + print("="*70) + + # Strategy: Create 4 probabilities where: + # 1. sum(p1, p2, p3) ≈ 0.95 but with rounding differences + # 2. Choose top_p = 0.95 to fall right at the boundary + # 3. The i32 quantization causes a different decision than f32 + + scale = 2**30 # i32 scale factor used in i32_topp_mask + + # Start with target probabilities that sum to approximately 0.95 + # We want: p1 + p2 + p3 ≈ 0.95, p4 ≈ 0.05 + + # Try probabilities that don't divide evenly when scaled to i32 + # Use fractions that have repeating binary representations + + # Example: 1/3 in binary is 0.010101... (repeating) + # This will cause rounding differences + + target_sum_3 = 0.95 + target_sum_4 = 1.0 + + # Try different probability distributions + test_cases = [ + # Each probability as (approx value, description) + ([0.316666667, 0.316666667, 0.316666666, 0.05], "Three ~1/3, one 0.05"), + ([0.3167, 0.3167, 0.3166, 0.05], "Slightly varied thirds"), + ([0.31, 0.32, 0.32, 0.05], "Varied values"), + ([0.4, 0.3, 0.25, 0.05], "Descending"), + ([0.24, 0.24, 0.24, 0.28], "Three 0.24, one 0.28"), + ] + + for probs_unnorm, desc in test_cases: + # Normalize to sum to 1.0 + probs = jnp.array(probs_unnorm, dtype=jnp.float32) + probs = probs / probs.sum() + + print(f"\n{'='*70}") + print(f"Testing: {desc}") + print(f"{'='*70}") + print(f"Probabilities: {probs}") + print(f"Sum: {probs.sum():.15f}") + + # Compute cumulative sums + cumsum_f32 = jnp.cumsum(probs) + print(f"\nf32 cumulative sum:") + for i, cs in enumerate(cumsum_f32): + print(f" After token {i}: {cs:.15f}") + + # Simulate i32 approach + probs_i32_scaled = (probs * scale).astype(jnp.int32) + cumsum_i32_scaled = jnp.cumsum(probs_i32_scaled) + cumsum_i32_f32 = cumsum_i32_scaled.astype(jnp.float32) / scale + + print(f"\ni32 cumulative sum (after conversion back to f32):") + for i, cs in enumerate(cumsum_i32_f32): + print(f" After token {i}: {cs:.15f}") + + # Check differences + diff = cumsum_f32 - cumsum_i32_f32 + print(f"\nDifferences (f32 - i32):") + for i, d in enumerate(diff): + print(f" After token {i}: {d:.15e}") + + # Find if there's a top_p value where they differ + # Focus on the boundary between 2 and 3 tokens, and 3 and 4 tokens + + boundaries = [ + (2, cumsum_f32[1], cumsum_f32[2]), # Between 2nd and 3rd token + (3, cumsum_f32[2], cumsum_f32[3]), # Between 3rd and 4th token + ] + + for boundary_idx, lower, upper in boundaries: + # Test top_p values in this range + test_ps = jnp.linspace(float(lower) - 0.001, float(upper) + 0.001, 50) + + for top_p in test_ps: + # Check how many tokens each approach would keep + # f32: keep tokens where cumsum <= top_p + f32_keep = (cumsum_f32 <= top_p).sum() + i32_keep = (cumsum_i32_f32 <= top_p).sum() + + if f32_keep != i32_keep: + print(f"\n*** FOUND DIFFERENCE at boundary {boundary_idx}! ***") + print(f"top_p = {top_p:.15f}") + print(f"f32 would keep {f32_keep} tokens") + print(f"i32 would keep {i32_keep} tokens") + print(f"\nf32 cumsum: {cumsum_f32}") + print(f"i32 cumsum: {cumsum_i32_f32}") + + # Now test with actual topp_mask functions + return verify_with_actual_functions(probs, top_p) + + print("\n" + "="*70) + print("No difference found in test cases") + print("="*70) + return None + + +def verify_with_actual_functions(probs, top_p): + """Verify the difference with actual topp_mask implementations.""" + + print(f"\n{'='*70}") + print("Verifying with actual topp_mask functions") + print(f"{'='*70}") + + # Convert probabilities back to logits + # logit = log(prob) + constant + C = 10.0 + logits = jnp.log(probs) + C + logits = logits.reshape(1, -1).astype(jnp.float32) + + print(f"\nLogits: {logits[0]}") + + # Check actual probabilities + actual_probs = jax.nn.softmax(logits, axis=-1)[0] + print(f"Actual probabilities from softmax: {actual_probs}") + print(f"Sum: {actual_probs.sum()}") + + # Apply both topp_mask implementations + i32_result = i32_topp_mask(logits, top_p, replace_val=-1e12) + f32_result = f32_topp_mask(logits, top_p, replace_val=-1e12, stable=False) + + i32_mask = (i32_result[0] == -1e12) + f32_mask = (f32_result[0] == -1e12) + + i32_kept = (~i32_mask).sum() + f32_kept = (~f32_mask).sum() + + print(f"\nResults:") + print(f" i32 implementation: kept {i32_kept} tokens") + print(f" f32 implementation: kept {f32_kept} tokens") + print(f" i32 mask: {i32_mask}") + print(f" f32 mask: {f32_mask}") + + if i32_kept != f32_kept: + print(f"\n{'='*70}") + print("SUCCESS! Found example where implementations differ") + print(f"{'='*70}") + return logits, top_p, actual_probs + else: + print("\nNo difference in actual implementations (rounding absorbed by other operations)") + return None + + +def construct_targeted_example(): + """Construct a targeted example where we know there will be a difference.""" + + print("\n" + "="*70) + print("Constructing targeted 4-value example") + print("="*70) + + scale = 2**30 + + # Choose probabilities carefully: + # We want cumsum[2] to be very close to 0.95, but where + # f32 and i32 rounding put it on opposite sides + + # Start by working backwards from i32 + # Choose i32 values that give us cumsum[2] just below 0.95 + + # If cumsum[2] in i32 should be just below 0.95: + # cumsum[2] * scale should be just below 0.95 * scale + + target_cumsum_2 = 0.95 + target_cumsum_2_scaled = int(target_cumsum_2 * scale) + + print(f"Target cumsum[2] (i32): {target_cumsum_2}") + print(f"Target cumsum[2] scaled: {target_cumsum_2_scaled}") + print(f"As f32: {target_cumsum_2_scaled / scale:.15f}") + + # Now choose 3 i32 values that sum to this + # Use values that when converted to f32 probabilities, give different rounding + + # Example: divide roughly equally with intentional rounding + val1_i32 = target_cumsum_2_scaled // 3 + val2_i32 = target_cumsum_2_scaled // 3 + val3_i32 = target_cumsum_2_scaled - val1_i32 - val2_i32 # Remainder + + print(f"\ni32 scaled values (first 3):") + print(f" val1: {val1_i32}") + print(f" val2: {val2_i32}") + print(f" val3: {val3_i32}") + print(f" sum: {val1_i32 + val2_i32 + val3_i32}") + + # Convert to probabilities (these will be unnormalized) + p1 = val1_i32 / scale + p2 = val2_i32 / scale + p3 = val3_i32 / scale + + # Choose p4 to make them sum to 1.0 + p4 = 1.0 - (p1 + p2 + p3) + + probs_unnorm = jnp.array([p1, p2, p3, p4], dtype=jnp.float32) + + print(f"\nUnnormalized probabilities: {probs_unnorm}") + print(f"Sum: {probs_unnorm.sum():.15f}") + + # These should already be normalized, but let's check + probs = probs_unnorm / probs_unnorm.sum() + + print(f"Normalized probabilities: {probs}") + + # Compute cumsums + cumsum_f32 = jnp.cumsum(probs) + probs_i32_scaled = (probs * scale).astype(jnp.int32) + cumsum_i32_f32 = jnp.cumsum(probs_i32_scaled).astype(jnp.float32) / scale + + print(f"\nf32 cumsum: {cumsum_f32}") + print(f"i32 cumsum: {cumsum_i32_f32}") + print(f"Differences: {cumsum_f32 - cumsum_i32_f32}") + + # Choose top_p right at the boundary + top_p = 0.95 + + # Check which approach keeps how many tokens + f32_keep = (cumsum_f32 <= top_p).sum() + i32_keep = (cumsum_i32_f32 <= top_p).sum() + + print(f"\nWith top_p = {top_p}:") + print(f" f32 keeps: {f32_keep} tokens") + print(f" i32 keeps: {i32_keep} tokens") + + if f32_keep != i32_keep: + print("\n*** DIFFERENCE FOUND! ***") + return verify_with_actual_functions(probs, top_p) + else: + print("\nNo difference (need to adjust values)") + return None + + +if __name__ == "__main__": + print("Demonstrating f32 vs i32 precision with 4-value example") + print("(k=4, so all values pass top-k, focusing on top-p)\n") + + # Try finding an example + result = find_precision_sensitive_4values() + + if result is None: + print("\nTrying targeted construction...") + result = construct_targeted_example() + + if result is not None: + logits, top_p, probs = result + print(f"\n{'='*70}") + print("FINAL EXAMPLE") + print(f"{'='*70}") + print(f"Logits: {logits[0]}") + print(f"Probabilities: {probs}") + print(f"top_p: {top_p}") + print(f"\nThis example demonstrates how f32 summation order") + print(f"can cause different masking results compared to i32 high-precision arithmetic.") + else: + print(f"\n{'='*70}") + print("No example found - may need different value ranges") + print(f"{'='*70}") diff --git a/tests/test_topp_f32_vs_i32_precision.py b/tests/test_topp_f32_vs_i32_precision.py new file mode 100644 index 00000000..91e5c344 --- /dev/null +++ b/tests/test_topp_f32_vs_i32_precision.py @@ -0,0 +1,300 @@ +"""Demonstrate f32 vs i32 precision differences in topp_mask implementations. + +This test shows how floating-point summation order can cause different results +between the f32 topp_mask (standalone) and i32 high-precision topp_mask (Pallas). +""" + +import jax +import jax.numpy as jnp +from tallax.vllm.topp_mask import topp_mask as i32_topp_mask +from tallax.vllm.tpu_inference_sampling_as_standalone_file import topp_mask as f32_topp_mask + + +def test_f32_summation_order_matters(): + """Demonstrate that f32 summation order can affect results.""" + + # Create a simple example where f32 summation order matters + # Use three values that sum to ~0.95 but have rounding differences + + # Strategy: Create logits where probabilities are approximately p1 ≈ p2 ≈ p3 ≈ 0.95/3 ≈ 0.3167 + # but with specific values chosen so that: + # - sum(p1, p2) + p3 ≠ sum(p2, p3) + p1 in f32 + # - And choose top_p to fall right at the boundary + + print("="*70) + print("Testing f32 summation order sensitivity") + print("="*70) + + # Start with equal logits + base_logit = 0.0 + + # Create three logits with small differences + # After softmax, these should be approximately equal but with rounding issues + logits = jnp.array([base_logit, base_logit + 1e-7, base_logit - 1e-7], dtype=jnp.float32) + + # Add more elements to make it realistic + logits = jnp.pad(logits, (0, 1021), constant_values=-10.0) + logits = logits.reshape(1, -1) + + print(f"\nLogits (first 3 values): {logits[0, :3]}") + + # Compute probabilities + probs = jax.nn.softmax(logits, axis=-1)[0] + print(f"Probabilities (first 3 values): {probs[:3]}") + print(f"Sum of first 3 probs: {probs[:3].sum()}") + + # Test different top_p values around the boundary + for top_p in [0.95, 0.949, 0.951, 0.9499, 0.9501]: + print(f"\n{'='*70}") + print(f"Testing with top_p = {top_p}") + print(f"{'='*70}") + + i32_result = i32_topp_mask(logits, top_p, replace_val=-1e12) + f32_result = f32_topp_mask(logits, top_p, replace_val=-1e12, stable=False) + + i32_mask = (i32_result == -1e12) + f32_mask = (f32_result == -1e12) + + i32_kept = (~i32_mask).sum() + f32_kept = (~f32_mask).sum() + + print(f" i32 implementation kept: {i32_kept} tokens") + print(f" f32 implementation kept: {f32_kept} tokens") + + if i32_kept != f32_kept: + print(f" *** DIFFERENCE FOUND! ***") + diff_indices = jnp.where(i32_mask[0] != f32_mask[0])[0] + print(f" Differing indices: {diff_indices}") + return True + + return False + + +def create_precision_sensitive_example(): + """Create a minimal example where f32 rounding causes different topp_mask results.""" + + print("\n" + "="*70) + print("Creating precision-sensitive example") + print("="*70) + + # Use the classic example: values that demonstrate f32 associativity issues + # Based on the fact that (a + b) + c ≠ a + (b + c) in floating point + + # Find three values where order matters + # Use values near the f32 precision limit + v1 = 1.0 + v2 = 1e-7 + v3 = -1e-7 + + # Test associativity + sum1 = (v1 + v2) + v3 # Left-to-right + sum2 = v1 + (v2 + v3) # Right-to-left + sum3 = (v1 + v3) + v2 # Different order + + print(f"\nDemonstrating f32 associativity issues:") + print(f" v1={v1}, v2={v2}, v3={v3}") + print(f" (v1 + v2) + v3 = {sum1}") + print(f" v1 + (v2 + v3) = {sum2}") + print(f" (v1 + v3) + v2 = {sum3}") + print(f" Difference: {abs(sum1 - sum2)}") + + # Now create logits that will expose this in topp_mask + # We want probabilities where the cumulative sum order matters + + # Strategy: Create a case where we have many small equal values + # and the cumulative sum crosses the threshold differently depending on order + + vocab_size = 1024 + + # Create logits that give us controlled probabilities + # Use a few large values and many small equal values + + # Main probability mass (3 values each with ~30% prob) + main_logits = jnp.array([10.0, 10.0 + 1e-6, 10.0 - 1e-6]) + + # Rest with very small probability + rest_logits = jnp.full(vocab_size - 3, 0.0) + + logits = jnp.concatenate([main_logits, rest_logits]).reshape(1, -1).astype(jnp.float32) + + # Compute actual probabilities + probs = jax.nn.softmax(logits, axis=-1)[0] + + print(f"\n3 main probabilities: {probs[:3]}") + print(f"Sum of 3 main probs: {probs[:3].sum():.15f}") + print(f"Sum of 2 first probs: {probs[:2].sum():.15f}") + print(f"Sum of 2 last probs: {probs[1:3].sum():.15f}") + + # Find a top_p that falls between 2 and 3 tokens + two_token_sum = probs[:2].sum() + three_token_sum = probs[:3].sum() + + print(f"\nCumulative probability analysis:") + print(f" 2 tokens: {two_token_sum:.15f}") + print(f" 3 tokens: {three_token_sum:.15f}") + + # Test with top_p values in between + for top_p in jnp.linspace(two_token_sum - 0.001, three_token_sum + 0.001, 20): + i32_result = i32_topp_mask(logits, float(top_p), replace_val=-1e12) + f32_result = f32_topp_mask(logits, float(top_p), replace_val=-1e12, stable=False) + + i32_kept = ((i32_result != -1e12).sum()) + f32_kept = ((f32_result != -1e12).sum()) + + if i32_kept != f32_kept: + print(f"\n*** FOUND DIFFERENCE at top_p={top_p:.15f} ***") + print(f" i32 kept: {i32_kept} tokens") + print(f" f32 kept: {f32_kept} tokens") + return logits, float(top_p) + + print("\nNo difference found with this example") + return None, None + + +def create_extreme_precision_example(): + """Create an extreme example with carefully crafted probabilities.""" + + print("\n" + "="*70) + print("Creating extreme precision-sensitive example") + print("="*70) + + # Use a more extreme approach: Create probabilities where the i32 scaling + # introduces quantization that changes the cumulative sum threshold + + # With i32 scaling to 2^30, we get precision of ~1/2^30 ≈ 1e-9 + # But f32 has precision of ~1e-7, so there's a gap + + scale = 2**30 + + # Create three probabilities that are affected by i32 quantization + # Each is approximately 1/3, but with specific f32 values that when + # converted to i32 and back, give slightly different sums + + # Start with a probability that doesn't divide evenly in i32 + p_approx_third = 1.0 / 3.0 + + # When scaled to i32 and back: + p_i32_quantized = int(p_approx_third * scale) / scale + + print(f"\nProbability quantization:") + print(f" f32: 1/3 = {p_approx_third:.15f}") + print(f" i32 scaled: {int(p_approx_third * scale)}") + print(f" i32 back to f32: {p_i32_quantized:.15f}") + print(f" Difference: {abs(p_approx_third - p_i32_quantized):.15e}") + + # Create logits that give us these probabilities + # Use the inverse softmax formula: logit = log(prob) + C + + # Three nearly equal probabilities that sum to slightly less than 0.95 + # This way top_p=0.95 should include all 3, but rounding might change that + + target_probs = jnp.array([0.316, 0.317, 0.317]) # Sum = 0.950 + + # Convert to logits (with arbitrary constant) + C = 10.0 + main_logits = jnp.log(target_probs) + C + + # Rest with tiny probability + vocab_size = 1024 + rest_logits = jnp.full(vocab_size - 3, -100.0) + + logits = jnp.concatenate([main_logits, rest_logits]).reshape(1, -1).astype(jnp.float32) + + # Check actual probabilities + actual_probs = jax.nn.softmax(logits, axis=-1)[0] + + print(f"\nActual probabilities:") + print(f" First 3: {actual_probs[:3]}") + print(f" Sum of first 3: {actual_probs[:3].sum():.15f}") + + # Test around 0.95 + test_values = [0.949, 0.9499, 0.95, 0.9501, 0.951] + + for top_p in test_values: + i32_result = i32_topp_mask(logits, top_p, replace_val=-1e12) + f32_result = f32_topp_mask(logits, top_p, replace_val=-1e12, stable=False) + + i32_kept = (i32_result != -1e12).sum() + f32_kept = (f32_result != -1e12).sum() + + print(f"\ntop_p={top_p:.4f}:") + print(f" i32: {i32_kept} tokens, f32: {f32_kept} tokens", end="") + + if i32_kept != f32_kept: + print(f" *** DIFFERENCE! ***") + return logits, top_p + else: + print() + + return None, None + + +def analyze_summation_approaches(): + """Analyze how the two approaches differ in their summation.""" + + print("\n" + "="*70) + print("Analyzing summation approaches") + print("="*70) + + # Create a simple example with known probabilities + probs = jnp.array([0.4, 0.3, 0.2, 0.09, 0.01], dtype=jnp.float32) + + print(f"\nTest probabilities: {probs}") + print(f"Sum: {probs.sum():.15f}") + + # Simulate f32 cumsum (what f32_topp_mask does) + f32_cumsum = jnp.cumsum(probs) + print(f"\nf32 cumulative sum: {f32_cumsum}") + + # Simulate i32 approach (what i32_topp_mask does) + scale = 2**30 + i32_scaled = (probs * scale).astype(jnp.int32) + print(f"\ni32 scaled values: {i32_scaled}") + + # Sum in i32 + i32_cumsum_scaled = jnp.cumsum(i32_scaled) + print(f"i32 cumulative sum (scaled): {i32_cumsum_scaled}") + + # Convert back to f32 + i32_cumsum_f32 = i32_cumsum_scaled.astype(jnp.float32) / scale + print(f"i32 cumsum converted to f32: {i32_cumsum_f32}") + + # Show differences + diff = f32_cumsum - i32_cumsum_f32 + print(f"\nDifference (f32 - i32): {diff}") + print(f"Max absolute difference: {jnp.abs(diff).max():.15e}") + + # Now test with top_p values at boundaries + for i, (top_p, cum_val) in enumerate(zip([0.39, 0.40, 0.41], f32_cumsum[:3])): + print(f"\n Testing top_p={top_p}:") + print(f" f32 cumsum[{i}] = {f32_cumsum[i]:.15f}") + print(f" i32 cumsum[{i}] = {i32_cumsum_f32[i]:.15f}") + + # Check which would include this token + f32_includes = f32_cumsum[i] <= top_p + i32_includes = i32_cumsum_f32[i] <= top_p + + if f32_includes != i32_includes: + print(f" *** DIFFERENT DECISIONS! f32={f32_includes}, i32={i32_includes} ***") + + +if __name__ == "__main__": + print("Demonstrating f32 vs i32 precision differences in topp_mask\n") + + # Analyze the fundamental difference + analyze_summation_approaches() + + # Try to find examples + print("\n" + "="*70) + print("Attempting to find real examples with different mask outcomes") + print("="*70) + + test_f32_summation_order_matters() + + logits1, top_p1 = create_precision_sensitive_example() + if logits1 is not None: + print(f"\n✓ Found example 1: logits shape={logits1.shape}, top_p={top_p1}") + + logits2, top_p2 = create_extreme_precision_example() + if logits2 is not None: + print(f"\n✓ Found example 2: logits shape={logits2.shape}, top_p={top_p2}") diff --git a/tests/test_topp_mask_comparison.py b/tests/test_topp_mask_comparison.py new file mode 100644 index 00000000..1c54d79e --- /dev/null +++ b/tests/test_topp_mask_comparison.py @@ -0,0 +1,122 @@ +"""Test to find differences between int dtype topp_mask and vLLM f32 topp_mask.""" + +import jax +import jax.numpy as jnp +from tallax.vllm.topp_mask import topp_mask as int_topp_mask +from tallax.vllm.tpu_inference_sampling_as_standalone_file import topp_mask as f32_topp_mask + + +def test_topp_masks_match(): + """Find cases where int dtype and f32 topp_mask give different results.""" + + # Test with various configurations to find a mismatch + configs = [ + (8, 1024, 42, 0.9), + (8, 1024, 123, 0.5), + (8, 1024, 456, 0.95), + (16, 2048, 42, 0.9), + (4, 512, 42, 0.8), + ] + + for batch_size, vocab_size, seed, p_val in configs: + key = jax.random.PRNGKey(seed) + logits_key, p_key = jax.random.split(key, 2) + + # Generate random logits + logits = jax.random.normal(logits_key, (batch_size, vocab_size)).astype(jnp.float32) + + # Test with scalar p + p_scalar = p_val + + # Apply both topp_mask implementations + int_result = int_topp_mask(logits, p_scalar, replace_val=-1e12) + f32_result = f32_topp_mask(logits, p_scalar, replace_val=-1e12, stable=False) + + # Compare masks (not values, since both replace with same value) + int_mask = (int_result == -1e12) + f32_mask = (f32_result == -1e12) + + masks_equal = jnp.array_equal(int_mask, f32_mask) + + if not masks_equal: + print(f"\n{'='*70}") + print(f"FOUND MISMATCH!") + print(f"Config: batch_size={batch_size}, vocab_size={vocab_size}, seed={seed}, p={p_val}") + print(f"{'='*70}") + + # Find which batch elements differ + diff_per_batch = (int_mask != f32_mask).sum(axis=1) + differing_batches = jnp.where(diff_per_batch > 0)[0] + + print(f"Batches with differences: {differing_batches}") + + # Focus on first differing batch + if len(differing_batches) > 0: + b = int(differing_batches[0]) + print(f"\nAnalyzing batch {b}:") + print(f" Int masked count: {int_mask[b].sum()}") + print(f" F32 masked count: {f32_mask[b].sum()}") + print(f" Difference count: {(int_mask[b] != f32_mask[b]).sum()}") + + # Find differing indices + diff_indices = jnp.where(int_mask[b] != f32_mask[b])[0] + print(f" First 10 differing indices: {diff_indices[:10]}") + + # Show the logits at those indices + for i in diff_indices[:3]: + print(f"\n Index {i}:") + print(f" Logit value: {logits[b, i]}") + print(f" Int masked: {int_mask[b, i]}") + print(f" F32 masked: {f32_mask[b, i]}") + + # Create a minimal test case + create_minimal_test_case(logits[b:b+1], p_scalar, b, seed) + + return False + + print("\nAll configs matched!") + return True + + +def create_minimal_test_case(logits, p, batch_idx, seed): + """Create a minimal reproducible test case.""" + print(f"\n{'='*70}") + print("MINIMAL TEST CASE") + print(f"{'='*70}") + + print(f""" +# Minimal test to reproduce the mismatch +import jax +import jax.numpy as jnp +from tallax.vllm.topp_mask import topp_mask as int_topp_mask +from tallax.vllm.tpu_inference_sampling_as_standalone_file import topp_mask as f32_topp_mask + +# Single batch element test +seed = {seed} +batch_idx = {batch_idx} +p = {p} +vocab_size = {logits.shape[1]} + +# Generate logits (using same seed to reproduce) +key = jax.random.PRNGKey(seed) +logits_key, _ = jax.random.split(key, 2) +all_logits = jax.random.normal(logits_key, ({logits.shape[0] + batch_idx}, vocab_size)).astype(jnp.float32) +logits = all_logits[{batch_idx}:{batch_idx+1}] # Shape (1, vocab_size) + +# Apply both implementations +int_result = int_topp_mask(logits, p, replace_val=-1e12) +f32_result = f32_topp_mask(logits, p, replace_val=-1e12, stable=False) + +# Compare +int_mask = (int_result == -1e12) +f32_mask = (f32_result == -1e12) +print(f"Masks equal: {{jnp.array_equal(int_mask, f32_mask)}}") +print(f"Int masked count: {{int_mask.sum()}}") +print(f"F32 masked count: {{f32_mask.sum()}}") +print(f"Difference count: {{(int_mask != f32_mask).sum()}}") +""") + + +if __name__ == "__main__": + print("Searching for differences between int dtype and f32 topp_mask...\n") + test_topp_masks_match() diff --git a/tests/test_topp_precision_demo.py b/tests/test_topp_precision_demo.py new file mode 100644 index 00000000..5c4a349f --- /dev/null +++ b/tests/test_topp_precision_demo.py @@ -0,0 +1,8 @@ +"""Demonstrate precision differences by exploiting softmax+cumsum ordering. + +The key difference between implementations: +- f32: softmax(logits) -> cumsum in f32 -> compare to threshold +- i32: softmax(logits) -> scale to i32 -> cumsum in i32 -> scale back -> compare + +We need logits where the softmax probabilities, when accumulated in different +precisions, cross the top_p threshold at different points. diff --git a/tests/topk_topp_mask_and_sample_test.py b/tests/topk_topp_mask_and_sample_test.py index dcc24770..0ee3ef2d 100644 --- a/tests/topk_topp_mask_and_sample_test.py +++ b/tests/topk_topp_mask_and_sample_test.py @@ -21,6 +21,14 @@ from tallax.tax.utils import is_cpu_platform +# Helper to call topk_topp_mask_and_sample with interpret=True on CPU +def _topk_topp_mask_and_sample_auto_interpret(*args, **kwargs): + """Wrapper that automatically sets interpret=True on CPU.""" + if is_cpu_platform(): + kwargs['interpret'] = True + return topk_topp_mask_and_sample(*args, **kwargs) + + def reference_topk_topp_mask_and_sample( logits: jax.Array, rng_key: jax.Array, @@ -63,8 +71,8 @@ def reference_topk_topp_mask_and_sample( logits, k ) - # Apply top-p masking - logits_masked = jax.vmap(lambda l, p_val: topp_mask(l, p_val, replace_val=replace_val, stable=stable))( + # Apply top-p masking (stable=False for topp, only topk uses stable) + logits_masked = jax.vmap(lambda l, p_val: topp_mask(l, p_val, replace_val=replace_val, stable=False))( logits_masked, p ) @@ -83,6 +91,10 @@ def reference_topk_topp_mask_and_sample( "shape", [ (8, 1024), + (8, 2048), + (8, 4096), + (8, 8192), + (8, 3570), # Non-aligned batch size (16, 2048), (24, 4096), (32, 8192), @@ -119,7 +131,7 @@ def test_topk_topp_mask_and_sample_vs_reference(shape, dtype, stable, seed): logits = jax.random.normal(logits_key, shape).astype(dtype) # Run Pallas implementation - pallas_result = topk_topp_mask_and_sample( + pallas_result = _topk_topp_mask_and_sample_auto_interpret( logits, sample_key, k, @@ -175,7 +187,7 @@ def test_topk_mask_correctness(shape, dtype, seed): logits = jax.random.normal(logits_key, shape).astype(dtype) # Run the full kernel - sampled = topk_topp_mask_and_sample( + sampled = _topk_topp_mask_and_sample_auto_interpret( logits, sample_key, k, @@ -218,7 +230,7 @@ def test_greedy_sampling_low_temperature(shape, seed): logits = jax.random.normal(logits_key, shape).astype(jnp.float32) # Run the kernel - sampled = topk_topp_mask_and_sample( + sampled = _topk_topp_mask_and_sample_auto_interpret( logits, sample_key, k, @@ -256,7 +268,7 @@ def test_padding_handling(batch_size, seed): logits = jax.random.normal(logits_key, (batch_size, vocab_size)).astype(jnp.float32) # Run the kernel - should handle padding internally - sampled = topk_topp_mask_and_sample( + sampled = _topk_topp_mask_and_sample_auto_interpret( logits, sample_key, k, @@ -288,10 +300,10 @@ def test_deterministic_sampling(seed): logits = jax.random.normal(logits_key, (batch_size, vocab_size)).astype(jnp.float32) # Run twice with same key - result1 = topk_topp_mask_and_sample( + result1 = _topk_topp_mask_and_sample_auto_interpret( logits, sample_key, k, p, temperature, stable=True, block_token=8 ) - result2 = topk_topp_mask_and_sample( + result2 = _topk_topp_mask_and_sample_auto_interpret( logits, sample_key, k, p, temperature, stable=True, block_token=8 )