Skip to content

Commit d40b187

Browse files
authored
[BugFix][Ops] chunk_scaled_dot_kkt_fwd_kernel recompile issues (vllm-project#11577)
### What this PR does / why we need it? This PR fixes recompilation issues in the `chunk_scaled_dot_kkt_fwd_kernel` Triton kernel. The parameters `bh_step`, `task_num`, and `num_core` were previously defined as `tl.constexpr`, which caused Triton to recompile the kernel whenever these values changed. By changing them to regular arguments, we avoid recompilation when these values change. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? - vLLM version: v0.23.0 - vLLM main: vllm-project/vllm@1f486d9 --------- Signed-off-by: Ascendyh <hw7osiris@outlook.com>
1 parent 0532c49 commit d40b187

1 file changed

Lines changed: 4 additions & 4 deletions

File tree

vllm_ascend/ops/triton/fla/chunk_scaled_dot_kkt.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
"USE_G": lambda args: args["g_cumsum"] is not None,
2424
}
2525
)
26-
@triton.jit(do_not_specialize=["T", "B"])
26+
@triton.jit(do_not_specialize=["T", "B", "bh_step", "task_num", "num_core"])
2727
def chunk_scaled_dot_kkt_fwd_kernel(
2828
k,
2929
beta, # [H, B, T]
@@ -33,16 +33,16 @@ def chunk_scaled_dot_kkt_fwd_kernel(
3333
chunk_indices,
3434
T,
3535
B,
36+
bh_step,
37+
task_num,
38+
num_core,
3639
H: tl.constexpr,
3740
Hg: tl.constexpr,
3841
K: tl.constexpr,
3942
BT: tl.constexpr,
4043
BK: tl.constexpr,
4144
IS_VARLEN: tl.constexpr,
4245
USE_G: tl.constexpr,
43-
bh_step: tl.constexpr,
44-
task_num: tl.constexpr,
45-
num_core: tl.constexpr,
4646
):
4747
bt_stride = B * T
4848
core_id = tl.program_id(0)

0 commit comments

Comments
 (0)