[Common] Use specialized unfused MXFP8 cast kernels by default

[Oleg-Goncharov]

Description

This PR enables the fast unfused MXFP8 cast kernels by default.

Previously, these kernels were gated behind an environment variable and therefore were not used unless explicitly enabled. This change makes the specialized cast-only path the default behavior.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Removed environment variable

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR promotes the specialized unfused MXFP8 cast kernels from opt-in (via ENABLE_CAST_ONLY env var) to the default path, while adding the runtime eligibility guards that make the promotion safe.

• specialized/quantize_mxfp8.cuh: Removes is_cast_only_enabled() and the ENABLE_CAST_ONLY environment variable; all four hasSpec specializations now return true unconditionally. The commented-out FIXME debugging block is also cleaned up.
• quantize_mxfp8.cuh: Adds scaling_type_has_specialized_support — guarding ROWWISE on cols % 128 == 0 (to prevent out-of-bounds reads on partial 128-element row tails) and both ROWWISE/BIDIMENSIONAL on grid-dim-Y ≤ 65535 — so the fast path is taken only when it is correct to do so. Removes the previously dead COLWISE case from the specialized switch (addressed from a prior review), wraps cudaFuncSetAttribute calls with NVTE_CHECK_CUDA, and consolidates per-case cudaGetLastError checks into a single post-switch call in the generic path.

Confidence Score: 5/5

Safe to merge — the specialized kernels are guarded by correct runtime eligibility checks before being invoked, and error detection is strengthened throughout.

The behavioral change is narrow and well-defended: the cols % 128 == 0 guard prevents the known out-of-bounds access on partial row tails, the grid-dim-Y checks prevent exceeding CUDA launch limits, and COLWISE is correctly excluded from the specialized path. The generic fallback is preserved for all shapes that do not meet the criteria. cudaFuncSetAttribute and cudaGetLastError are now properly checked. The only finding is a formatting nit where the error-check call shares a line with the switch's closing brace.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh	Removes the is_cast_only_enabled() env-var gate and makes all four hasSpec specializations return true unconditionally, enabling the fast cast-only path by default.
transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh	Adds runtime eligibility guards (cols % 128 == 0, grid-dim-Y limits) before taking the specialized path; removes the now-unreachable COLWISE dead case; wraps cudaFuncSetAttribute and kernel launches with NVTE_CHECK_CUDA; consolidates per-case cudaGetLastError checks into one post-switch call — minor style concern on line 839 where the check sits on the same line as the closing brace.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[quantize called] --> B{hasSpec AND\nnot swizzled scales?}
    B -- No --> G[Generic kernel path]
    B -- Yes --> C{scaling_type_has_specialized_support?}
    C -- No --> G
    C -- Yes --> D{scaling_type?}
    D -- ROWWISE\ncols%128==0 AND\ngrid fits --> E[specialized rowwise\ncast-only kernel]
    D -- BIDIMENSIONAL\ngrid fits --> F[specialized bidimensional\ncast-only kernel]
    E --> CUDA_CHECK[NVTE_CHECK_CUDA\ncudaGetLastError]
    F --> CUDA_CHECK
    CUDA_CHECK --> RETURN[return]
    G --> SW{scaling_type?}
    SW -- ROWWISE --> GR[generic ROWWISE kernel]
    SW -- COLWISE --> GC[generic COLWISE kernel]
    SW -- BIDIMENSIONAL --> GB[generic BIDIMENSIONAL kernel]
    GR & GC & GB --> GE[NVTE_CHECK_CUDA\ncudaGetLastError]

Loading

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into pr_fast_default..." | Re-trigger Greptile}

[ksivaman]

LGTM

[Oleg-Goncharov]

/te-ci

[Oleg-Goncharov]

/te-ci

[Oleg-Goncharov]

/te-ci

[Oleg-Goncharov]

/te-ci

[ptrendx]

For the future work we could think about doing the swizzling support for that kernel, but not sure how needed it really is.

[ptrendx]

/te-ci