diff --git a/src/target/source/codegen_webgpu.cc b/src/target/source/codegen_webgpu.cc
index 32b32063940b..8c34c6144fe3 100644
--- a/src/target/source/codegen_webgpu.cc
+++ b/src/target/source/codegen_webgpu.cc
@@ -185,10 +185,9 @@ runtime::FunctionInfo CodeGenWebGPU::AddFunction(const PrimFunc& f, bool skip_re
   }
 
   // setup thread tags and param access in launch param tags;
-  if (auto opt = f->GetAttr<Array<tir::IterVar>>(tir::attr::kDeviceThreadAxis)) {
-    auto thread_axis = opt.value();
-    for (size_t i = 0; i < thread_axis.size(); ++i) {
-      func_info.launch_param_tags.push_back(thread_axis[i]->thread_tag);
+  if (auto opt = f->GetAttr<Array<String>>(tir::attr::kKernelLaunchParams)) {
+    for (const auto& thread_tag : opt.value()) {
+      func_info.launch_param_tags.push_back(thread_tag);
     }
   }
   os_param_access << "]";