feat(wgpu/msl): register bf16 when the device probe-compiles bfloat - #1509
Open
scott-wyatt wants to merge 1 commit into
Open
scott-wyatt wants to merge 1 commit into
scott-wyatt wants to merge 1 commit into
Conversation
The msl passthrough's Metal backend registers f16/f32 but never bf16 — yet the bfloat MSL codegen already ships in cubecl-cpp's metal dialect (reductions fall back through float, shuffles through ushort, float-returning intrinsics cast back). On macOS 14+ wgpu-hal compiles MSL at 3.1+, where `bfloat` is a native type, so on those systems the missing capability registration is the only gap. Rather than encoding the version/family table here (wgpu-hal keeps its resolved msl_version private), the gate probe-compiles a one-line bfloat kernel on the adapter's raw device at init: if the device+OS pair compiles it, bf16 is registered with full usage; any error fails closed to today's behavior. A few milliseconds, once per device init. macOS 13 and earlier (MSL 3.0 — no bfloat) refuse the probe and are unchanged. Measured on an M1 Max (Apple7, macOS 26): a 10-step rank-8 LoRA training run over an 8B Llama-arch model in bf16 completes with finite loss through this path — load, forward, backward, optimizer. Throughput note, stated plainly: bf16 ran ~16x slower than the same run at f16 on this GPU family (5.3 vs 84.3 tok/s), consistent with per-op conversion and no native bf16 SIMD on Apple7; newer families were not measured and are expected to fare better. The value on Apple7 is correctness (bf16's exponent range), not speed. The objc2-metal/objc2-foundation deps are macOS-only, optional behind the msl feature, and version-matched to wgpu-hal's own objc2 stack, so they add no new lockfile entries.
Collaborator
|
I added a 3.2 check yesterday because the new compiler requires lambdas so the canary shader is no longer required. |
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.
The msl passthrough's Metal backend registers f16/f32 but never bf16 — yet the bfloat MSL codegen already ships in
cubecl-cpp's metal dialect (reductions fall back throughfloat, shuffles throughushort, float-returning intrinsics cast back). On macOS 14+, wgpu-hal compiles MSL at 3.1+, wherebfloatis a native type, so on those systems the missing capability registration is the only gap.Approach
Rather than encoding the OS/family version table here (wgpu-hal keeps its resolved
msl_versionprivate), the gate probe-compiles a one-linebfloatkernel on the adapter's raw device at init: if the device+OS pair compiles it, bf16 is registered with full usage; any error fails closed to today's behavior. It costs a few milliseconds, once per device init. macOS 13 and earlier (MSL 3.0 — nobfloat) refuse the probe and are unchanged.Happy to re-gate differently if you'd prefer another idiom (e.g., exposing
msl_versionvia your wgpu fork).Measured
On an M1 Max (Apple7 family): a 10-step rank-8 LoRA training run over an 8B Llama-arch model in bf16 completes with finite loss through this path — load (16.5 GB, 4 shards), forward, backward, optimizer step.
Throughput, stated plainly: bf16 ran ~16× slower than the same run at f16 on this GPU family (5.3 vs 84.3 tok/s), consistent with per-op conversion and no native bf16 SIMD on Apple7. Newer families were not measured and are expected to fare better. The value on Apple7 is correctness (bf16's exponent range), not speed.
Dependencies
objc2-metal/objc2-foundationare macOS-only, optional behind themslfeature, and version-matched to wgpu-hal's own objc2 stack — no new lockfile entries.🤖 Generated with Claude Code