feat(quant): implement fp4_act_quant Triton kernel for DeepSeek-V4 (#807)#1479
Open
haowu1234 wants to merge 1 commit into
Open
feat(quant): implement fp4_act_quant Triton kernel for DeepSeek-V4 (#807)#1479haowu1234 wants to merge 1 commit into
haowu1234 wants to merge 1 commit into
Conversation
…OCm#807) Add production BF16 → packed FP4 e2m1 quantize kernel for the DeepSeek-V4 Indexer and Compressor paths on AMD MI300X. Highlights: - _fp4_act_quant_kernel: Triton JIT kernel with column-wise load, per-row amax + ue8m0 scale, FP4 e2m1 snapping to 8 magnitudes, 4-bit nibble encode, and pairwise pack into uint8 - _fp4_act_quant_torch: pure-torch reference with identical output - fp4_act_quant: main API returning (packed_uint8, scale_e8m0) - fp4_act_quant_inplace: QAT round-trip shim preserved for compat - tools/test_fp4_act_quant.py: standalone correctness + benchmark Pack format (matching dequant_fp4_e2m1): byte = (nibble_high << 4) | nibble_low, 2 FP4 per byte nibbles 0-7 = positive magnitudes, 8-15 = same negative scale = float8_e8m0fnu, per-1×32 block Validated on MI300X (gfx942, ROCm 7.1, Triton 3.6): - 13 shapes × 3 block_sizes all packed_diff=0, scale_diff=0 - roundtrip (quant + dequant) off_grid=0.00 - 44x speedup on DSv4 Indexer kv (1,128,256) - 182x peak on (16,128,256) Refs: ROCm#807
Contributor
|
Hello, we plan to use the kernel that combines Hadamard transformation and fp4 quantization in this PR ROCm/aiter#4029 |
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.
Summary
Implement the
fp4_act_quantkernel (BF16 → packed FP4 e2m1 uint8 + E8M0 scale) as identified in the #807 kernel gap analysis. This replaces the torch fallback for the DeepSeek-V4 Indexer Q and Compressor KV quantization paths.Changes
atom/model_ops/quant_v4.py(+395/-23): Add Triton JIT kernel_fp4_act_quant_kernelwith column-wise load, per-row amax + ue8m0 scale, FP4 e2m1 snapping, 4-bit nibble encoding, and pairwise pack into uint8. New public APIfp4_act_quant(x) → (packed_uint8, scale_e8m0).fp4_act_quant_inplacepreserved as QAT round-trip shim.tools/test_fp4_act_quant.py(new, +343): Standalone correctness test (Triton vs torch, roundtrip, pack format, edge cases) and benchmark via direct module load (no aiter dependency).Pack format
Matches existing
dequant_fp4_e2m1:Validation (MI300X VF, ROCm 7.1, Triton 3.6)
Correctness:
Performance vs torch fallback:
(1,128,256)(16,128,256)(32,512)Closes: #807 (fp4_act_quant sub-task)