Fix private memory size too large in sample_recovered_tokens_kernel

[leex404]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

fix error of RuntimeError: Triton Error [MACA]: memory size or pointer value too large to fit in 32 bit

Test Plan

Test Result

(Optional) Documentation Update

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[gemini-code-assist]

Code Review

This pull request aims to fix a memory size or pointer value too large error in a Triton kernel. The main change is the introduction of a new, blocked implementation for sample_recovered_tokens_kernel to manage memory usage better. My review focuses on ensuring the new kernel is robust against potential integer overflows during pointer arithmetic, which is a likely cause of the original error. I've identified two areas where 32-bit integer overflows could still occur and have provided suggestions to cast to 64-bit integers to prevent this, enhancing the correctness and reliability of the fix.

[gemini-code-assist]

The pointer offset calculation (start_idx + pos) * vocab_size is repeated and may suffer from 32-bit integer overflow if start_idx is loaded as a 32-bit integer and multiplied by a large vocab_size. This can lead to incorrect memory access and is a likely cause for the pointer value too large to fit in 32 bit error.

To ensure correctness and improve readability, it's better to calculate the base offset once outside the loop, explicitly casting to tl.int64 to prevent any potential overflow.

    token_idx = start_idx + pos
    # Cast to int64 to prevent overflow when calculating pointer offsets.
    base_offset = token_idx.to(tl.int64) * vocab_size

    max_prob = -float('inf')
    best_token_id = 0

    for block_start in range(0, PADDED_VOCAB_SIZE, BLOCK_SIZE):
        block_end = min(block_start + BLOCK_SIZE, vocab_size)

        vocab_offset = tl.arange(0, BLOCK_SIZE)
        mask = vocab_offset < block_end - block_start

        if NO_DRAFT_PROBS:
            draft_token_id = tl.load(draft_token_ids_ptr + token_idx)
            prob = tl.load(
                target_probs_ptr + base_offset + block_start + vocab_offset,
                mask=(mask & (vocab_offset + block_start != draft_token_id)),
                other=0)

        else:
            draft_prob = tl.load(draft_probs_ptr + base_offset + block_start +
                                 vocab_offset,
                                 mask=mask,
                                 other=0)
            target_prob = tl.load(target_probs_ptr + base_offset +
                                  block_start + vocab_offset,
                                  mask=mask,
                                  other=0)
            prob = tl.maximum(target_prob - draft_prob, 0)

[gemini-code-assist]

For consistency and to prevent potential 32-bit integer overflows, the offset calculation for q_ptr should also use 64-bit integers. While req_idx is likely small, multiplying by a large vocab_size could still pose a risk on some platforms or with very large batches. Using tl.int64 ensures the calculation is safe.

Suggested change

	q = tl.load(q_ptr + req_idx * vocab_size + block_start + vocab_offset,
	mask=mask,
	other=float("-inf"))
	q = tl.load(q_ptr + req_idx.to(tl.int64) * vocab_size + block_start + vocab_offset,
	mask=mask,
	other=float("-inf"))

[gemini-code-assist]

Code Review

This pull request aims to fix a Triton Error related to memory size and pointer values in sample_recovered_tokens_kernel. The approach of iterating over the vocabulary in blocks is a good solution for the private memory size issue. However, the fix seems incomplete as it doesn't address the potential for 32-bit integer overflow in pointer offset calculations. I've added a suggestion to explicitly use 64-bit integers for these calculations to make the fix robust. The other changes in the pull request are correct.

[gemini-code-assist]

The PR title and description mention fixing RuntimeError: Triton Error [MACA]: memory size or pointer value too large to fit in 32 bit. While iterating over the vocabulary in blocks fixes the "private memory size too large" issue, the "pointer value too large" issue can still occur.

The pointer offset calculations like (start_idx + pos) * vocab_size and req_idx * vocab_size can overflow a 32-bit integer if the total number of tokens or the vocabulary size is large, leading to the error.

This suggestion ensures 64-bit arithmetic for offset calculations by explicitly casting to tl.int64 before multiplication. This prevents potential overflows and fully resolves the reported error, making the kernel more robust.

    token_idx = start_idx + pos
    # Ensure 64-bit arithmetic for large offsets to prevent overflow.
    probs_base_offset = token_idx.to(tl.int64) * vocab_size
    q_base_offset = req_idx.to(tl.int64) * vocab_size

    max_prob = -float('inf')
    best_token_id = 0

    for block_start in range(0, PADDED_VOCAB_SIZE, BLOCK_SIZE):
        block_end = min(block_start + BLOCK_SIZE, vocab_size)

        vocab_offset = tl.arange(0, BLOCK_SIZE)
        mask = vocab_offset < block_end - block_start

        if NO_DRAFT_PROBS:
            draft_token_id = tl.load(draft_token_ids_ptr + token_idx)
            prob = tl.load(
                target_probs_ptr + probs_base_offset + block_start +
                vocab_offset,
                mask=(mask & (vocab_offset + block_start != draft_token_id)),
                other=0)

        else:
            draft_prob = tl.load(draft_probs_ptr + probs_base_offset +
                                 block_start + vocab_offset,
                                 mask=mask,
                                 other=0)
            target_prob = tl.load(target_probs_ptr + probs_base_offset +
                                  block_start + vocab_offset,
                                  mask=mask,
                                  other=0)
            prob = tl.maximum(target_prob - draft_prob, 0)

            # NOTE(woosuk): We don't need `prob = prob / tl.sum(prob)` here because
            # `tl.argmax` will select the maximum value.

        q = tl.load(q_ptr + q_base_offset + block_start + vocab_offset,
                    mask=mask,
                    other=float("-inf"))