Skip to content

[MetaxGPU][quantize] Use portable C++ MXFP4 dequant on Maca #114

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ventijing wants to merge 13 commits into
base: dev
Choose a base branch
from
Open

[MetaxGPU][quantize] Use portable C++ MXFP4 dequant on Maca #114

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
1d335ad
[MetaxGPU][quantize] Use portable C++ MXFP4 dequant on Maca
Jun 22, 2026
b3ac85a
[MetaxGPU][regression] Default regression_all scan root to examples
Jun 22, 2026
1a6a8cb
[MetaxGPU][regression] Default regression_all scan root to examples
Jun 22, 2026
3393e0b
[CI] Update flash_linear_attention version to 0.4.0+metax3.7.2.0torch2.8
Jun 22, 2026
a09c835
[CI] Update flash_linear_attention version
Jun 22, 2026
3ce4d43
[CI] Update flash_linear_attention version
Jun 22, 2026
3eca72d
[CI] Update flash_linear_attention version
Jun 22, 2026
a66f8b6
[MetaxGPU][quantize] Use portable C++ MXFP4 dequant on Maca
Jun 23, 2026
6959071
[CI] Update flash_linear_attention version
Jun 23, 2026
eb23897
[MetaxGPU][quantize] Use portable C++ MXFP4 dequant on Maca
Jun 23, 2026
e78ebdd
[CI] Update flash_linear_attention version
Jun 23, 2026
3404026
[CI] Update flash_linear_attention version
Jun 24, 2026
6ed645f
[CI] Update flash_linear_attention version
Jun 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 12 additions & 3 deletions .github/workflows/pr-regression-test-bot.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,7 +207,10 @@ jobs:
fi
uv pip install -v -r requirements-test.txt
if [[ "${{ matrix.runner.toolkit }}" == *"MACA"* ]]; then
uv pip install -v --no-deps --python-version 3.10.0 flash_linear_attention==0.4.0+metax3.5.3.9torch2.8 -i https://repos.metax-tech.com/r/maca-pypi/simple --trusted-host repos.metax-tech.com
uv pip uninstall -y flash-linear-attention flash_linear_attention 2>/dev/null || true
uv pip install --force-reinstall --no-deps --python-version 3.10.0 \
flash_linear_attention==0.4.0 \
-i https://repos.metax-tech.com/r/maca-pypi/simple --trusted-host repos.metax-tech.com
fi
uv pip install -v .
if [[ "${{ matrix.runner.toolkit }}" == *"MACA"* ]]; then
Expand All @@ -229,7 +232,10 @@ jobs:
fi
uv pip install -v -r requirements-test.txt
if [[ "${{ matrix.runner.toolkit }}" == *"MACA"* ]]; then
uv pip install -v --no-deps --python-version 3.10.0 flash_linear_attention==0.4.0+metax3.5.3.9torch2.8 -i https://repos.metax-tech.com/r/maca-pypi/simple --trusted-host repos.metax-tech.com
uv pip uninstall -y flash-linear-attention flash_linear_attention 2>/dev/null || true
uv pip install --force-reinstall --no-deps --python-version 3.10.0 \
flash_linear_attention==0.4.0 \
-i https://repos.metax-tech.com/r/maca-pypi/simple --trusted-host repos.metax-tech.com
fi
uv pip install -v .
if [[ "${{ matrix.runner.toolkit }}" == *"MACA"* ]]; then
Expand All @@ -244,7 +250,10 @@ jobs:
fi
uv pip install -v -r requirements-test.txt
if [[ "${{ matrix.runner.toolkit }}" == *"MACA"* ]]; then
uv pip install -v --no-deps --python-version 3.10.0 flash_linear_attention==0.4.0+metax3.5.3.9torch2.8 -i https://repos.metax-tech.com/r/maca-pypi/simple --trusted-host repos.metax-tech.com
uv pip uninstall -y flash-linear-attention flash_linear_attention 2>/dev/null || true
uv pip install --force-reinstall --no-deps --python-version 3.10.0 \
flash_linear_attention==0.4.0 \
-i https://repos.metax-tech.com/r/maca-pypi/simple --trusted-host repos.metax-tech.com
fi

- name: Clear uv cache for self-hosted runners (if setup failed)
Expand Down
2 changes: 1 addition & 1 deletion maint/scripts/regression_all.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ def _parse_table(output: str) -> dict[str, float]:


def _examples_root() -> Path:
return Path(__file__).resolve().parents[2] / "examples" / "maca"
return Path(__file__).resolve().parents[2] / "examples"

@coderabbitai coderabbitai Bot Jun 22, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Expanded examples scope can silently drop benchmark results due to name collisions.

By switching discovery to the repo-wide examples tree (Line 60), regression_all() now has a much higher chance of encountering duplicate benchmark names across different drivers. The current merge logic (if k not in merged) silently ignores later entries, which can hide regressions and undercount totals.

A safer approach is to either fail on duplicate names or namespace keys by file path when merging.

Suggested fix (fail fast on duplicate benchmark names) 
@@
-        for k, v in parsed.items():
-            if k not in merged:
-                merged[k] = v
-                _RESULTS.append(PerfResult(name=k, latency=v))
+        for k, v in parsed.items():
+            if k in merged:
+                failures.append(
+                    f"{rel_path}\nDuplicate benchmark name detected: {k!r}. "
+                    "Benchmark names must be globally unique across discovered drivers."
+                )
+                continue
+            merged[k] = v
+            _RESULTS.append(PerfResult(name=k, latency=v))
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@maint/scripts/regression_all.py` at line 60, The merge logic using `if k not
in merged` silently skips duplicate benchmark names, which can hide regressions
when the expanded examples scope encounters multiple benchmarks with the same
name across different drivers. Modify the merge operation in the
regression_all() function to detect when a key already exists in the merged
dictionary and raise an exception or error with details about the collision,
rather than silently ignoring the duplicate entry. Alternatively, namespace the
keys by including the source file path along with the benchmark name to ensure
uniqueness while preserving all results.



def _discover_bench_files(examples_root: Path) -> list[Path]:
Expand Down
50 changes: 35 additions & 15 deletions tilelang/quantize/mxfp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,10 @@
from typing import Literal

from tvm.target import Target

from tilelang import language as T
from tilelang.backend.target import determine_target
from tilelang.rocm.target import target_is_gfx950, target_is_hip

# Implementation asm for fp4 to bf16, using twiddling
# Reference: https://github.com/triton-lang/triton/blob/main/python/triton_kernels/triton_kernels/tensor_details/layout_details/hopper_value.py#L11-L18
Expand Down Expand Up @@ -157,6 +162,28 @@
"""


def _resolve_mxfp_target(target):
if target is not None and target != "auto":
return target
current = Target.current(allow_none=True)
if current is not None:
return current
return determine_target("auto", return_object=True)
Comment thread
ventijing marked this conversation as resolved.


def _target_uses_portable_mxfp_dequant(target) -> bool:
"""Return True for targets that cannot compile CUDA PTX inline asm (e.g. Maca, AMD gfx950)."""
if target is None:
return False
if not isinstance(target, Target):
target = Target(target)
if target.kind.name == "maca":
return True
if target_is_hip(target):
return target_is_gfx950(target)
return False


def get_mxfp_intrin_group(
out_dtype: Literal[T.float16, T.bfloat16] = T.bfloat16,
source_format: Literal[T.int, T.uint] = T.uint,
Expand Down Expand Up @@ -195,33 +222,26 @@ def get_mxfp_intrin_group(
assert source_format in [T.int, T.uint], f"Invalid source_format: {source_format}. Expected 'int' or 'uint'."
assert storage_dtype in [T.int32, T.int8, T.uint8], f"Invalid storage_dtype: {storage_dtype}. Expected 'int32' or 'int8' or 'uint8'."

# Detect AMD gfx950 target to select the HIP C++ dequantization implementation.
# All other targets (NV, RDNA, MI300) use the default CUDA PTX path below.
_is_gfx950 = False
if target is not None:
try:
from tilelang.rocm.target import target_is_gfx950

_is_gfx950 = target_is_gfx950(target)
except (ImportError, ModuleNotFoundError, AttributeError):
# target_is_gfx950 unavailable in this build; assume non-gfx950.
pass
# Maca and AMD gfx950 cannot compile CUDA PTX; use portable C++ below.
# All other targets (NV, RDNA, MI300) use the default CUDA PTX path.
# target=None keeps the CUDA PTX default; only target="auto" resolves from context.
_resolved = _resolve_mxfp_target(target) if target == "auto" else target
_use_portable = _target_uses_portable_mxfp_dequant(_resolved)

dtype_map = {T.float16: "f16", T.bfloat16: "bf16"}
func_name = f"decode_fp{source_bit}_to_{dtype_map[out_dtype]}"
if use_twiddling:
func_name += "_twiddling"

if _is_gfx950:
# AMD gfx950 path: use portable HIP C++ implementations.
# The function name stays the same so the call site is unchanged.
if _use_portable:
# Portable C++ path (Maca / AMD gfx950). Function name unchanged for call sites.
if use_twiddling and source_bit == 4 and out_dtype == T.bfloat16:
return {"func_name": func_name, "c_source": decode_f4_to_bf16_twiddling_hip}
elif not use_twiddling and source_bit == 4 and out_dtype == T.bfloat16:
return {"func_name": func_name, "c_source": decode_f4_to_bf16_simple_hip}
else:
raise AssertionError(
f"AMD gfx950 MXFP dequant only supports source_bit=4 and out_dtype=bfloat16, "
f"Portable MXFP dequant only supports source_bit=4 and out_dtype=bfloat16, "
f"got source_bit={source_bit}, out_dtype={out_dtype}"
)

@Five-HZ Five-HZ Jun 24, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add if branch in here like gfx950 to resolve same issue with maca target.


Expand Down
Loading