[Refactor] Dynamic registration of FP8 data type for compatibility with older PyTorch versions

[LeiWang1999]

This pull request introduces a dynamic approach to handling FP8 data types in the tilelang/language/v2/dtypes.py module and updates a default argument in maint/gemm_v2/latency_mha_fwd.py. The main improvement is making FP8 dtype support conditional on their availability in the installed version of PyTorch, improving compatibility and robustness.

Dynamic FP8 dtype handling and default parameter update:

Data type conversion improvements:

• Refactored the dtype conversion list by splitting the static base types into _dtype_cvt_base and moving FP8-related types to a dynamic section. This ensures FP8 types are only added to _dtype_cvt if they exist in the current PyTorch installation, preventing import errors on unsupported versions. [1] [2]

Parameter default update:

• Changed the default value of the --dim argument in maint/gemm_v2/latency_mha_fwd.py from 512 to 256.

Summary by CodeRabbit

• Chores

  • Updated default latency benchmark dimension from 512 to 256 in command-line parameters.

• Improvements

  • Enhanced FP8 data type support with runtime detection based on available framework capabilities.

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

Walkthrough

Two independent modifications: a CLI parameter default adjustment in a latency measurement script (dim from 512 to 256), and internal dtype mapping refactoring to support conditional FP8 dtype registration based on torch availability.

Changes

Cohort / File(s)	Summary
Latency measurement configuration maint/gemm_v2/latency_mha_fwd_bhsd.py	Changed --dim argument default value from 512 to 256, altering the default dimensionality for Q/K/V tensors in MHA forward pass latency benchmarks.
Internal dtype system refactoring tilelang/language/v2/dtypes.py	Renamed existing dtype conversion list from _dtype_cvt to _dtype_cvt_base, introduced _fp8_dtype_mappings for FP8 types, and reconstructed _dtype_cvt to conditionally append FP8 mappings only when corresponding torch attributes are available.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

• File 1 is a straightforward parameter default change with no logic impact
• File 2 involves basic list operations and conditional appending; verify that the FP8 dtype mappings align with intended torch version compatibility checks

Possibly related PRs

• [BugFix] Correct direct copy from bf16 to fp8 #1090: Adds FP8 types and conversion machinery at the dtype representation level, complementing this PR's conditional FP8 dtype registration logic.

Poem

🐰 A script dims its gaze from 512 to 256,
While dtypes dance conditionally in the mist—
FP8 joins when torch permits,
Base and mapped, refactored bits,
Precision tuned, the benchmark's gift! ✨

Pre-merge checks and finishing touches

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main objective of the pull request—refactoring dynamic FP8 dtype registration for PyTorch compatibility, which aligns with the primary change in tilelang/language/v2/dtypes.py.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a9d823b and 00e8c73.

📒 Files selected for processing (2)

• maint/gemm_v2/latency_mha_fwd_bhsd.py (1 hunks)
• tilelang/language/v2/dtypes.py (2 hunks)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)

• GitHub Check: Test for Python 3.12 with ROCm-6.3 (on self-hosted-amd)
• GitHub Check: Test for Python 3.12 with Metal (on macos-latest)

🔇 Additional comments (3)

maint/gemm_v2/latency_mha_fwd_bhsd.py (2)

15-15: LGTM: Default dimension reduced.

The default dim value change from 512 to 256 aligns with the PR objectives and will affect the default tensor dimensionality in benchmarks.

---

18-18: Flag undocumented functional change.

The --use_v2 flag and its usage (lines 22, 73-76, 90-93) introduce v1/v2 execution path selection but are not mentioned in the PR objectives, which only describe the dim default change and FP8 dtype refactoring.

tilelang/language/v2/dtypes.py (1)

13-56: LGTM: Dynamic FP8 dtype registration implemented correctly.

The refactoring splits the base types into _dtype_cvt_base and dynamically extends _dtype_cvt with FP8 types only when they exist in the installed PyTorch via hasattr checks. This prevents import errors on older PyTorch versions that lack FP8 support, achieving the stated PR objective.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[LeiWang1999]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Nice work!

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".