fix layernorm import (Dao-AILab#92)

zhuochenKIDD · web-flow · commit 171b15bcd2fe · 2026-04-01T00:00:54.000-04:00
diff --git a/benchmarks/benchmark_layernorm.py b/benchmarks/benchmark_layernorm.py
@@ -8,7 +8,7 @@
 import cutlass
 import cutlass.torch as cutlass_torch
 from cutlass.cute.runtime import from_dlpack
-from quack.layernorm import layernorm, layernorm_ref, rstd_ref, mean_ref
+from quack.rmsnorm import layernorm_fwd as layernorm, layernorm_ref
 import cutlass.cute as cute
 
 try:
