feat(sm100): add FP4 activation support - #75
Draft
mgoin wants to merge 1 commit into
Draft
Conversation
Add native MXFP4 and NVFP4 activation execution across dense and MoE GEMM layouts while preserving the existing weight packing. Normalize explicit input scales and use the packed physical row stride for logical K tails and indexed rows. Co-authored-by: Codex <codex@openai.com> Signed-off-by: mgoin <mgoin64@gmail.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
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.
Adds native SM100-family E2M1 FP4 activations for MXFP4 (group32/E8M0) and NVFP4 (group16/E4M3), with BF16 outputs. Dense, indexed, grouped-contiguous and grouped-masked GEMMs reuse existing weight packing and output handling. Coverage includes supported smaller integer weights dequantized to FP4, channel scales and NVFP4 secondary tensor scales.
Packed FP4 activations stay in shared memory; canonical weight fragments and scales are staged into TMEM for native block-scaled UMMA. The initial schedule uses BM64/BM128 and K128/K256 stages with bounded occupancy choices. Existing integer A4 and SM120 execution retain their paths. Explicit input scales and indexed gathers use the physical packed row stride, including logical K tails.
Integration status: this is the independently validated A4 branch against main. It overlaps A8 #73 in the new MXUMMA executor, eligibility and geometry selection. Keep it draft until those shared changes are reconciled; the two drafts cannot simply be merged as written. The A8 native-output optimization #74 has not been ported here. This preserves the tested A4 implementation for review without claiming a validated combined A4/A8 executor.
11 files, +564/-15, including 295 added lines in the existing MXMMA test module. The publication tree is identical to validated
83c616a. Open Humming PRs were checked; current SM90 tuning work does not duplicate FP4 activation support on SM100.Validation on B300/SM103, PyTorch 2.13.0+cu130, CUDA 13.0, generated sm_100f (CUDA >=12.9):
The support campaign passed 20 native cases across its recorded runs, public internal/prequantized input tests, compiler eligibility checks, five BF16 pipeline regressions and eight legacy integer-A4 comparisons. The later scale-stride correction passed 24 strict cases, six tensor-operation memchecks with zero errors and both retained K2880 public regressions. All active outputs are explicitly checked against independent references. Ruff and
git diff --checkpassed.Identical-input dense MXFP4 projection, N=K=4096, CUPTI/CUDA graphs/cold L2: at M512, Humming 1201 TFLOP/s versus DeepGEMM 1588; at M2048, 1786 versus 2818. Quantization, packing and routing are excluded. This is initial kernel/library support, not native performance parity; no A4 vLLM model evaluation or distributed EP measurement was performed.
AI assistance: implementation, review and validation used OpenAI Codex, including GPT-6 Astra.