fix

shihaobai · shihaobai · commit b0f6332ef6e8 · 2025-09-02T18:43:53.000+08:00
diff --git a/lightllm/models/llama/triton_kernel/rmsnorm.py b/lightllm/models/llama/triton_kernel/rmsnorm.py
@@ -53,8 +53,8 @@ def get_test_configs():
     ]
 
 
-def get_static_key(x, out):
-    return {"N": x.shape[-1], "out_dtype": str(out.dtype)}
+def get_static_key(x):
+    return {"N": x.shape[-1], "out_dtype": str(x.dtype)}
 
 
 @autotune(

Original file line number	Diff line number	Diff line change
`@@ -53,8 +53,8 @@ def get_test_configs():`
`53`	`53`	`]`
`54`	`54`
`55`	`55`
`56`		`-def get_static_key(x, out):`
`57`		`- return {"N": x.shape[-1], "out_dtype": str(out.dtype)}`
	`56`	`+def get_static_key(x):`
	`57`	`+ return {"N": x.shape[-1], "out_dtype": str(x.dtype)}`
`58`	`58`
`59`	`59`
`60`	`60`	`@autotune(`