Refactor debug output to return full intermediate arrays instead of checksums

[oliverdutton]

Summary

This PR refactors the debug output mechanism in the top-k/top-p sampling kernels to return full intermediate arrays instead of checksums or partial values. This enables more comprehensive debugging and validation of the sampling pipeline.

Key Changes

Arbitrary-k Kernel (tallax/vllm/arbitrary_k/kernel.py)

• Renamed debug_results_ref parameter to debug_arrays_ref to clarify that full arrays are returned
• Replaced the _write_debug_results() helper function that computed checksums with direct array assignments
• Changed debug output from SMEM-based checksums to full arrays:
  • greedy_sampled: Full [batch, 1] array instead of [1, 1] checksum
  • topk_logits_unsorted: Full [batch, vocab_size] array (new)
  • topk_topp_unnorm_probs_i32_unsorted: Full [batch, vocab_size] array instead of nonzero count
  • random_unnorm_cdf_sampled: Full [batch, 1] array (new, extracted as int64)
  • next_tokens: Full [batch, 1] array instead of [1, 1] checksum
• Updated BlockSpec definitions to use proper memory layout for full arrays instead of SMEM
• Added trimming of padded outputs to match original batch size in debug results

Bounded-k Kernel (tallax/vllm/bounded_k/top_p_and_sample.py)

• Added debug parameter to top_p_and_sample_arrays() and related functions
• Refactored to store intermediate values before transposition for proper debug output
• Added greedy_sampled computation before temperature scaling
• Returns (tokens, debug_results) tuple when debug=True
• Debug output includes full arrays matching the arbitrary-k format:
  • greedy_sampled: [batch_size]
  • topk_logits_unsorted: [batch_size, k]
  • topk_topp_unnorm_probs_i32_unsorted: [batch_size, k]
  • random_unnorm_cdf_sampled: [batch_size]
  • next_tokens: [batch_size]

Reference Implementation (tallax/vllm/reference.py)

• Updated debug output to use OrderedDict for consistent ordering
• Renamed debug keys to match kernel implementation:
  • topp_unnorm_probs_i32 → topk_topp_unnorm_probs_i32_unsorted
  • topp_nonzero_count → removed (now returning full array)
  • total_sum → removed (now returning random_unnorm_cdf_sampled)

Tests (tests/vllm/arbitrary_k/kernel_test.py)

• Updated assertions to work with full arrays instead of checksums
• Added new assertions for:
  • topk_logits_unsorted: Validates full top-k logits match
  • topk_topp_unnorm_probs_i32_unsorted: Validates full unnormalized probabilities
  • random_unnorm_cdf_sampled: Validates the sampled CDF value
• Removed assertions for checksums that are no longer computed

Implementation Details

• Uses OrderedDict to maintain consistent ordering of debug outputs across implementations
• Full arrays are now stored in HBM instead of SMEM, enabling larger intermediate values
• Debug output shapes match the full computation shapes for easier validation
• Padding is trimmed from debug outputs to match the original batch size

https://claude.ai/code/session_01NJso6RoxkCzL1d3hSTnU5R