recover triton kernel name in ctx.call_kernel

staugust · staugust · commit 744554bd2a58 · 2025-11-28T10:56:32.000+08:00
Signed-off-by: augusto.yjh &lt;augusto.yjh@antgroup.com&gt;
diff --git a/vllm/attention/ops/common.py b/vllm/attention/ops/common.py
@@ -177,7 +177,7 @@ def correct_attn_out(
         cp_rank,
     )
     const_args = {"HEAD_DIM": D, "N_ROUNDED": N, "IS_BASE_E": is_lse_base_on_e}
-    ctx.call_kernel(correct_attn_kernel, grid, *regular_args, **const_args)
+    ctx.call_kernel(_correct_attn_cp_out_kernel, grid, *regular_args, **const_args)
     return out, lse
 
 

Original file line number	Diff line number	Diff line change
`@@ -177,7 +177,7 @@ def correct_attn_out(`
`177`	`177`	`cp_rank,`
`178`	`178`	`)`
`179`	`179`	`const_args = {"HEAD_DIM": D, "N_ROUNDED": N, "IS_BASE_E": is_lse_base_on_e}`
`180`		`- ctx.call_kernel(correct_attn_kernel, grid, regular_args, *const_args)`
	`180`	`+ ctx.call_kernel(_correct_attn_cp_out_kernel, grid, regular_args, *const_args)`
`181`	`181`	`return out, lse`
`182`	`182`
`183`	`183`