transformer_engine/jax/dense.py · 127b6d3ab3088c403f3b38b8405b70fc33ee3f34 · OpenDAS / TransformerEngine

"git@developer.sourcefind.cn:kecinstone/2024-pra-vllm.git" did not exist on "fbd80ad4092c4bc48ce672f0435c1d1362aee052"

[JAX] Activation/Normalization to output amax for later quantization in CurrentScaling (#2238) · 127b6d3a

Phuong Nguyen authored Oct 07, 2025



* reuse amax for current scaling
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

127b6d3a

dense.py 21.3 KB

Replace dense.py