From 7f89cc5286c34debc98a77b447022387311945f8 Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Fri, 4 Sep 2026 15:00:38 +0800 Subject: [PATCH] [AMD] Skip unused TOPK v2 plan kernel on ROCm (#37580) Co-authored-by: kk <43161300+kkHuang-amd@users.noreply.github.com> --- python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh index 656b9ca40..b09a15364 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh @@ -467,8 +467,13 @@ struct TopKKernel { .with_device(device_) .verify(metadata); - const auto batch_size = static_cast(B.unwrap()); RuntimeCheck(Bp1.unwrap() == B.unwrap() + 1, "invalid metadata shape"); +#ifdef USE_ROCM + // ROCm compiles out the cluster path, the only consumer of this plan. + (void)static_cluster_threshold; + return; +#else + const auto batch_size = static_cast(B.unwrap()); const auto device = device_.unwrap(); LaunchKernel(1, kBlockSize, device)( // topk_plan, @@ -476,6 +481,7 @@ struct TopKKernel { static_cast(metadata.data_ptr()), batch_size, static_cluster_threshold); +#endif } static void transform_paged(