Skip to content

fix(kernel): correct fp4 cache scale rounding - #720

Merged
jason-fxz merged 2 commits into
FlashML-org:mainfrom
LZ-QWQ:fix/dsv41-fp4-scale-rounding
Oct 16, 2026
Merged

jason-fxz merged 2 commits into
FlashML-org:mainfrom
LZ-QWQ:fix/dsv41-fp4-scale-rounding

Conversation

@LZ-QWQ

@LZ-QWQ LZ-QWQ commented Oct 16, 2026

Copy link
Copy Markdown
Contributor

Based on the analysis and proposed fix in #673, fix DSV4.1 FP4 cache scale rounding reported on RTX 3090 Ti and reproduced on Radeon GPUs and RTX4090 with software emulation.

The compiler can fuse reciprocal multiplication with the emulated e4m3 magic-add into an FMA, skipping intermediate FP32 rounding and changing subnormal ties. Use tl.math.div_rn(amax, 6.0) only for software emulation; retain the existing multiplication for native conversion through compile-time e4m3_native_cx().

Validation:

python -m pytest tests/kernels/test_dsv41_pack.py

GPU / path Before After
W7900D 8 passed, 2 failed 10 passed
R9700 8 passed, 2 failed 10 passed
RTX4090 default 10 passed 10 passed
RTX4090 emulated 8 passed, 2 failed 10 passed

For RTX4090 emulation, prefix the command with FREETOKEN_FORCE_E4M3_EMU=1. KV pool and sparse-attention tests also pass (15/15 per path).

AMD validation used ROCm 10.1 with #535's torchvision prerelease lower bound.

The checked RTX4090 default specializations produce byte-identical executable code before and after the change.

@jason-fxz

Copy link
Copy Markdown
Collaborator

LGTM.

@jason-fxz
jason-fxz merged commit d5eb18f into FlashML-org:main Oct 16, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants