fix(perf): recognize RTX PRO 6000 for MFU

[ralovets]

What does this PR do ?

Register the RTX PRO 6000 Workstation and Max-Q dense BF16 peak throughput values so automatic device detection reports MFU with the correct denominator for each variant.

Problem

AutoMFU detects the current CUDA device to select the theoretical peak used by MFU logging. Neither RTX PRO 6000 workstation variant is currently present in the device table, so their full CUDA device names are treated as unknown. A shared RTX PRO 6000 key is not sufficient because substring matching would assign the 503.8 TFLOP/s Workstation peak to the 438.9 TFLOP/s Max-Q variant.

NVIDIA RTX Blackwell GPU Architecture v1.1, Table 4 lists dense BF16 Tensor Core peaks of 503.8 TFLOP/s for the Workstation Edition and 438.9 TFLOP/s for the Max-Q Workstation Edition. The device table uses these dense values.

Validation

RTX PRO 6000 GPU smoke test

Ran the clean-main and patched MFU modules on GPU 0 of an RTX PRO 6000 Blackwell workstation. With fixed calculation inputs of 100 total TFLOPs, two GPUs, and one second:

device=NVIDIA RTX PRO 6000 Blackwell Workstation Edition
before: peak_tflops=inf, mfu=0.0
after: peak_tflops=503.8, mfu=9.924573243350537

This smoke test exercises CUDA device-name detection and the resulting MFU denominator; it does not run a model training workload.

After narrowing the lookup key, the same physical GPU still resolves correctly:

device=NVIDIA RTX PRO 6000 Blackwell Workstation Edition
peak_tflops=503.8

Max-Q hardware was not available. A CPU regression test covers both expected full CUDA device strings and verifies that Max-Q resolves to 438.9 while the Workstation Edition resolves to 503.8.

Unit and style checks

uv run pytest tests/unit_tests/utils/test_flops_utils.py -q
uv run ruff check nemo_automodel/_transformers/mfu.py tests/unit_tests/utils/test_flops_utils.py
uv run ruff format --check nemo_automodel/_transformers/mfu.py tests/unit_tests/utils/test_flops_utils.py

Changelog

• Add distinct RTX PRO 6000 Workstation and Max-Q entries to _DEVICE_FLOPS.
• Add regression coverage for both full device names.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests? Added focused coverage for the two RTX PRO 6000 variants and ran a targeted GPU smoke test.
• Did you add or update any necessary documentation? No repository documentation changes are needed; the NVIDIA specification is linked above.

Additional Information

• No related issue.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yuhezhang-ai]

Thanks for adding this support and for including the hardware smoke test—very helpful!

One small concern: RTX PRO 6000 is also a substring of the Max-Q product name, but the same NVIDIA table lists its dense BF16 peak as 438.9 TFLOP/s rather than 503.8. Could you please narrow this key to RTX PRO 6000 Blackwell Workstation Edition so the current automatic substring matching does not classify Max-Q incorrectly?

If you have access to a Max-Q card, could you also check the exact value returned by torch.cuda.get_device_name() and add its 438.9 entry? If hardware is not available, a small unit test with the two expected full device strings would still be very helpful.

Thanks again for the contribution!

[ralovets]

Thanks—good catch. I reproduced the collision: both expected product names resolved to 503.8 TFLOP/s with the shared key. Commit f4eaa06 replaces it with distinct full Workstation (503.8) and Max-Q Workstation (438.9) entries and adds a two-case regression test. I do not have Max-Q hardware available, but I rechecked the narrowed Workstation key on the physical RTX PRO 6000 and it still resolves to 503.8. The full MFU unit file passes (34 tests), along with Ruff check and format.

[yuhezhang-ai]

/ok to test f4eaa06

[akoumpa]

/ok to test e473443