Repository navigation
fix(w4a): gate the FP8/NVFP4 kernels off ROCm - #3554
Merged
Merged
Conversation
The W4A kernels cannot be lowered by the AMD backend: - `_grouped_fp8_gemm` uses `NVMMADistributedLayout` and Gluon `mma_v2`, so Triton emits `ttg.nvidia_mma` and gfx1201 aborts with "Unsupported DotOp found when converting TritonGPU to LLVM" (`PassManager::run failed`). - The INT4->E4M3 packers round through the PTX instruction `cvt.rn.satfinite.e4m3x2.f32`. The unpack kernel does compile on gfx1201, but it returns wrong codes (zeros against the Torch decode), so leaving it to fail at runtime would corrupt accuracy silently instead. `torch.cuda.is_available()` is true under HIP and `get_device_capability()` reports a synthetic `(12, 0)` on gfx1201 that clears the FP8 floor, so neither can gate this backend. Kernel selection is already covered: `SUPPORTS_DEVICES` declares these kernels NVIDIA-only and the device family resolves through `normalize_device()`. Two paths still reached the kernels: - `validate(device=None)`: the AUTO path skips the base family check, so a ROCm host was accepted and the first forward aborted inside Triton. - Direct construction and the packer entry points, including the unpack that returns wrong codes instead of failing. So `validate()` now resolves the requested target with `normalize_device()` (the centralized HIP resolver) and reports an actionable error, and `native_fp8_gemm_error()`, `unpack_int4_to_e4m3()`, `fp8_pack()` and `nvfp4_pack_and_swizzle()` refuse HIP up front. W4A tests gate on `w4a_kernels_available()`, which compares the resolved default family against the kernel's own `SUPPORTS_DEVICES` declaration, instead of `torch.cuda.is_available()`. No weight representation or lifetime changes: the guards only reject a platform, and the lifecycle tests still assert that no expanded weight is retained. Measured on AMD Radeon AI PRO R9700 (gfx1201, HIP 7.16.26385, Triton 3.9.0, GPUs 2/3) before the fix: 7 failed / 46 passed in test_w4a_cache_lifecycle.py. After: 47 passed / 9 skipped, where every skip is an NVIDIA-only kernel case.
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
tests/kernels/test_w4a_cache_lifecycle.pyfails 7/53 on gfx1201:The W4A kernels are NVIDIA-only by construction:
_grouped_fp8_gemmusesNVMMADistributedLayoutand Gluonmma_v2, so Triton emitsttg.nvidia_mmaand the AMD backend aborts.cvt.rn.satfinite.e4m3x2.f32).torch.cuda.is_available()is true under HIP, andget_device_capability()reports a synthetic(12, 0)on gfx1201 that clears the FP8 floor, so neither can gate this backend.Second, worse failure mode: the INT4->E4M3 unpack kernel does compile on gfx1201, but it returns wrong codes (all zeros against the Torch decode). That is silent accuracy corruption, not a crash — it must be refused, not left to fail at runtime.
Fix
#3550 already resolves
cuda/cuda:N/ordinals toDEVICE.ROCMon HIP and enforcesSUPPORTS_DEVICES, so kernel selection no longer dispatches to W4A. Two paths still reached the kernels:validate(device=None)(AUTO)normalize_device(), returns an actionableNotImplementedErrornative_fp8_gemm_error(),unpack_int4_to_e4m3(),fp8_pack(),nvfp4_pack_and_swizzle()refuse HIP up frontPer
.agents/skills/gptqmodel-device-families/SKILL.md, the decision comes from the resolved family (normalize_device()), not a scatteredIS_ROCMcheck, and the kernels keep their staticSUPPORTS_DEVICES = [DEVICE.CUDA]declaration.W4A tests now gate on
w4a_kernels_available(), which compares the resolved default family against the kernel's ownSUPPORTS_DEVICES, instead oftorch.cuda.is_available()._capability()no longer probes an NVIDIA SM version on AMD.No weight representation or lifetime changes: the guards only reject a platform, and the lifecycle tests still assert that no expanded weight is retained.
Tests
AMD Radeon AI PRO R9700 (gfx1201, HIP 7.16.26385, Triton 3.9.0,
HIP_VISIBLE_DEVICES=2,3):Final set on this branch —
test_w4a_cache_lifecycle.py,test_w4a_activation_policy.py,test_rocm_kernel_contracts.py,test_w4afp8_gb10.py,test_w4a_nvfp4_gb10.py,test_w4a_stream.py:Live checks on the same host:
w4a_kernels_available() -> False;W4AFP8Linear.validate(device=None) -> (False, NotImplementedError);select_quant_linear(GPTQ_W4AFP8)raises aValueErrorwhose text is the platform reason.ruffclean, 119-column limit respected.NVIDIA behaviour is unchanged by construction (
normalize_device(0)isDEVICE.CUDA, which is inSUPPORTS_DEVICES), but this branch was not executed on an NVIDIA GPU.Note on
maintests/test_kernel_device_error_reason.py::test_rocm_rejection_reason_survivesis red onorigin/mainat5cf7b7bd: #3550 changed the message to "is CUDA-only and does not run on ROCm" but the assertion still expects "not supported on ROCm". I reproduced it on a pristineorigin/mainworktree; it is unrelated to this branch and not fixed here.