[kainzhong]

Description

We need fail cases where we expect CuTeDSL kernels are used to avoid silent fallbacks like #3743 catches. This also introduces a breaking change that NVTE_WARN_IF_CUTEDSL_BACKEND_NOT_CHOSEN is changed to NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN which can warn or fail the test.

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

Please list the changes introduced in this PR:

• Report CuTeDSL kernels falling back with a warn or fail option based on the env var
• 2D quantization is currently unsupported so for this it won't trigger warn or fail and I added a TODO to remind me of adding this kernel later
• Add cuda-bindings in setup.py since it's used by quantize_mxfp8.py and fused_mla_q_uproj
• Make CI scripts fail if falling back to CUDA

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]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no new actionable issues were established.

Summary

Adds opt-in warnings or errors when an enabled CuTeDSL backend falls back to CUDA.

• CuTeDSL misses can now warn or stop instead of silently falling back.
• C++, JAX, and PyTorch test runs stop when CuTeDSL fallback is reported.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[MXFP8 request] --> B{CuTeDSL enabled?}
  B -->|No| C[Use CUDA]
  B -->|Yes| D{2D quantization?}
  D -->|Yes| C
  D -->|No| E{CuTeDSL kernel available?}
  E -->|Yes| F[Use CuTeDSL]
  E -->|No| G{Report setting}
  G -->|Unset| C
  G -->|WARN| H[Warn and use CUDA]
  G -->|ERROR| I[Raise an error]

Loading

_{Reviews (6) · Last reviewed commit: "Merge branch 'main' into warn_or_fail_wh..." · Reviewed by Greptile}

[kainzhong]

/te-ci

[kainzhong]

/te-ci

[ksivaman]

This is a breaking change due to the envvar change, so we should mark that in the description box, which would be later useful for relnotes

[ksivaman]

Why?

[kainzhong]

It's used by quantize_mxfp8.py and fused_mla_q_uproj.py. nvidia-cutlass-dsl should always depends on this but I think it's clearer to explicitly mark this as a dependency

[ksivaman]

In that case we should also mark nvidia-cutlass-dsl and tvi-ffi as required dependencies. Currently those are user installed, which is fine because this is an optional backend. So for consistency either all or none should be included.

[kainzhong]

Actually fused_mla_q_uproj.py uses cuda-bindings for passing the stream to cudnn-fe's kernel so I guess we can make this a dependency for all. For nvidia-cutlass-dsl and tvm-ffi, you can use CUDA kernels if you don't install these so I think we should make them dependency when we move to 3.0 and get rid of CUDAS kernels?

[kainzhong]

OK I removed this

[ksivaman]

LGtM

[ksivaman]

/te-ci

[kainzhong]

Actually we should merge #3743 first

[kainzhong]

/te-ci

[kainzhong]

/te-ci