Motivation
Diffusion and multimodal transformer models commonly execute closely related memory-bound patterns that go beyond standalone LayerNorm/RMSNorm:
updated = residual + gate * update
normalized = layer_norm_or_rms_norm(updated, weight, bias, eps)
modulated = normalized * (1 + scale) + shift
return modulated, updated
Depending on the model, the normalization may be LayerNorm or RMSNorm, affine or non-affine, and gate/scale/shift may use broadcast layouts such as [B, 1, D], [B, T, D], or [B, T, 1, 1, D].
AutoKernel already supports standalone LayerNorm and RMSNorm, so it looks like a promising foundation for profiling and autonomously tuning these fused transformer operations.
Requested support
Would you be open to supporting one or both of the following?
- Built-in kernel types for:
- modulated LayerNorm / RMSNorm:
norm(x) * (1 + scale) + shift
- gated residual:
residual + gate * update
- combined gated-residual + modulated normalization, optionally returning the updated residual
- A generic custom-kernel extension interface where users can provide:
- a PyTorch reference function
- input shape/dtype/broadcast generators
- correctness tolerances and edge cases
- a model replacement/integration hook
- optional forward and backward verification
Useful validation would include fp16/bf16 with fp32 accumulation, arbitrary leading dimensions, multiple broadcast layouts, deterministic forward/backward parity, and compatibility checks with torch.compile.
Contribution
I’d be happy to contribute and open a PR if I’m able to implement this in a way that fits AutoKernel’s design. Guidance on whether you would prefer dedicated kernel types or a general custom-operation/plugin path would be very helpful.
Motivation
Diffusion and multimodal transformer models commonly execute closely related memory-bound patterns that go beyond standalone LayerNorm/RMSNorm:
Depending on the model, the normalization may be LayerNorm or RMSNorm, affine or non-affine, and
gate/scale/shiftmay use broadcast layouts such as[B, 1, D],[B, T, D], or[B, T, 1, 1, D].AutoKernel already supports standalone LayerNorm and RMSNorm, so it looks like a promising foundation for profiling and autonomously tuning these fused transformer operations.
Requested support
Would you be open to supporting one or both of the following?
norm(x) * (1 + scale) + shiftresidual + gate * updateUseful validation would include fp16/bf16 with fp32 accumulation, arbitrary leading dimensions, multiple broadcast layouts, deterministic forward/backward parity, and compatibility checks with
torch.compile.Contribution
I’d be happy to contribute and open a PR if I’m able to implement this in a way that fits AutoKernel’s design. Guidance on whether you would prefer dedicated kernel types or a general custom-operation/plugin path would be very helpful.