examples: Fast Hadamard + MXFP4 Quant fused kernel - #237
Draft
Mocchibird wants to merge 2 commits into
Draft
Conversation
Two examples, each fusing a Hadamard rotation with MXFP4 quantization into a single dav-c310-vec launch. Unfused, the butterfly writes the rotated tile out and the quantizer reads it straight back; fused, the tile never leaves UB and only the nibbles and scales are written, which is 6.53 B/element against 2.53. fused_hadamard_quant_a5 rotates the whole row. Sylvester factors as H_K = H_(K/256) (x) H_256, so the transform runs in two phases: an order-256 transform inside each 256-element window, then log2(K/256) stages pairing windows elementwise. Neither phase holds more than one window in registers, so the row width does not enter the register budget -- ten powers of two from 32 to 16384. A single-phase butterfly keeping a whole row in registers stops at 4096. 2.45x over the same work as two launches, and 1.55-1.58x a device-to-device copy at 1382-1398 GB/s, flat within 1.4 us across a 16x range of K. fused_hadamard_quant_b32_a5 rotates independent 32-element blocks, so the rotation is 32 wide however long the row is. K therefore carries neither a power-of-two constraint nor a register-budget limit: 28 widths from 32 to 16384, covering 4096 and the 11008-style widths together. 2.45-2.54x over two launches at 1448-1450 GB/s. Its 32-wide rotation also matches MXFP4's scale granularity, and on heavy-tailed data it measured 4-5% lower quantization error than the full-row rotation at K=4096. K is a template parameter in both, so one .so per example holds an instantiation per width and the launcher dispatches on it. Rows per tile is closed-form rather than a countdown template recursion, which at K=32 would instantiate 768 deep. The butterfly is the unnormalised Sylvester matrix in both, sqrt(K) and sqrt(32) respectively, left to the caller because MXFP4's E8M0 scale is a power of two and cannot always absorb it. 43 and 48 tests. The full-row reference is an independent strided transform, itself pinned against an explicitly constructed Sylvester matrix, since that matrix is a gigabyte at K=16384. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016qbbp2Kjzv1AhUkU6w1mi4
Mocchibird
changed the title
Prospector runs pylint, and a coded suppression only reaches the tool that owns the code: `# noqa: F401` silences pyflakes but leaves pylint's unused-import, and `# noqa: E402` silences pycodestyle but leaves wrong-import-position. A bare `# noqa` covers every tool prospector runs, which is what the kernels already in the repo use (fast_hadamard/standard/bench_hadamard.py). torch_npu / sibling imports coded noqa -> bare noqa benchmark.py record builders dict(...) -> a dict literal the `seeded` fixture autouse, so it is not an unused argument row_quantum() deleted; it had no caller in either tree The seeded fixture now applies to every test in its module rather than the eight that named it, so torch.npu.set_device(0) runs for all of them. Verified: prospector reports 0 messages on the whole tree, and on an Ascend950PR_9589 with CANN 9.1.0 both suites still pass, 43 and 48 tests. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016qbbp2Kjzv1AhUkU6w1mi4
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
1 participant
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.
fused_hadamard_quant_a5, fused_hadamard_quant_b32_a5: a Hadamard and MXFP4 quantization in one launch
Adds fused bf16 Hadamard + MXFP4 quantization kernels for A5:
fused_hadamard_quant_a5: full-row rotation, power-of-two K from 32 to 16384.fused_hadamard_quant_b32_a5: independent 32-element rotations, including non-power-of-two widths such as 11008.One launch, packed FP4 output with E8M0 scales per 32 elements. Hadamard is unnormalized.
Measured on Ascend950PR_9589, M=16384. Times in µs.
At 64Mi elements, both reach ~1.4 TB/s. Input d2d copy takes 1.55–1.65x as long.
Tests: 43 passed for full-row, 42 for block-32.
Requires CANN with MXFP4 support: 9.1.0 / 9.2.0; 9.0.0 is unsupported.