Skip to content

test(moe): grant one bf16 ulp only at rounding ties in the ds_fp4 transcription check - #668

Open
void-mckenzie wants to merge 1 commit into
FlashML-org:mainfrom
void-mckenzie:fix/dsfp4-gemv-bf16-tie-tolerance
Open

void-mckenzie wants to merge 1 commit into
FlashML-org:mainfrom
void-mckenzie:fix/dsfp4-gemv-bf16-tie-tolerance

Conversation

@void-mckenzie

Copy link
Copy Markdown
Contributor

Closes #649.

The transcription check rounds each route's output to bf16 before the fp32 sum. On
sm_120 the kernel's accumulation order puts token 1 / route 2 / column 493 of the
seeded inputs at -97.75000x, 1e-5 past the bf16 tie between -97.5 and -98.0, while
the torch reference lands 1e-5 short of it: the same value rounds to -98.0 on one
side and -97.5 on the other. The four routes cancel to -4.5, so the one-ulp (0.5)
difference exceeds the allowed 0.11 (atol + rtol*|sum| with the flat 2e-2).

reference() now also returns a per-element tie budget: the largest bf16 ulp among
routes whose raw output sits within 0.1 ulp of a rounding tie, which is the only
place a ~1e-5 relative difference in the fp32 accumulation can change the rounded
result. _assert_close widens atol + rtol*|want| by that budget, so every element
away from a tie keeps exactly the 2e-2 check. The grouped-prefill, cpu-executor and
inactive-routes tests compare the same sum and go through the same helper.

Tested on RTX 5090 (sm_120), driver 580.95.05, torch 2.11.0+cu130, python 3.10.19;
synthetic banks, no checkpoint:

.venv/bin/python -m pytest tests/moe/test_dsfp4_reference_semantics.py -q
  • main 95ac86f: test_gemv_path_matches_the_transcription[32] and [128] fail,
    1 / 3072 elements, abs 0.5 at (1, 493) - identical to the issue
  • this branch: 7 passed

This branch has not been deployed

No deployments
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.

test_gemv_path_matches_the_transcription fails on sm_120 (bf16 rounding tie exceeds the tolerance)

1 participant