Add Snake activation functor for EVT#3184
Open
emre570 wants to merge 1 commit intoNVIDIA:mainfrom
Open
Conversation
Introduces cutlass::epilogue::thread::Snake, a two-operand activation functor implementing Snake_a(x) = x + (1/a) * sin^2(a*x) from Ziyin et al. 2020 (arXiv:2006.08195). The per-channel learnable frequency `a` flows through an EVT child (e.g. Sm90RowBroadcast), composing into Sm90EVT<Sm90Compute<Snake, ...>, x_node, alpha_node> for fused GEMM+Snake epilogues used in neural vocoders. Adds unit tests in test/unit/epilogue/thread/activation.cu covering f32 and bf16 paths, validated against float64 reference goldens. Closes NVIDIA#3141
Author
|
@hwu36 Hi sir, is it possible to take a look in your free time? |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
cutlass::epilogue::thread::Snake, a two-operand activation functorimplementing the Snake activation from Ziyin, Hartwig, & Ueda
(NeurIPS 2020, arXiv:2006.08195):
The per-channel learnable frequency
ais passed as a second operand,flowing through an EVT child (typically
Sm90RowBroadcast). This enablesfused GEMM/Conv + Snake epilogues for neural vocoders, where Snake is the
standard activation in the WaveGen path. See issue #3141 for production
benchmarks (~2.1× median speedup on H100 vocoder shapes).
Implementation
include/cutlass/epilogue/thread/activation.h— scalarSnake<T>andSnake<Array<T, N>>specializations, usingcutlass::fast_sinforIEEE-precise device/host math.
Sm90Computeprimary; no newinfrastructure needed.
CUTLASS_HOST_DEVICE,kIsHeavy = true,CUTLASS_PRAGMA_UNROLLon the Array specialization.Tests
Adds unit tests in
test/unit/epilogue/thread/activation.cu:Epilogue_thread_snake.device_f32— tolerance 1e-5Epilogue_thread_snake.device_bf16— tolerance 2e-2Both pass on SM90 (H100). Goldens are generated from a
math.sin-basedfloat64 reference over 256 samples each of x, α, and expected output,
with α ∈ (0.1, 2.0) and x ~ N(0, 1).
Usage
Closes #3141