Describe the bug
Linear, LayernormLinear, LayernormMLPandGroupedLinearunconditionally setweight_quantizer.optimize_for_gemm=True` for non-primary quantized weights. This is unsafe when weight caching and CUDA graphs are enabled but the quantization kernel cannot emit GEMM-swizzled scales directly, for example:
NVFP4 with unaligned weight shapes
NVFP4 with 2D weight quantization disabled
Any MXFP8/NVFP4 path that falls back to a standalone swizzle
The standalone swizzle reallocates and rebinds the cached scale-factor tensor. A CUDA graph captured with is_first_microbatch=False retains the old scale pointer; a later eager is_first_microbatch=True cache update can replace/free it, causing corruption or an illegal memory access during replay.
Proposed fix: Enable weight optimize_for_gemm only when the selected quantization kernel supports direct swizzle fusion for the weight shape/configuration. Otherwise, keep compact cached scales and let GEMM perform the lazy swizzle.
This currently affects TE modules; the TE op-fuser grouped MLP does not enable weight optimize_for_gemm or maintain an is_first_microbatch weight cache so it's not affected.
Steps/Code to reproduce bug
Please list minimal steps or code snippet for us to be able to reproduce the bug.
A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.
Expected behavior
A clear and concise description of what you expected to happen.
Environment overview (please complete the following information)
- Environment location: [Bare-metal, Docker, Cloud(specify cloud provider - AWS, Azure, GCP, Collab)]
- Method of Transformer Engine install: [pip install or from source]. Please specify exact commands you used to install.
- If method of install is [Docker], provide
docker pull & docker run commands used
Environment details
If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
- OS version
- PyTorch version
- Python version
- Transformer Engine version
- CUDA version
- CUDNN version
Device details
Additional context
Add any other context about the problem here.
Describe the bug
Linear, LayernormLinear,LayernormMLPandGroupedLinearunconditionally setweight_quantizer.optimize_for_gemm=True` for non-primary quantized weights. This is unsafe when weight caching and CUDA graphs are enabled but the quantization kernel cannot emit GEMM-swizzled scales directly, for example:The standalone swizzle reallocates and rebinds the cached scale-factor tensor. A CUDA graph captured with
is_first_microbatch=Falseretains the old scale pointer; a later eageris_first_microbatch=Truecache update can replace/free it, causing corruption or an illegal memory access during replay.Proposed fix: Enable weight
optimize_for_gemmonly when the selected quantization kernel supports direct swizzle fusion for the weight shape/configuration. Otherwise, keep compact cached scales and let GEMM perform the lazy swizzle.This currently affects TE modules; the TE op-fuser grouped MLP does not enable weight
optimize_for_gemmor maintain anis_first_microbatchweight cache so it's not affected.Steps/Code to reproduce bug
Please list minimal steps or code snippet for us to be able to reproduce the bug.
A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.
Expected behavior
A clear and concise description of what you expected to happen.
Environment overview (please complete the following information)
docker pull&docker runcommands usedEnvironment details
If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
Device details
Additional context
Add any other context about the problem here.