onnxruntime
e2856db8 - [CUDA] Address feedback for block-scaled MatMul and MoE GEMV (#29910)

Commit
24 days ago
[CUDA] Address feedback for block-scaled MatMul and MoE GEMV (#29910) ## Description This PR addresses follow-up review feedback from recently merged CUDA PRs covering block-scaled FP4/FP8 MatMul and MoE GEMV paths. The update consolidates duplicated CUDA type helpers into a shared header, tightens kernel contracts and comments, and adds edge-case validation tests for FP4 block-scaled MatMul behavior. The goal is to improve maintainability and correctness without changing intended operator semantics. ## Summary of Changes ### CUDA helper consolidation and kernel cleanup | File | Change | |------|--------| | onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh | Expanded shared CUDA helper utilities (float/bfloat16 conversion helpers, packed vec2 traits, bit_cast helper, FP8 E4M3 raw-byte conversion helpers) and removed anonymous-namespace scoping so helpers are reusable across translation units. | | onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu | Switched to shared helper usage and removed local duplicate helper patterns; updated call sites to shared conversion and vec2 helper API. | | onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu | Switched to shared helper usage and updated FP8 GEMV path call sites/comments to use the consolidated helper API. | | onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4_sm120.cu | Switched to shared helper usage and kept SM120 quantization path aligned with consolidated type conversion helpers. | ### FP4/NVFP4 behavior and API cleanups | File | Change | |------|--------| | onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc | Updated PrePack path and argument handling based on feedback; removed obsolete weight-scale handoff usage in native SM120 launch path. | | onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h | Clarified launcher contracts/documentation and aligned signatures with current runtime behavior. | ### MoE GEMV and related style/documentation updates | File | Change | |------|--------| | onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu | Applied coding-style consistency updates and readability improvements requested in review. | | onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.h | Applied coding-style consistency updates requested in review. | | onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu | Applied coding-style consistency updates requested in review. | | onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h | Applied coding-style consistency updates requested in review. | | onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu | Added clarifying comments around invariants/assumptions raised in review. | | onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh | Applied targeted const-ordering style alignment for PR-introduced lines. | ### Tests | File | Change | |------|--------| | onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc | Added edge-case tests covering partial trailing scale block handling, K multiple-of-16 but not 32 path, and zero-K behavior with and without bias. | ## Testing - Build and install: - bash /home/tianlei/git/onnxruntime/.env/cuda130_fp4_bench.sh --build --install - Additional verification for sources excluded by the default build config: - bash /tmp/sm120_check.sh - bash /tmp/test_syntax.sh - GPU functional checks: - python /tmp/verify_fp4_tests.py - Formatting verification: - lintrunner -a --take CLANGFORMAT onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4_sm120.cu ## Motivation and Context In https://github.com/microsoft/onnxruntime/pull/29818 and https://github.com/microsoft/onnxruntime/pull/29887 and https://github.com/microsoft/onnxruntime/pull/29896, reviewer feedback identified duplicate helper logic and opportunities to tighten contracts/comments across the CUDA block-scaled MatMul and MoE GEMV code paths. Consolidating type/packing helpers in one shared CUDA header reduces duplication and future drift, while the added tests lock in edge-case expectations for FP4 block-scaled behavior. ## Checklist - [x] Tests added/updated - [x] No breaking changes - [ ] Documentation updated (not required for this internal CUDA refactor/cleanup)
Author
Parents
Loading