Fix encoder CPU dwte bucket width - #3
Draft
SaiSandeepKantareddy wants to merge 1 commit into
Draft
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
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. Comment |
Owner
|
Hi @SaiSandeepKantareddy! Thanks for reporting this issue via draft PR. I will have a fix that includes a regression tests up for this shortly. |
ulmentflam
added a commit
that referenced
this pull request
Aug 13, 2026
Reported in #3 against the CPU custom-op path on Apple Silicon / macOS, where encoder backward left dwte at zero for token rows present in the batch. Bucket metadata is emitted against the module-level contract `WTE_C_PER_WARP = 32 * WTE_BWD_SIMD_WIDTH` (128), and the equivalence test's Python reference builder hardcodes the same 128, but the CPU consumer re-derived the span as `WARP_SIZE * simd_width_of[dtype]()`. The GPU kernel uses that identical expression and is right there only because on a GPU target it evaluates to exactly 32 * 4 == 128; the CPU path inherited the expression without the property that makes it true, and WARP_SIZE is meaningless on a CPU target anyway. The span is wrong on every platform, but only wrong in RESULT when it comes out smaller than 128. On this x86-64 box it is 32 * 16 = 512: group 0 covers [0, 512), every higher group computes c_len <= 0 and is skipped, and the union is still exactly [0, channels) covered once, so the answer is correct and merely leaves one bucket doing all the work. At a stride of 0, which is what the reporter's sandboxed CPU compile produces, nothing is written at all. Worth noting the mechanism in the report does not fire on Apple Silicon by itself: NEON fp32 gives width == 4 and 32 * 4 == 128 would have matched, so the operative fact is WARP_SIZE resolving to 0 there, not the SIMD width. Two independent gaps hid this. `make test` on a GPU box auto-detects the GPU and runs test-python-cuda, so the CPU custom-op path is never executed. And every encoder case had channels <= 128, so ceildiv(channels, 128) == 1, channel_group was always 0, and any wrong stride cancelled at c_base == 0. Adds a channels=256 case (fp32_multi_channel_group) so the multi-group path is covered at all. Be clear about what it buys: on x86-64 it passes both before and after this fix, because a too-large stride still tiles correctly. It is real coverage for a previously untested path and it catches this bug where the bug is live, but it would not have failed here. Catching a too-small stride on x86 needs an assertion on the per-bucket span itself rather than end-to-end gradients. Verified: encoder equivalence 8 passed on the CPU custom-op path and 8 passed on the GPU path; full gate green (21 mojo test files, 245 passed / 5 skipped, zero warnings). `make test-python` on this box is separately broken by the upstream MAX CPU custom-op SIGSEGV (docs/ai/max_cpu_custom_op_crash_2026-07-24.md); confirmed pre-existing by reproducing the adamw crash at fa74e79 with this change stashed. Analysis: docs/ai/wte_backward_cpu_channel_span_bug.md. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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
2 participants
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.
Why
While setting up the repo to look into possible benchmark contributions, I hit a failing encoder backward equivalence test on Apple Silicon/macOS.
make test-pythonwas failing intests/test_encoder_equivalence.pybecause the CPU custom-op path for encoder backward left token embedding gradients (dwte) as zero for token rows that appeared in the batch.The issue was in
wte_backward_cpu: it derived each bucket's channel span fromWARP_SIZE * width, but the bucket metadata is produced using the module-levelWTE_C_PER_WARPcontract. On the CPU custom-op path this made the channel span empty, so the accumulation loop never wrotedwte.What changed
Use
WTE_C_PER_WARPdirectly when mapping a bucketchannel_groupto its channel range in the CPUwte_backward_cpupath.Also added a short
docs/ai/bug note describing the symptom, root cause, fix, verification, and AI-use disclosure.Verification
tests/test_encoder_equivalence.py:6 passed12 passedmake test-python:233 passed, 15 skippedmake build: passedmake lint: passedmake check: passedNote: Mojo printed Crashpad warnings under sandboxed macOS, but the relevant commands exited successfully. One pre-push attempt hit a transient Mojo/Metal
build-profilecompiler failure; runningmake build-profiledirectly passed, and the subsequent pre-pushmake checkpassed.