[CPU] Skip Triton kernel monkey-patches when Triton-CPU is available#44991
Conversation
When Triton-CPU backend is active (HAS_TRITON=True), let the native Triton @jit kernels run directly on CPU instead of replacing them with C++ fallback implementations. The C++ monkey-patches remain as fallback for environments without Triton-CPU. Signed-off-by: jmamou <[email protected]>
Triton-CPU's compiler is stricter than CUDA Triton about type consistency across if/else branches. The pattern `start_idx = 0 if cond else tl.load(ptr)` produces int32 in the then-branch and int64 in the else-branch, causing a CompilationError. Fix by using `tl.zeros([], dtype=tl.int64)` instead of the literal `0` so both branches produce int64. This is backward-compatible with CUDA Triton (which handles implicit widening). Kernels fixed: - expand_kernel - rejection_greedy_sample_kernel - rejection_random_sample_kernel - sample_recovered_tokens_kernel - eagle_prepare_inputs_padded_kernel Signed-off-by: jmamou <[email protected]>
Verify that the fixed Triton kernels (expand_kernel, rejection_greedy_sample_kernel, rejection_random_sample_kernel, sample_recovered_tokens_kernel, eagle_prepare_inputs_padded_kernel) compile and produce correct results on CPU with Triton-CPU. Tests are skipped when HAS_TRITON is False (no CPU Triton backend). Signed-off-by: jmamou <[email protected]>
Add tests for eagle_step_slot_mapping_metadata_kernel, eagle_prepare_next_token_padded_kernel, and copy_and_expand_eagle_inputs_kernel on Triton-CPU. All 8 monkey-patched kernels are now verified to compile and run correctly on Triton-CPU. Signed-off-by: jmamou <[email protected]>
bigPYJ1151
left a comment
There was a problem hiding this comment.
Hi @jmamou thanks for the fix :) I also opened a issue on Triton-CPU side triton-lang/triton-cpu#262
I think the test is not required for now. After the spec decoding on CPU become stable we can reuse some existing tests in CI.
Removing the Triton CPU kernel unit tests as requested by reviewer. Signed-off-by: jmamou <[email protected]>
The cu_num_draft_tokens tensor is int32, so tl.zeros with int64 creates a type mismatch between the if/else branches. Use ptr.dtype.element_ty to always match the pointer's actual dtype. Signed-off-by: jmamou <[email protected]>
Head branch was pushed to by a user without write access
…llm-project#44991) Signed-off-by: jmamou <[email protected]> Co-authored-by: Li, Jiang <[email protected]>
…llm-project#44991) Signed-off-by: jmamou <[email protected]> Co-authored-by: Li, Jiang <[email protected]> Signed-off-by: divineearthly <[email protected]>
…llm-project#44991) Signed-off-by: jmamou <[email protected]> Co-authored-by: Li, Jiang <[email protected]>
|
Hi @jmamou , according to the discussion triton-lang/triton-cpu#262 (comment) with Triton CPU developer the error about type mismatch is as expected. |
|
Hi @bigPYJ1151, thanks for digging into this with the Triton-CPU team. Looking again at The second part of the PR is still relevant: skipping the C++ monkey-patches when |
…llm-project#44991) Signed-off-by: jmamou <[email protected]> Co-authored-by: Li, Jiang <[email protected]>
…llm-project#44991) Signed-off-by: jmamou <[email protected]> Co-authored-by: Li, Jiang <[email protected]>
…llm-project#44991) Signed-off-by: jmamou <[email protected]> Co-authored-by: Li, Jiang <[email protected]>
…llm-project#44991) Signed-off-by: jmamou <[email protected]> Co-authored-by: Li, Jiang <[email protected]>
Purpose
When Triton-CPU is installed and active (the default in x86 Docker images since 0.22.0), the
@triton.jitkernels inrejection_sampler.py,block_table.py, andllm_base_proposer.pycan run natively on CPU — but they fail to compile due to int32/int64 type mismatches that Triton-CPU's stricter compiler catches.This PR:
0(int32) withtl.zeros([], dtype=tl.int64)in conditional branches where the else-path produces int64 fromtl.load(). This is backward-compatible with CUDA Triton.HAS_TRITONis True (i.e., Triton-CPU backend is detected), letting native Triton kernels execute directly. The C++ fallback implementations remain for environments without Triton-CPU.Kernels fixed
expand_kernelrejection_greedy_sample_kernelrejection_random_sample_kernelsample_recovered_tokens_kerneleagle_prepare_inputs_padded_kernelRoot cause
Triton-CPU requires type consistency across if/else branches:
Ref: #42638 (comment)
Test plan
vllm/vllm-openai-cpu:v0.22.0Docker image)_compute_slot_mapping_kernelalready works on Triton-CPU (no fix needed)