Unverified Commit e290303e authored by Lianmin Zheng's avatar Lianmin Zheng Committed by GitHub
Browse files

[Auto Sync] Update elementwise.py (20250923) (#10823)


Co-authored-by: default avatargithub-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: default avatarCheng Wan <54331508+ch-wan@users.noreply.github.com>
parent aab35bcc
...@@ -187,7 +187,9 @@ fused_dual_residual_rmsnorm_kernel_autotune = rmsnorm_autotune( ...@@ -187,7 +187,9 @@ fused_dual_residual_rmsnorm_kernel_autotune = rmsnorm_autotune(
def fused_dual_residual_rmsnorm(x, residual, weight1, weight2, eps, autotune=False): def fused_dual_residual_rmsnorm(x, residual, weight1, weight2, eps, autotune=False):
assert len(x.shape) == 2 assert len(x.shape) == 2
assert x.shape == residual.shape and x.dtype == residual.dtype assert (
x.shape == residual.shape and x.dtype == residual.dtype
), f"{x.shape=} {residual.shape=} {x.dtype=} {residual.dtype=}"
output, mid = torch.empty_like(x), torch.empty_like(x) output, mid = torch.empty_like(x), torch.empty_like(x)
bs, hidden_dim = x.shape bs, hidden_dim = x.shape
if autotune: if autotune:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment