Fix apply_top_k_top_p_triton called by non-cuda logits Tensor

[xli]

Summary:
#33538 add apply_top_k_top_p_triton function; inside apply_top_k_top_p_triton, there is assert logits.is_cuda to ensure input is cuda logits tensor, so we should guard the call with logits.is_cuda

[dosubot]

Related Documentation

Checked 0 published document(s) in 1 knowledge base(s). No updates required.

^{How did I do? Any feedback?}  

[gemini-code-assist]

Code Review

This pull request correctly fixes a bug where the Triton-based apply_top_k_top_p_triton kernel could be called with a non-CUDA tensor, which would lead to an assertion failure. The fix introduces a logits.is_cuda check, ensuring that the optimized Triton kernel is only used for CUDA tensors. For non-CUDA tensors or small batch sizes, the execution correctly falls back to the more general PyTorch implementation. This change is a necessary and well-implemented bug fix that improves the robustness of the sampling logic.

[WoosukKwon]

@xli Thanks for the PR.

@njhill I actually think we can fix this by removing assert logits.is_cuda (plus others like k.is_cuda). WDYT?

[xli]

@xli Thanks for the PR.

@njhill I actually think we can fix this by removing assert logits.is_cuda (plus others like k.is_cuda). WDYT?

Removing assertion does not work in my local test with our internal hardware test, so I kept this diff simple.

[njhill]

@xli Thanks for the PR.
@njhill I actually think we can fix this by removing assert logits.is_cuda (plus others like k.is_cuda). WDYT?

Removing assertion does not work in my local test with our internal hardware test, so I kept this diff simple.

Thanks @xli and apologies for not catching this before. I agree we should remove the is_cudas, could you give more information about the failure that you encountered? Was it related to torch.cuda.get_device_properties? Ideally we should fix the HAS_TRITON logic to exclude your case if needed rather than the additional check here.

[houseroad]

oh, I didn't see @njhill 's comments. @xli , could you address @njhill 's comments as follow?

[njhill]

Please see also parallel discussion in #35011, @jikunshang is also working on this.

[xli]

@xli Thanks for the PR.
@njhill I actually think we can fix this by removing assert logits.is_cuda (plus others like k.is_cuda). WDYT?

Removing assertion does not work in my local test with our internal hardware test, so I kept this diff simple.

Thanks @xli and apologies for not catching this before. I agree we should remove the is_cudas, could you give more information about the failure that you encountered? Was it related to torch.cuda.get_device_properties? Ideally we should fix the HAS_TRITON logic to exclude your case if needed rather than the additional check here.

Hi, I just found my previous internal hardware test was running on a nvidia gpu host (applying this pr change does fix the test), the error is related to get_device_properties, let me borrow hardware to redo test, will report back once I got it.

[jikunshang]

@xli please check whether this also works for you, thanks! #35042