[Perf] Replace cudaMemsetAsync with in-kernel cleanup for persistent_topk

[LopezCastroRoberto]

Motivation

Replaces the per-call cudaMemsetAsync (PR #41444 and #41665) with an in-kernel cleanup at the end of persistent_topk_kernel, eliminating a 3-5us overhead per launch (decode step).

This is particularly relevant in low-latency scenarios and when max_seq_len <= RADIX_THRESHOLD, where the workspace is not required. Despite this, the persistent kernel currently invokes cudaMemsetAsync unconditionally, incurring unnecessary overhead.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

Hi @LopezCastroRoberto, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[gemini-code-assist]

Code Review

This pull request moves the zero-initialization of the RadixRowState workspace from a host-side cudaMemsetAsync to an in-kernel cleanup at the end of the persistent TopK kernel. This change aims to improve performance and support CUDA graph replays by avoiding stream-ordered host calls. However, the current in-kernel implementation introduces a performance bottleneck due to excessive memory fences in a loop and a race condition where the arrival_counter could be reset before all histogram zeros are globally visible. I have provided feedback on how to parallelize the zeroing and use a proper synchronization barrier.

[LopezCastroRoberto]

@claude review

[LopezCastroRoberto]

cc: @zyongye