From 4ce84a76a21f457b42432b1e6bc4c0de9f56b854 Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Wed, 8 Jul 2026 09:08:04 -0700 Subject: [PATCH 1/5] Skip MinLength logits processor construction when eos_token_id is negative (#29604) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What Skip constructing the `MinLengthLogitsProcessor` when `eos_token_id` is negative. A negative `eos_token_id` is the "no-EOS" sentinel used by greedy/sampling generation (it defaults to `-1`). With no EOS token there is nothing to demote, so a MinLength processor built with a negative eos can only ever be a guaranteed no-op. This change guards its construction at the list level in `LogitsProcessorInitImpl` (`logits_processor.h`), so we do not build a processor that would do nothing. ```cpp if (parameters.min_length > 0 && parameters.eos_token_id >= 0) { ... // add MinLength processor } ``` ### Why This is a small **defense-in-depth / code-clarity** improvement, not a behavior change and not a correctness or security fix: - For a valid `eos_token_id >= 0`, behavior is unchanged — the processor is still constructed and enforces the minimum length exactly as before. - For a negative eos, the added guard skips a processor that would be a no-op anyway (`SetScore` already ignores negative token ids), so this is a minor performance and clarity optimization. - It mirrors the existing conditional-adds in the same function (e.g. `RepetitionPenalty`), keeping the construction logic consistent. CPU-only: the CUDA path is already inherently a no-op for a negative eos, so no CUDA code is touched. ### Tests Adds 4 unit tests in `min_length_logits_processor_test.cc` (run via `onnxruntime_provider_test --gtest_filter='*MinLengthLogitsProcessorTest*'`): - `SetScoreIgnoresNegativeTokenId` — documents the pre-existing inline backstop that ignores negative token ids. - `ListInitSkipsProcessorForNegativeEosTokenId` — drives `LogitsProcessorList::Init` with a negative eos and confirms a below-min-length run leaves scores unchanged (the processor is skipped as a guaranteed no-op). - `ListInitDemotesEosBelowMinLength` — positive control / enforcement path: with a valid eos the processor is constructed and demotes the eos score below min length. - `ListInitLeavesScoresUnchangedAtMinLength` — at the min length, no demotion. The tests only reference exported (`LogitsProcessorList::Init`/`Process`) and header-inline symbols so they link in both static and shared-library builds. --------- Signed-off-by: Tita Wang Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cpu/transformers/logits_processor.h | 5 +- .../min_length_logits_processor_test.cc | 136 ++++++++++++++++++ 2 files changed, 140 insertions(+), 1 deletion(-) create mode 100644 onnxruntime/test/contrib_ops/min_length_logits_processor_test.cc diff --git a/onnxruntime/contrib_ops/cpu/transformers/logits_processor.h b/onnxruntime/contrib_ops/cpu/transformers/logits_processor.h index fb9638bb09d59..7e186e1a87e91 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/logits_processor.h +++ b/onnxruntime/contrib_ops/cpu/transformers/logits_processor.h @@ -332,7 +332,10 @@ class LogitsProcessorList : public ILogitsProcessorList { processor_list_.push_back(prefix_vocab_mask_processor_.get()); } - if (parameters.min_length > 0) { + // For a negative "no eos" sentinel there is no token to demote, so a MinLength processor here + // would be a guaranteed no-op (SetScore ignores negative token ids). Skip constructing it as a + // defensive, minor performance optimization, consistent with the other conditional-adds above. + if (parameters.min_length > 0 && parameters.eos_token_id >= 0) { min_length_processor_ = std::make_unique>(parameters.min_length, parameters.eos_token_id); processor_list_.push_back(min_length_processor_.get()); diff --git a/onnxruntime/test/contrib_ops/min_length_logits_processor_test.cc b/onnxruntime/test/contrib_ops/min_length_logits_processor_test.cc new file mode 100644 index 0000000000000..d93858710387d --- /dev/null +++ b/onnxruntime/test/contrib_ops/min_length_logits_processor_test.cc @@ -0,0 +1,136 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include + +#include "gtest/gtest.h" +#include +#include "contrib_ops/cpu/transformers/logits_processor.h" + +namespace onnxruntime { +namespace contrib { +namespace transformers { +namespace test { + +namespace { + +// Minimal ISequences stub reporting a fixed sequence length. The MinLength path only reads +// GetSequenceLength(); the device/sequence accessors are unused here and return empty spans. +class FixedLengthSequences : public ISequences { + public: + explicit FixedLengthSequences(int sequence_length) : sequence_length_(sequence_length) {} + + gsl::span GetSequence(int /*beam_index*/) const override { return {}; } + gsl::span GetCurrentDeviceSequences() const override { return {}; } + gsl::span GetNextDeviceSequences() override { return {}; } + int GetSequenceLength() const override { return sequence_length_; } + int GetMaxLength() const override { return sequence_length_; } + + private: + int sequence_length_; +}; + +constexpr int kBatchBeamSize = 2; +constexpr int kVocabSize = 4; + +// Builds a minimal GreedySearchParameters that activates only the MinLength processor path, so the +// list-level tests exercise the MinLength construction condition without other processors interfering. +GreedySearchParameters MakeMinLengthOnlyParameters(int min_length, int eos_token_id) { + GreedySearchParameters parameters{}; + parameters.model_type = IGenerationParameters::kModelTypeGpt; + parameters.logits_processor = 0; + parameters.eos_token_id = eos_token_id; + parameters.min_length = min_length; + parameters.no_repeat_ngram_size = 0; + parameters.repetition_penalty = 1.0f; // 1.0 means no penalty, so that processor is skipped + parameters.temperature = 0.0f; // <= 0 skips the temperature processor + parameters.batch_size = kBatchBeamSize; // GreedySearchParameters::BatchBeamSize() == batch_size + parameters.vocab_size = kVocabSize; + return parameters; +} + +} // namespace + +// Backstop: SetScore ignores a negative token id, which is the "no eos" sentinel. This guards the +// runtime path against indexing scores with a negative token id even if a processor is reached. +TEST(MinLengthLogitsProcessorTest, SetScoreIgnoresNegativeTokenId) { + std::vector scores(kBatchBeamSize * kVocabSize, 1.0f); + gsl::span scores_span(scores); + NextTokenScores next_token_scores{scores_span, kBatchBeamSize, kVocabSize}; + + next_token_scores.SetScore(/*token_id=*/-1, std::numeric_limits::lowest()); + + for (float value : scores) { + EXPECT_FLOAT_EQ(value, 1.0f); + } +} + +// With a negative "no eos" sentinel there is no token to demote, so the Init guard skips constructing +// the MinLength processor as a guaranteed no-op (a defensive/performance skip). The scores would be +// unchanged here regardless, because SetScore also ignores a negative token id; the enforcement path +// itself is covered by the eos >= 0 positive-control test below. +TEST(MinLengthLogitsProcessorTest, ListInitSkipsProcessorForNegativeEosTokenId) { + GreedySearchParameters parameters = MakeMinLengthOnlyParameters(/*min_length=*/5, /*eos_token_id=*/-1); + LogitsProcessorList processor_list; + processor_list.Init(parameters); + + std::vector scores(kBatchBeamSize * kVocabSize, 1.0f); + gsl::span scores_span(scores); + FixedLengthSequences sequences(/*sequence_length=*/1); // below min_length + processor_list.Process(&sequences, scores_span, /*step=*/1); + + for (float value : scores) { + EXPECT_FLOAT_EQ(value, 1.0f); + } +} + +// Enforcement path: with a valid eos_token_id the Init call site adds the MinLength processor, so a +// below-min-length run demotes the eos score. This is the discriminating test for the eos >= 0 branch +// of the guard and for the min-length enforcement behavior. +TEST(MinLengthLogitsProcessorTest, ListInitDemotesEosBelowMinLength) { + constexpr int kEosTokenId = 2; + GreedySearchParameters parameters = MakeMinLengthOnlyParameters(/*min_length=*/5, kEosTokenId); + LogitsProcessorList processor_list; + processor_list.Init(parameters); + + std::vector scores(kBatchBeamSize * kVocabSize, 1.0f); + gsl::span scores_span(scores); + FixedLengthSequences sequences(/*sequence_length=*/1); // below min_length + processor_list.Process(&sequences, scores_span, /*step=*/1); + + const float lowest = std::numeric_limits::lowest(); + for (int beam = 0; beam < kBatchBeamSize; ++beam) { + for (int token = 0; token < kVocabSize; ++token) { + const float value = scores[static_cast(beam) * kVocabSize + token]; + if (token == kEosTokenId) { + EXPECT_FLOAT_EQ(value, lowest); + } else { + EXPECT_FLOAT_EQ(value, 1.0f); + } + } + } +} + +// Once the sequence reaches min_length, a valid eos score is left untouched: min_length is no longer +// enforced. This locks in the boundary behavior and confirms the guard change is limited to the +// negative-sentinel case (no regression for valid eos ids). +TEST(MinLengthLogitsProcessorTest, ListInitLeavesScoresUnchangedAtMinLength) { + GreedySearchParameters parameters = MakeMinLengthOnlyParameters(/*min_length=*/5, /*eos_token_id=*/2); + LogitsProcessorList processor_list; + processor_list.Init(parameters); + + std::vector scores(kBatchBeamSize * kVocabSize, 1.0f); + gsl::span scores_span(scores); + FixedLengthSequences sequences(/*sequence_length=*/5); // == min_length: no demotion expected + processor_list.Process(&sequences, scores_span, /*step=*/1); + + for (float value : scores) { + EXPECT_FLOAT_EQ(value, 1.0f); + } +} + +} // namespace test +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime From 5dc8d4c6353ef27a0e9aace8dd6e3e6cede479ec Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 8 Jul 2026 10:06:30 -0700 Subject: [PATCH 2/5] Speed up CUDA build by splitting SM80 MoE GEMM generated files (#29614) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description The SM80 fused MoE GEMM launchers were generated as two monolithic translation units — `fused_moe_gemm_sm80_bf16.generated.cu` and `fused_moe_gemm_sm80_f16.generated.cu` — each packing **30 CUTLASS kernel instantiations** (5 tile shapes × 3 stages × 2 epilogues) into a single `.cu` file. nvcc's `--threads` flag only parallelizes across GPU architectures (`80;90`) within a file; it does **not** parallelize template instantiations inside one translation unit. As a result each of these two objects was a long, serial compile on the build's critical path (each ~5–6 min in a Debug `80;90` build), and Ninja could only overlap the two of them. This PR updates the kernel generator to split the SM80 instantiations by tile shape, producing **10 smaller files** (5 tile shapes × 2 dtypes, 6 kernels each) that Ninja compiles in parallel. ### Key Changes - `generate_moe_kernels.py`: the SM80 loop now groups instantiations by `(element_type, tile_shape)` and emits one file per group, named `fused_moe_gemm_sm80_{dtype}_m{M}_n{N}_k{K}.generated.cu`. Added `glob`-based cleanup that removes stale SM80 generated files (the old monolithic per-dtype files and any files from a previous tile configuration) so they are not compiled. - Regenerated the launcher directory: removed the 2 monolithic files and added the 10 split files. Total number of compiled kernel instantiations is unchanged. cmake already discovers these via `GLOB_RECURSE ... CONFIGURE_DEPENDS`, so no cmake changes are needed — the new files are picked up automatically on the next configure. ### Testing - Verified generator idempotency (re-running reports all files up to date) and stale-file cleanup. - Built the new objects in a Debug `CMAKE_CUDA_ARCHITECTURES="80;90"` build: - one 6-kernel split object compiles in ~67s; - all 10 split objects build in ~182s wall-clock in parallel. - Previously the two 30-kernel objects were ~5× the per-file cost each and only 2-way parallel, so this substantially shortens the MoE portion of the critical path on multi-core machines. Total compiled kernels (and thus runtime coverage) is unchanged. To regenerate after future tile-shape changes: ``` python onnxruntime/contrib_ops/cuda/llm/generate_moe_kernels.py -a "80;90" -o onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers ``` --- .../cuda/llm/generate_moe_kernels.py | 29 +- .../fused_moe_gemm_sm80_bf16.generated.cu | 141 ---------- ..._gemm_sm80_bf16_m128_n128_k64.generated.cu | 45 +++ ...e_gemm_sm80_bf16_m16_n128_k64.generated.cu | 45 +++ ...e_gemm_sm80_bf16_m16_n256_k64.generated.cu | 45 +++ ...e_gemm_sm80_bf16_m32_n128_k64.generated.cu | 45 +++ ...e_gemm_sm80_bf16_m64_n128_k64.generated.cu | 45 +++ .../fused_moe_gemm_sm80_f16.generated.cu | 260 ------------------ ...e_gemm_sm80_f16_m128_n128_k64.generated.cu | 68 +++++ ...oe_gemm_sm80_f16_m16_n128_k64.generated.cu | 68 +++++ ...oe_gemm_sm80_f16_m16_n256_k64.generated.cu | 68 +++++ ...oe_gemm_sm80_f16_m32_n128_k64.generated.cu | 68 +++++ ...oe_gemm_sm80_f16_m64_n128_k64.generated.cu | 68 +++++ 13 files changed, 586 insertions(+), 409 deletions(-) delete mode 100644 onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16.generated.cu create mode 100644 onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m128_n128_k64.generated.cu create mode 100644 onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m16_n128_k64.generated.cu create mode 100644 onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m16_n256_k64.generated.cu create mode 100644 onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m32_n128_k64.generated.cu create mode 100644 onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m64_n128_k64.generated.cu delete mode 100644 onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16.generated.cu create mode 100644 onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m128_n128_k64.generated.cu create mode 100644 onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m16_n128_k64.generated.cu create mode 100644 onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m16_n256_k64.generated.cu create mode 100644 onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m32_n128_k64.generated.cu create mode 100644 onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m64_n128_k64.generated.cu diff --git a/onnxruntime/contrib_ops/cuda/llm/generate_moe_kernels.py b/onnxruntime/contrib_ops/cuda/llm/generate_moe_kernels.py index c7f9788ad786f..b32676dad29d2 100644 --- a/onnxruntime/contrib_ops/cuda/llm/generate_moe_kernels.py +++ b/onnxruntime/contrib_ops/cuda/llm/generate_moe_kernels.py @@ -6,6 +6,7 @@ # python generate_moe_kernels.py -a "80;90" -o ./moe_gemm/launchers import argparse +import glob import os from itertools import product @@ -333,19 +334,31 @@ def has_arch(sm): if has_arch(80) or has_arch(90): # SM90 also uses SM80 kernels for non-TMA path operations = generate_sm80_moe_operations() - # Group by element type for separate files to reduce compile time + # Split by (element type, tile shape) into separate files. Each SM80 template + # instantiation is a full CUTLASS GEMM kernel that nvcc compiles serially within + # a translation unit, so packing all of them into one file per dtype creates a + # long single-file critical path. Emitting one file per tile shape lets the build + # system (Ninja) compile these kernels in parallel and shrinks the slowest object. groups = {} for op in operations: - key = op["element_type"] - if key not in groups: - groups[key] = [] - groups[key].append(op) - - for dtype, ops in groups.items(): - output_file = os.path.join(output_dir, f"fused_moe_gemm_sm80_{dtype}.generated.cu") + key = (op["element_type"], op["tile_m"], op["tile_n"], op["tile_k"]) + groups.setdefault(key, []).append(op) + + current_files = set() + for (dtype, tile_m, tile_n, tile_k), ops in groups.items(): + file_name = f"fused_moe_gemm_sm80_{dtype}_m{tile_m}_n{tile_n}_k{tile_k}.generated.cu" + current_files.add(file_name) + output_file = os.path.join(output_dir, file_name) content = get_sm80_file_content(ops, 80) write_file(content, output_file) + # Remove stale SM80 generated files (e.g. the old monolithic per-dtype files or + # split files from a previous tile configuration) so they are not compiled. + for stale in glob.glob(os.path.join(output_dir, "fused_moe_gemm_sm80_*.generated.cu")): + if os.path.basename(stale) not in current_files: + os.remove(stale) + print(f"Removed stale {stale}") + # Generate SM90 TMA Warp Specialized Grouped GEMM kernels if has_arch(90): operations = generate_sm90_tma_ws_operations() diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16.generated.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16.generated.cu deleted file mode 100644 index ff06ab76912e2..0000000000000 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16.generated.cu +++ /dev/null @@ -1,141 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT License. - * - * Auto-generated MoE GEMM kernel instantiations for SM80. - * DO NOT EDIT MANUALLY. - */ - -#ifndef EXCLUDE_SM_80 -#include "contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_launcher_sm80.inl" - -namespace onnxruntime::llm::kernels::cutlass_kernels { - -#ifdef ENABLE_BF16 -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -#else -// BF16 not enabled, only instantiate FP16 variants - -#endif - -} // namespace onnxruntime::llm::kernels::cutlass_kernels -#endif // EXCLUDE_SM_80 diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m128_n128_k64.generated.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m128_n128_k64.generated.cu new file mode 100644 index 0000000000000..15c589b89a170 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m128_n128_k64.generated.cu @@ -0,0 +1,45 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + * + * Auto-generated MoE GEMM kernel instantiations for SM80. + * DO NOT EDIT MANUALLY. + */ + +#ifndef EXCLUDE_SM_80 +#include "contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_launcher_sm80.inl" + +namespace onnxruntime::llm::kernels::cutlass_kernels { + +#ifdef ENABLE_BF16 +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#else +// BF16 not enabled, only instantiate FP16 variants + +#endif + +} // namespace onnxruntime::llm::kernels::cutlass_kernels +#endif // EXCLUDE_SM_80 diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m16_n128_k64.generated.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m16_n128_k64.generated.cu new file mode 100644 index 0000000000000..610104e822c69 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m16_n128_k64.generated.cu @@ -0,0 +1,45 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + * + * Auto-generated MoE GEMM kernel instantiations for SM80. + * DO NOT EDIT MANUALLY. + */ + +#ifndef EXCLUDE_SM_80 +#include "contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_launcher_sm80.inl" + +namespace onnxruntime::llm::kernels::cutlass_kernels { + +#ifdef ENABLE_BF16 +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#else +// BF16 not enabled, only instantiate FP16 variants + +#endif + +} // namespace onnxruntime::llm::kernels::cutlass_kernels +#endif // EXCLUDE_SM_80 diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m16_n256_k64.generated.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m16_n256_k64.generated.cu new file mode 100644 index 0000000000000..f1a071f8177db --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m16_n256_k64.generated.cu @@ -0,0 +1,45 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + * + * Auto-generated MoE GEMM kernel instantiations for SM80. + * DO NOT EDIT MANUALLY. + */ + +#ifndef EXCLUDE_SM_80 +#include "contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_launcher_sm80.inl" + +namespace onnxruntime::llm::kernels::cutlass_kernels { + +#ifdef ENABLE_BF16 +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#else +// BF16 not enabled, only instantiate FP16 variants + +#endif + +} // namespace onnxruntime::llm::kernels::cutlass_kernels +#endif // EXCLUDE_SM_80 diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m32_n128_k64.generated.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m32_n128_k64.generated.cu new file mode 100644 index 0000000000000..43837b91c4cac --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m32_n128_k64.generated.cu @@ -0,0 +1,45 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + * + * Auto-generated MoE GEMM kernel instantiations for SM80. + * DO NOT EDIT MANUALLY. + */ + +#ifndef EXCLUDE_SM_80 +#include "contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_launcher_sm80.inl" + +namespace onnxruntime::llm::kernels::cutlass_kernels { + +#ifdef ENABLE_BF16 +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#else +// BF16 not enabled, only instantiate FP16 variants + +#endif + +} // namespace onnxruntime::llm::kernels::cutlass_kernels +#endif // EXCLUDE_SM_80 diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m64_n128_k64.generated.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m64_n128_k64.generated.cu new file mode 100644 index 0000000000000..1f7a97b69f32b --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_bf16_m64_n128_k64.generated.cu @@ -0,0 +1,45 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + * + * Auto-generated MoE GEMM kernel instantiations for SM80. + * DO NOT EDIT MANUALLY. + */ + +#ifndef EXCLUDE_SM_80 +#include "contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_launcher_sm80.inl" + +namespace onnxruntime::llm::kernels::cutlass_kernels { + +#ifdef ENABLE_BF16 +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, bool, cutlass::bfloat16_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#else +// BF16 not enabled, only instantiate FP16 variants + +#endif + +} // namespace onnxruntime::llm::kernels::cutlass_kernels +#endif // EXCLUDE_SM_80 diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16.generated.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16.generated.cu deleted file mode 100644 index 1c950e47dbbae..0000000000000 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16.generated.cu +++ /dev/null @@ -1,260 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT License. - * - * Auto-generated MoE GEMM kernel instantiations for SM80. - * DO NOT EDIT MANUALLY. - */ - -#ifndef EXCLUDE_SM_80 -#include "contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_launcher_sm80.inl" - -namespace onnxruntime::llm::kernels::cutlass_kernels { - -#ifdef ENABLE_BF16 -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -#else -// BF16 not enabled, only instantiate FP16 variants -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -template void sm80_generic_fused_moe_gemm_kernelLauncher( - cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, - int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); - -#endif - -} // namespace onnxruntime::llm::kernels::cutlass_kernels -#endif // EXCLUDE_SM_80 diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m128_n128_k64.generated.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m128_n128_k64.generated.cu new file mode 100644 index 0000000000000..a917305c788f4 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m128_n128_k64.generated.cu @@ -0,0 +1,68 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + * + * Auto-generated MoE GEMM kernel instantiations for SM80. + * DO NOT EDIT MANUALLY. + */ + +#ifndef EXCLUDE_SM_80 +#include "contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_launcher_sm80.inl" + +namespace onnxruntime::llm::kernels::cutlass_kernels { + +#ifdef ENABLE_BF16 +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#else +// BF16 not enabled, only instantiate FP16 variants +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#endif + +} // namespace onnxruntime::llm::kernels::cutlass_kernels +#endif // EXCLUDE_SM_80 diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m16_n128_k64.generated.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m16_n128_k64.generated.cu new file mode 100644 index 0000000000000..dd67705f8064b --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m16_n128_k64.generated.cu @@ -0,0 +1,68 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + * + * Auto-generated MoE GEMM kernel instantiations for SM80. + * DO NOT EDIT MANUALLY. + */ + +#ifndef EXCLUDE_SM_80 +#include "contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_launcher_sm80.inl" + +namespace onnxruntime::llm::kernels::cutlass_kernels { + +#ifdef ENABLE_BF16 +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#else +// BF16 not enabled, only instantiate FP16 variants +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#endif + +} // namespace onnxruntime::llm::kernels::cutlass_kernels +#endif // EXCLUDE_SM_80 diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m16_n256_k64.generated.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m16_n256_k64.generated.cu new file mode 100644 index 0000000000000..58ce8c4c631ea --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m16_n256_k64.generated.cu @@ -0,0 +1,68 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + * + * Auto-generated MoE GEMM kernel instantiations for SM80. + * DO NOT EDIT MANUALLY. + */ + +#ifndef EXCLUDE_SM_80 +#include "contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_launcher_sm80.inl" + +namespace onnxruntime::llm::kernels::cutlass_kernels { + +#ifdef ENABLE_BF16 +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#else +// BF16 not enabled, only instantiate FP16 variants +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#endif + +} // namespace onnxruntime::llm::kernels::cutlass_kernels +#endif // EXCLUDE_SM_80 diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m32_n128_k64.generated.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m32_n128_k64.generated.cu new file mode 100644 index 0000000000000..b865cbcf636a0 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m32_n128_k64.generated.cu @@ -0,0 +1,68 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + * + * Auto-generated MoE GEMM kernel instantiations for SM80. + * DO NOT EDIT MANUALLY. + */ + +#ifndef EXCLUDE_SM_80 +#include "contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_launcher_sm80.inl" + +namespace onnxruntime::llm::kernels::cutlass_kernels { + +#ifdef ENABLE_BF16 +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#else +// BF16 not enabled, only instantiate FP16 variants +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#endif + +} // namespace onnxruntime::llm::kernels::cutlass_kernels +#endif // EXCLUDE_SM_80 diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m64_n128_k64.generated.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m64_n128_k64.generated.cu new file mode 100644 index 0000000000000..24765bf9f3333 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_sm80_f16_m64_n128_k64.generated.cu @@ -0,0 +1,68 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + * + * Auto-generated MoE GEMM kernel instantiations for SM80. + * DO NOT EDIT MANUALLY. + */ + +#ifndef EXCLUDE_SM_80 +#include "contrib_ops/cuda/llm/moe_gemm/launchers/fused_moe_gemm_launcher_sm80.inl" + +namespace onnxruntime::llm::kernels::cutlass_kernels { + +#ifdef ENABLE_BF16 +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#else +// BF16 not enabled, only instantiate FP16 variants +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +template void sm80_generic_fused_moe_gemm_kernelLauncher( + cutlass::half_t const*, cutlass::half_t const*, cutlass::half_t const*, bool, cutlass::half_t*, + int64_t const*, int64_t, int64_t, int64_t, int, int, cudaStream_t, int*); + +#endif + +} // namespace onnxruntime::llm::kernels::cutlass_kernels +#endif // EXCLUDE_SM_80 From 682936c1ac80e4addc84f231fe062a4ac43620e3 Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Wed, 8 Jul 2026 10:14:09 -0700 Subject: [PATCH 3/5] Validate DynamicQuantizeLSTM recurrence quantization parameter shapes (#29254) ### Description The `R_zero_point` / `R_scale` (recurrence) quantization-parameter shape validation in `DynamicQuantizeLSTM` was inadvertently checking the `W` (input) quantization parameters instead of the `R` ones. This change validates the `R` parameters' own shapes symmetrically with `W`, so malformed recurrence quantization parameters are rejected with a clear error. ### Changes - Fix the shape checks so `R_zero_point` and `R_scale` are validated against the `R` tensor's expected shape (previously bound to the `W` tensor). - Add two expect-failure unit tests covering inconsistent recurrence zero-point and scale shapes. ### Motivation Improves input validation and error diagnostics for malformed `DynamicQuantizeLSTM` recurrence quantization parameters. CPU-only; no behavior change for valid inputs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../test/contrib_ops/quantize_lstm_op_test.cc | 178 ++++++++++++++++-- 1 file changed, 158 insertions(+), 20 deletions(-) diff --git a/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc b/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc index 93a63ce19eb03..f2ba2eb57b819 100644 --- a/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc +++ b/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc @@ -105,7 +105,8 @@ static void ComputeRefOutput(std::vector& Y_data, const std::vector initial_c_data, const std::string& direction, const std::vector& activations, - bool per_channel) { + bool w_per_channel, + bool r_per_channel) { OpTester test("LSTM", 7 /*opset_version*/, onnxruntime::kOnnxDomain /*domain*/, false /*verify_output*/); test.AddAttribute>("activations", activations); @@ -120,8 +121,8 @@ static void ComputeRefOutput(std::vector& Y_data, std::vector R_dims = {num_directions, 4 * hidden_size, hidden_size}; test.AddInput("X", X_dims, ApplyQDQ(X_data, 1)); - test.AddInput("W", W_dims, ApplyQDQ(W_data, per_channel ? num_directions * 4 * hidden_size : num_directions, per_channel)); - test.AddInput("R", R_dims, ApplyQDQ(R_data, per_channel ? num_directions * 4 * hidden_size : num_directions, per_channel)); + test.AddInput("W", W_dims, ApplyQDQ(W_data, w_per_channel ? num_directions * 4 * hidden_size : num_directions, w_per_channel)); + test.AddInput("R", R_dims, ApplyQDQ(R_data, r_per_channel ? num_directions * 4 * hidden_size : num_directions, r_per_channel)); if (B_data) { std::vector B_dims = {num_directions, 8 * hidden_size}; @@ -186,7 +187,8 @@ static void RunQuantLSTM(int64_t input_size, bool has_P, bool is_initializer_W, bool is_initializer_R, - bool per_channel, + bool w_per_channel, + bool r_per_channel, const std::string& direction) { OpTester test("DynamicQuantizeLSTM", 1 /*opset_version*/, onnxruntime::kMSDomain /*domain*/); @@ -219,7 +221,7 @@ static void RunQuantLSTM(int64_t input_size, std::vector w_scale; std::vector w_zp; std::vector w_quant; - QuantizeWeight(w_quant, w_scale, w_zp, W_data, num_directions, 4 * hidden_size, input_size, per_channel); + QuantizeWeight(w_quant, w_scale, w_zp, W_data, num_directions, 4 * hidden_size, input_size, w_per_channel); test.AddInput("W", W_dims, w_quant, is_initializer_W); // R @@ -229,7 +231,7 @@ static void RunQuantLSTM(int64_t input_size, std::vector r_scale; std::vector r_zp; std::vector r_quant; - QuantizeWeight(r_quant, r_scale, r_zp, R_data, num_directions, 4 * hidden_size, hidden_size, per_channel); + QuantizeWeight(r_quant, r_scale, r_zp, R_data, num_directions, 4 * hidden_size, hidden_size, r_per_channel); test.AddInput("R", R_dims, r_quant, is_initializer_R); std::vector B_data; @@ -266,11 +268,11 @@ static void RunQuantLSTM(int64_t input_size, std::vector per_tensor_dims = {num_directions}; std::vector per_channel_dims = {num_directions, 4 * hidden_size}; - test.AddInput("W_scale", per_channel ? per_channel_dims : per_tensor_dims, w_scale); - test.AddInput("W_zero_point", per_channel ? per_channel_dims : per_tensor_dims, w_zp); + test.AddInput("W_scale", w_per_channel ? per_channel_dims : per_tensor_dims, w_scale); + test.AddInput("W_zero_point", w_per_channel ? per_channel_dims : per_tensor_dims, w_zp); - test.AddInput("R_scale", per_channel ? per_channel_dims : per_tensor_dims, r_scale); - test.AddInput("R_zero_point", per_channel ? per_channel_dims : per_tensor_dims, r_zp); + test.AddInput("R_scale", r_per_channel ? per_channel_dims : per_tensor_dims, r_scale); + test.AddInput("R_zero_point", r_per_channel ? per_channel_dims : per_tensor_dims, r_zp); std::vector Y_data; std::vector Y_h_data; @@ -281,7 +283,7 @@ static void RunQuantLSTM(int64_t input_size, has_bias ? &B_data : nullptr, has_P ? &P_data : nullptr, initial_h_data, initial_c_data, - direction, activations, per_channel); + direction, activations, w_per_channel, r_per_channel); std::vector Y_dims = {seq_len, num_directions, batch_size, hidden_size}; test.AddOutput("Y", Y_dims, Y_data); @@ -305,49 +307,49 @@ static void RunQuantLSTM(int64_t input_size, RunQuantLSTM(input_size, batch_size, hidden_size, false /*has_bias*/, false /*has_P*/, false /*is_initializer_W*/, false /*is_initializer_R*/, - per_channel, "forward"); + per_channel /*w_per_channel*/, per_channel /*r_per_channel*/, "forward"); // bias + P: 0, prepacking: 0, bidirectional: 1 RunQuantLSTM(input_size, batch_size, hidden_size, false /*has_bias*/, false /*has_P*/, false /*is_initializer_W*/, false /*is_initializer_R*/, - per_channel, "bidirectional"); + per_channel /*w_per_channel*/, per_channel /*r_per_channel*/, "bidirectional"); // bias + P: 0, prepacking: 1, bidirectional: 0 RunQuantLSTM(input_size, batch_size, hidden_size, false /*has_bias*/, false /*has_P*/, true /*is_initializer_W*/, true /*is_initializer_R*/, - per_channel, "forward"); + per_channel /*w_per_channel*/, per_channel /*r_per_channel*/, "forward"); // bias + P: 0, prepacking: 1, bidirectional: 1 RunQuantLSTM(input_size, batch_size, hidden_size, false /*has_bias*/, false /*has_P*/, true /*is_initializer_W*/, true /*is_initializer_R*/, - per_channel, "bidirectional"); + per_channel /*w_per_channel*/, per_channel /*r_per_channel*/, "bidirectional"); // bias + P: 1, prepacking: 0, bidirectional: 0 RunQuantLSTM(input_size, batch_size, hidden_size, true /*has_bias*/, true /*has_P*/, false /*is_initializer_W*/, false /*is_initializer_R*/, - per_channel, "forward"); + per_channel /*w_per_channel*/, per_channel /*r_per_channel*/, "forward"); // bias + P: 1, prepacking: 0, bidirectional: 1 RunQuantLSTM(input_size, batch_size, hidden_size, true /*has_bias*/, true /*has_P*/, false /*is_initializer_W*/, false /*is_initializer_R*/, - per_channel, "bidirectional"); + per_channel /*w_per_channel*/, per_channel /*r_per_channel*/, "bidirectional"); // bias + P: 1, prepacking: 1, bidirectional: 0 RunQuantLSTM(input_size, batch_size, hidden_size, true /*has_bias*/, true /*has_P*/, true /*is_initializer_W*/, true /*is_initializer_R*/, - per_channel, "forward"); + per_channel /*w_per_channel*/, per_channel /*r_per_channel*/, "forward"); // bias + P: 1, prepacking: 1, bidirectional: 1 RunQuantLSTM(input_size, batch_size, hidden_size, true /*has_bias*/, true /*has_P*/, true /*is_initializer_W*/, true /*is_initializer_R*/, - per_channel, "bidirectional"); + per_channel /*w_per_channel*/, per_channel /*r_per_channel*/, "bidirectional"); } TEST(DynamicQuantLSTMTest, SmallSize) { @@ -442,7 +444,7 @@ TEST(DynamicQuantLSTMTest, SharedPrepackedWeights) { nullptr, nullptr, initial_h_data, initial_c_data, - "forward", activations, false); + "forward", activations, false, false); std::vector Y_dims = {seq_len, num_directions, batch_size, hidden_size}; test.AddOutput("Y", Y_dims, Y_data); @@ -525,5 +527,141 @@ TEST(DynamicQuantLSTMTest, SharedPrepackedWeights) { } #endif +// Builds a minimal DynamicQuantizeLSTM and runs it expecting the kernel's input validation to reject +// the recurrence quantization parameters. The caller supplies the R_scale and R_zero_point shapes +// (and optionally the R_zero_point values) so a shape that is inconsistent with R (e.g. first dim +// != num_directions, or a per-channel second dim != 4*hidden_size), or a per-channel zero point that +// violates the constant/zero requirement, can be exercised. The recurrence parameters must be +// validated symmetrically with the input (W) ones. num_directions selects forward (1) or +// bidirectional (2). +static void RunQuantLSTMExpectInvalidRecurrenceQuantParam(const std::vector& r_scale_dims, + const std::vector& r_zp_dims, + const std::string& expected_error, + int64_t num_directions = 1, + const std::vector& r_zp_values = {}) { + OpTester test("DynamicQuantizeLSTM", 1 /*opset_version*/, onnxruntime::kMSDomain /*domain*/); + + constexpr int64_t input_size = 2; + constexpr int64_t hidden_size = 2; + constexpr int64_t batch_size = 1; + constexpr int64_t seq_len = 1; + + auto num_elements = [](const std::vector& dims) { + int64_t count = 1; + for (int64_t dim : dims) { + count *= dim; + } + return static_cast(count); + }; + + const std::string direction = (num_directions == 2) ? "bidirectional" : "forward"; + std::vector activations = {"sigmoid", "tanh", "tanh"}; + if (num_directions == 2) { + activations = {"sigmoid", "tanh", "tanh", "sigmoid", "tanh", "tanh"}; + } + + test.AddAttribute>("activations", activations); + test.AddAttribute("direction", direction); + test.AddAttribute("hidden_size", hidden_size); + test.AddAttribute("input_forget", static_cast(0)); + + // X: [seq_length, batch_size, input_size] + test.AddInput("X", {seq_len, batch_size, input_size}, + std::vector(num_elements({seq_len, batch_size, input_size}), 0.0f)); + + // W / R quantized weight values are irrelevant: validation fails before any dequantization. + test.AddInput("W", {num_directions, input_size, 4 * hidden_size}, + std::vector(num_elements({num_directions, input_size, 4 * hidden_size}), 0)); + test.AddInput("R", {num_directions, hidden_size, 4 * hidden_size}, + std::vector(num_elements({num_directions, hidden_size, 4 * hidden_size}), 0)); + + test.AddOptionalInputEdge(); // B + test.AddOptionalInputEdge(); // sequence_lens + test.AddInput("initial_h", {num_directions, batch_size, hidden_size}, + std::vector(num_elements({num_directions, batch_size, hidden_size}), 0.0f)); + test.AddInput("initial_c", {num_directions, batch_size, hidden_size}, + std::vector(num_elements({num_directions, batch_size, hidden_size}), 0.0f)); + test.AddOptionalInputEdge(); // P + + // Valid per-tensor quantization parameters for the input weights. + test.AddInput("W_scale", {num_directions}, std::vector(num_elements({num_directions}), 1.0f)); + test.AddInput("W_zero_point", {num_directions}, std::vector(num_elements({num_directions}), 0)); + + // Recurrence parameters with caller-supplied (possibly inconsistent) shapes and zero-point values. + const std::vector r_zp = r_zp_values.empty() + ? std::vector(num_elements(r_zp_dims), 0) + : r_zp_values; + test.AddInput("R_scale", r_scale_dims, std::vector(num_elements(r_scale_dims), 1.0f)); + test.AddInput("R_zero_point", r_zp_dims, r_zp); + + // Placeholder outputs (not validated: the run fails during input validation). + test.AddOutput("Y", {seq_len, num_directions, batch_size, hidden_size}, + std::vector(num_elements({seq_len, num_directions, batch_size, hidden_size}), 0.0f)); + test.AddOutput("Y_h", {num_directions, batch_size, hidden_size}, + std::vector(num_elements({num_directions, batch_size, hidden_size}), 0.0f)); + test.AddOutput("Y_c", {num_directions, batch_size, hidden_size}, + std::vector(num_elements({num_directions, batch_size, hidden_size}), 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, expected_error); +} + +TEST(DynamicQuantLSTMTest, RejectsInconsistentRecurrenceZeroPointShape) { + // R_zero_point's first dim must equal num_directions (1); {2} is inconsistent and must be rejected + // rather than silently validated against the input zero point's shape. + RunQuantLSTMExpectInvalidRecurrenceQuantParam(/*r_scale_dims=*/{1}, /*r_zp_dims=*/{2}, + "Input R_zero_point must have shape"); +} + +TEST(DynamicQuantLSTMTest, RejectsInconsistentRecurrenceScaleShape) { + // R_scale's first dim must equal num_directions (1); {2} is inconsistent and must be rejected + // rather than silently validated against the input scale's shape. + RunQuantLSTMExpectInvalidRecurrenceQuantParam(/*r_scale_dims=*/{2}, /*r_zp_dims=*/{1}, + "Input R_scale must have shape"); +} + +TEST(DynamicQuantLSTMTest, RejectsInconsistentPerChannelRecurrenceShape) { + // Per-channel (2D) recurrence quantization parameters must have second dim 4*hidden_size (== 8 here). + // {1, 3} has the correct first dim (num_directions) but a wrong per-channel dim, exercising the + // per-channel branch of the shape check (distinct from the wrong-first-dim {2} cases above). + RunQuantLSTMExpectInvalidRecurrenceQuantParam(/*r_scale_dims=*/{1, 8}, /*r_zp_dims=*/{1, 3}, + "Input R_zero_point must have shape"); + RunQuantLSTMExpectInvalidRecurrenceQuantParam(/*r_scale_dims=*/{1, 3}, /*r_zp_dims=*/{1, 8}, + "Input R_scale must have shape"); +} + +TEST(DynamicQuantLSTMTest, RejectsInconsistentRecurrenceQuantParamBidirectional) { + // With two directions the recurrence quantization parameters' first dim must equal num_directions + // (2); a {1} shape is inconsistent and must be rejected in the bidirectional case as well. + RunQuantLSTMExpectInvalidRecurrenceQuantParam(/*r_scale_dims=*/{2}, /*r_zp_dims=*/{1}, + "Input R_zero_point must have shape", + /*num_directions=*/2); +} + +TEST(DynamicQuantLSTMTest, RejectsNonConstantPerChannelRecurrenceZeroPoint) { + // A per-channel R_zero_point with the correct shape {num_directions, 4*hidden_size} (== {1, 8} here) + // passes the shape check and reaches the per-element zero-point validation. Unsigned recurrence + // weights require a constant zero point, so non-constant values must be rejected. This exercises the + // zero-point iteration against R's own shape and values. + RunQuantLSTMExpectInvalidRecurrenceQuantParam( + /*r_scale_dims=*/{1, 8}, /*r_zp_dims=*/{1, 8}, + "RecurrentWeight point must be constant", + /*num_directions=*/1, + /*r_zp_values=*/{0, 1, 0, 0, 0, 0, 0, 0}); +} + +TEST(DynamicQuantLSTMTest, AcceptsMixedGranularityRecurrenceQuantParam) { + // Mixed quantization granularity between the input (W) and recurrence (R) weights is valid and must + // run to completion producing correct output. Exercise both orderings: per-tensor W with per-channel + // R, and per-channel W with per-tensor R. + RunQuantLSTM(2, 1, 16, + false /*has_bias*/, false /*has_P*/, + false /*is_initializer_W*/, false /*is_initializer_R*/, + false /*w_per_channel*/, true /*r_per_channel*/, "forward"); + RunQuantLSTM(2, 1, 16, + false /*has_bias*/, false /*has_P*/, + false /*is_initializer_W*/, false /*is_initializer_R*/, + true /*w_per_channel*/, false /*r_per_channel*/, "forward"); +} + } // namespace test } // namespace onnxruntime From bcc9d30b8615d5890a8847cd43e2df13cfcbae4c Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Wed, 8 Jul 2026 10:31:11 -0700 Subject: [PATCH 4/5] Handle non-finite ROI coordinates in CropAndResize (#29605) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description The `CropAndResize` CPU contrib op (`com.microsoft::CropAndResize`) computes bilinear/nearest interpolation indices from the ROI box coordinates. The bounds guards were written as `if (in_y < 0 || in_y > height - 1)` / `if (in_x < 0 || in_x > width - 1)`. Because every comparison involving a NaN is false, a non-finite (NaN or ±inf) ROI coordinate slipped past both comparisons, skipped the extrapolation `continue`, and reached the integer index computation (`(int)floorf(NaN)`) with an invalid value. This is a robustness/correctness gap: ROI coordinates are a runtime input, and non-finite values are not handled gracefully. ### Fix - **NaN-safe bounds guards.** Rewrite both guards into negated-conjunction form: `if (!(in_y >= 0 && in_y <= height - 1))` / `if (!(in_x >= 0 && in_x <= width - 1))`. This is logically identical for all finite coordinates, but is `true` for NaN/±inf, so non-finite coordinates now take the extrapolation branch and are filled with `extrapolation_value` (matching the documented behavior for out-of-range coordinates). - **crop_size validation.** Add a Status-based `ORT_RETURN_IF_NOT(crop_height > 0 && crop_width > 0, ...)` check in `Compute` so non-positive crop sizes are rejected with a clear error rather than producing degenerate work. - **Index arithmetic hardening.** Use `SafeInt` for the interpolation index computation, which allows the now-unnecessary MSVC 26451 (arithmetic-overflow) suppression pragma to be removed. No behavior change for valid finite ROIs. CPU-only — there is no CUDA `CropAndResize` kernel. ### Tests Added to `onnxruntime/test/contrib_ops/crop_and_resize_op_test.cc`: - NaN ROI coordinate → extrapolation (bilinear height, bilinear width, and nearest). - ±inf ROI coordinate → extrapolation. - Finite boundary values (exact `[0, 1]` identity crop and just-outside `1.0001`) — no regression. - Non-positive `crop_size` (`{0, 2}` and `{-1, 2}`) rejected. - Out-of-range batch index rejected. Run with: ``` onnxruntime_provider_test --gtest_filter='CropAndResizeTest.*' ``` All 10 CropAndResize tests pass (5 new + 5 pre-existing, no regression). An AddressSanitizer build can additionally corroborate the index handling in CI. ### Motivation and Context Handles non-finite ROI coordinates gracefully and makes the CropAndResize index arithmetic robust, consistent with how out-of-range coordinates are already treated (extrapolation). Signed-off-by: Tita Wang Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cpu/crop_and_resize.cc | 34 ++--- .../contrib_ops/crop_and_resize_op_test.cc | 119 ++++++++++++++++++ 2 files changed, 139 insertions(+), 14 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/crop_and_resize.cc b/onnxruntime/contrib_ops/cpu/crop_and_resize.cc index ec24aa8e9c52c..47f164fde8322 100644 --- a/onnxruntime/contrib_ops/cpu/crop_and_resize.cc +++ b/onnxruntime/contrib_ops/cpu/crop_and_resize.cc @@ -17,16 +17,12 @@ limitations under the License. #include "contrib_ops/cpu/crop_and_resize.h" #include +#include "core/common/safeint.h" #include "core/util/math_cpuonly.h" #include "core/common/common.h" #include "core/framework/tensor.h" #include "core/platform/threadpool.h" #include "core/providers/cpu/object_detection/roialign.h" -// TODO: fix the warnings -#if defined(_MSC_VER) && !defined(__clang__) -// Chance of arithmetic overflow could be reduced -#pragma warning(disable : 26451) -#endif using namespace onnxruntime::concurrency; namespace onnxruntime { @@ -97,7 +93,11 @@ void CropAndResizeForward(const TensorShape& output_shape, ? roi_start_h * (height - 1) : 0.5 * (roi_start_h + roi_end_h) * (height - 1)); } - if (in_y < 0 || in_y > height - 1) { + // Route non-finite (NaN/inf) or out-of-image coordinates to the extrapolation + // branch. The negated-conjunction form is NaN-safe: every comparison with NaN is + // false, so !(in_y >= 0 && in_y <= height - 1) is true and NaN cannot reach the + // integer index math below. For finite values this is identical to the < / > form. + if (!(in_y >= 0 && in_y <= static_cast(height - 1))) { for (int64_t pw = 0; pw < pooled_width; pw++) { for (int64_t c = 0; c < channels; c++) { int64_t index_n_c = index_n + c * pooled_width * pooled_height; @@ -126,7 +126,8 @@ void CropAndResizeForward(const TensorShape& output_shape, ? roi_start_w * (width - 1) : 0.5 * (roi_start_w + roi_end_w) * (width - 1)); } - if (in_x < 0 || in_x > width - 1) { + // NaN-safe bounds guard (see the in_y guard above for rationale). + if (!(in_x >= 0 && in_x <= static_cast(width - 1))) { for (int64_t c = 0; c < channels; c++) { int64_t index_n_c = index_n + c * pooled_width * pooled_height; int64_t index = index_n_c + ph * pooled_width + pw; @@ -140,16 +141,17 @@ void CropAndResizeForward(const TensorShape& output_shape, const int left_x_index = static_cast(floorf(static_cast(in_x))); const int right_x_index = static_cast(ceilf(static_cast(in_x))); const float x_lerp = static_cast(in_x - left_x_index); - auto top_left_index = top_y_index * width + left_x_index; - auto top_right_index = top_y_index * width + right_x_index; - auto bottom_left_index = bottom_y_index * width + left_x_index; - auto bottom_right_index = bottom_y_index * width + right_x_index; + auto top_left_index = SafeInt(top_y_index) * width + left_x_index; + auto top_right_index = SafeInt(top_y_index) * width + right_x_index; + auto bottom_left_index = SafeInt(bottom_y_index) * width + left_x_index; + auto bottom_right_index = SafeInt(bottom_y_index) * width + right_x_index; for (auto c = 0; c < channels; c++) { int64_t index_n_c = index_n + c * pooled_width * pooled_height; int64_t index = index_n_c + ph * pooled_width + pw; const T* offset_bottom_data = - bottom_data + static_cast((roi_batch_ind * channels + c) * height * width); + bottom_data + + static_cast((SafeInt(roi_batch_ind) * channels + c) * height * width); const float top_left(static_cast(offset_bottom_data[top_left_index])); const float top_right(static_cast(offset_bottom_data[top_right_index])); const float bottom_left(static_cast(offset_bottom_data[bottom_left_index])); @@ -162,13 +164,14 @@ void CropAndResizeForward(const TensorShape& output_shape, } else { // mode == "nearest" const int closest_x_index = static_cast(roundf(static_cast(in_x))); const int closest_y_index = static_cast(roundf(static_cast(in_y))); - auto closest_index = closest_y_index * width + closest_x_index; + auto closest_index = SafeInt(closest_y_index) * width + closest_x_index; for (auto c = 0; c < channels; c++) { int64_t index_n_c = index_n + c * pooled_width * pooled_height; int64_t index = index_n_c + ph * pooled_width + pw; const T* offset_bottom_data = - bottom_data + static_cast((roi_batch_ind * channels + c) * height * width); + bottom_data + + static_cast((SafeInt(roi_batch_ind) * channels + c) * height * width); top_data[index] = static_cast(offset_bottom_data[closest_index]); } } @@ -217,6 +220,9 @@ Status CropAndResize::Compute(OpKernelContext* context) const { auto crop_height = crop_size_data[0]; auto crop_width = crop_size_data[1]; + ORT_RETURN_IF_NOT(crop_height > 0 && crop_width > 0, + "crop_size values must be positive; got [", crop_height, ", ", crop_width, "]"); + auto status = CheckROIAlignValidInput(X_ptr, rois_ptr, batch_indices_ptr); if (status != Status::OK()) { return status; diff --git a/onnxruntime/test/contrib_ops/crop_and_resize_op_test.cc b/onnxruntime/test/contrib_ops/crop_and_resize_op_test.cc index b7bf7386e72d5..68d3e4332a45c 100644 --- a/onnxruntime/test/contrib_ops/crop_and_resize_op_test.cc +++ b/onnxruntime/test/contrib_ops/crop_and_resize_op_test.cc @@ -4,6 +4,8 @@ #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include + namespace onnxruntime { namespace test { @@ -123,5 +125,122 @@ TEST(CropAndResizeTest, CropAndResize_1133) { test3.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); } +// Non-finite ROI coordinates (NaN) must route to the extrapolation branch instead of +// reaching the integer index math, which would otherwise compute an invalid image index. +TEST(CropAndResizeTest, CropAndResize_NaN_roi_extrapolates) { + const float nan = std::numeric_limits::quiet_NaN(); + + // NaN start/end height -> in_y is NaN for every row -> whole output extrapolates. + OpTester test_y("CropAndResize", 1, onnxruntime::kMSDomain); + test_y.AddInput("X", {1, 1, 2, 2}, {1.1f, 2.2f, 3.3f, 4.4f}); + test_y.AddInput("rois", {1, 4}, {nan, 0.0f, nan, 1.0f}); + test_y.AddInput("batch_indices", {1}, {0}); + test_y.AddInput("crop_size", {2}, {2, 2}); + test_y.AddAttribute("extrapolation_value", 5.5f); + test_y.AddOutput("output", {1, 1, 2, 2}, {5.5f, 5.5f, 5.5f, 5.5f}); + test_y.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); + + // Finite height but NaN width -> in_y is in range, in_x is NaN -> every pixel extrapolates. + OpTester test_x("CropAndResize", 1, onnxruntime::kMSDomain); + test_x.AddInput("X", {1, 1, 2, 2}, {1.1f, 2.2f, 3.3f, 4.4f}); + test_x.AddInput("rois", {1, 4}, {0.0f, nan, 1.0f, nan}); + test_x.AddInput("batch_indices", {1}, {0}); + test_x.AddInput("crop_size", {2}, {2, 2}); + test_x.AddAttribute("extrapolation_value", 5.5f); + test_x.AddOutput("output", {1, 1, 2, 2}, {5.5f, 5.5f, 5.5f, 5.5f}); + test_x.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); + + // Same NaN-height case in nearest mode. + OpTester test_nearest("CropAndResize", 1, onnxruntime::kMSDomain); + test_nearest.AddInput("X", {1, 1, 2, 2}, {1.1f, 2.2f, 3.3f, 4.4f}); + test_nearest.AddInput("rois", {1, 4}, {nan, 0.0f, nan, 1.0f}); + test_nearest.AddInput("batch_indices", {1}, {0}); + test_nearest.AddInput("crop_size", {2}, {2, 2}); + test_nearest.AddAttribute("mode", "nearest"); + test_nearest.AddAttribute("extrapolation_value", 5.5f); + test_nearest.AddOutput("output", {1, 1, 2, 2}, {5.5f, 5.5f, 5.5f, 5.5f}); + test_nearest.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); +} + +// Infinite ROI coordinates must also land in the extrapolation branch. +TEST(CropAndResizeTest, CropAndResize_Inf_roi_extrapolates) { + const float inf = std::numeric_limits::infinity(); + + OpTester test_pos("CropAndResize", 1, onnxruntime::kMSDomain); + test_pos.AddInput("X", {1, 1, 2, 2}, {1.1f, 2.2f, 3.3f, 4.4f}); + test_pos.AddInput("rois", {1, 4}, {inf, 0.0f, inf, 1.0f}); + test_pos.AddInput("batch_indices", {1}, {0}); + test_pos.AddInput("crop_size", {2}, {2, 2}); + test_pos.AddAttribute("extrapolation_value", 5.5f); + test_pos.AddOutput("output", {1, 1, 2, 2}, {5.5f, 5.5f, 5.5f, 5.5f}); + test_pos.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); + + OpTester test_neg("CropAndResize", 1, onnxruntime::kMSDomain); + test_neg.AddInput("X", {1, 1, 2, 2}, {1.1f, 2.2f, 3.3f, 4.4f}); + test_neg.AddInput("rois", {1, 4}, {-inf, 0.0f, -inf, 1.0f}); + test_neg.AddInput("batch_indices", {1}, {0}); + test_neg.AddInput("crop_size", {2}, {2, 2}); + test_neg.AddAttribute("extrapolation_value", 5.5f); + test_neg.AddOutput("output", {1, 1, 2, 2}, {5.5f, 5.5f, 5.5f, 5.5f}); + test_neg.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); +} + +// Finite coordinates on and just outside the image boundary must behave exactly as the +// original guard: coordinates in [0, dim-1] interpolate, coordinates just past dim-1 extrapolate. +TEST(CropAndResizeTest, CropAndResize_finite_boundary_no_regression) { + // Exact boundary [0, 1] maps to the identity crop (no extrapolation at 0 or height-1). + OpTester test_exact("CropAndResize", 1, onnxruntime::kMSDomain); + test_exact.AddInput("X", {1, 1, 2, 2}, {1.1f, 2.2f, 3.3f, 4.4f}); + test_exact.AddInput("rois", {1, 4}, {0.0f, 0.0f, 1.0f, 1.0f}); + test_exact.AddInput("batch_indices", {1}, {0}); + test_exact.AddInput("crop_size", {2}, {2, 2}); + test_exact.AddAttribute("extrapolation_value", 9.0f); + test_exact.AddOutput("output", {1, 1, 2, 2}, {1.1f, 2.2f, 3.3f, 4.4f}); + test_exact.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); + + // End just past the boundary (1.0001) extrapolates only the out-of-range row/column. + OpTester test_outside("CropAndResize", 1, onnxruntime::kMSDomain); + test_outside.AddInput("X", {1, 1, 2, 2}, {1.1f, 2.2f, 3.3f, 4.4f}); + test_outside.AddInput("rois", {1, 4}, {0.0f, 0.0f, 1.0001f, 1.0001f}); + test_outside.AddInput("batch_indices", {1}, {0}); + test_outside.AddInput("crop_size", {2}, {2, 2}); + test_outside.AddAttribute("extrapolation_value", 9.0f); + test_outside.AddOutput("output", {1, 1, 2, 2}, {1.1f, 9.0f, 9.0f, 9.0f}); + test_outside.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); +} + +// crop_size values must be positive; non-positive crop dimensions are rejected via Status. +TEST(CropAndResizeTest, CropAndResize_rejects_nonpositive_crop_size) { + OpTester test_zero("CropAndResize", 1, onnxruntime::kMSDomain); + test_zero.AddInput("X", {1, 1, 2, 2}, {1.1f, 2.2f, 3.3f, 4.4f}); + test_zero.AddInput("rois", {1, 4}, {0.0f, 0.0f, 1.0f, 1.0f}); + test_zero.AddInput("batch_indices", {1}, {0}); + test_zero.AddInput("crop_size", {2}, {0, 2}); + test_zero.AddOutput("output", {1, 1, 1, 1}, {0.0f}); + test_zero.Run(OpTester::ExpectResult::kExpectFailure, + "crop_size values must be positive", {kTensorrtExecutionProvider}); + + OpTester test_negative("CropAndResize", 1, onnxruntime::kMSDomain); + test_negative.AddInput("X", {1, 1, 2, 2}, {1.1f, 2.2f, 3.3f, 4.4f}); + test_negative.AddInput("rois", {1, 4}, {0.0f, 0.0f, 1.0f, 1.0f}); + test_negative.AddInput("batch_indices", {1}, {0}); + test_negative.AddInput("crop_size", {2}, {-1, 2}); + test_negative.AddOutput("output", {1, 1, 1, 1}, {0.0f}); + test_negative.Run(OpTester::ExpectResult::kExpectFailure, + "crop_size values must be positive", {kTensorrtExecutionProvider}); +} + +// A batch index outside [0, batch_size) must be rejected rather than computing an invalid image index. +TEST(CropAndResizeTest, CropAndResize_rejects_out_of_range_batch_index) { + OpTester test("CropAndResize", 1, onnxruntime::kMSDomain); + test.AddInput("X", {1, 1, 2, 2}, {1.1f, 2.2f, 3.3f, 4.4f}); + test.AddInput("rois", {1, 4}, {0.0f, 0.0f, 1.0f, 1.0f}); + test.AddInput("batch_indices", {1}, {5}); + test.AddInput("crop_size", {2}, {2, 2}); + test.AddOutput("output", {1, 1, 2, 2}, {0.0f, 0.0f, 0.0f, 0.0f}); + test.Run(OpTester::ExpectResult::kExpectFailure, + "is out of range [0, 1)", {kTensorrtExecutionProvider}); +} + } // namespace test } // namespace onnxruntime From dd32f3598a62d56a434439aa134176c96005ad69 Mon Sep 17 00:00:00 2001 From: Copilot <198982749+Copilot@users.noreply.github.com> Date: Wed, 8 Jul 2026 11:15:28 -0700 Subject: [PATCH 5/5] Fix libcudart.so.13 hard dependency in pybind module breaking import on CPU-only Linux (#29590) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description `onnxruntime-gpu` 1.27 introduced a hard `NEEDED libcudart.so.13` entry in `onnxruntime_pybind11_state.so`, causing `ImportError` at `import onnxruntime` on CPU-only Linux machines — before any provider is selected. **Root cause:** `cmake/onnxruntime_python.cmake` was changed to compile `fpA_intB_gemm_adaptor.cu` and `fpA_intB_gemm_preprocessors_impl.cu` directly into `onnxruntime_pybind11_state.so` and link `CUDA::cudart` (dynamic). This embeds a load-time CUDA dependency in the Python module itself. **Fix:** Move the CUDA weight-preprocessing entry point (`pack_weights_for_cuda_mixed_gemm`) out of the main pybind module and into a **standalone extension module**, `onnxruntime_cuda_quant_preprocess`, that links `CUDA::cudart` on its own. The main `onnxruntime_pybind11_state.so` no longer compiles or links any CUDA code, so `import onnxruntime` has no `libcudart` dependency. The new module is imported **lazily** by `onnxruntime/python/tools/quantization/cuda_quantizer.py` only when weight prepacking is actually requested — never at `import onnxruntime` time. These preprocessing APIs are **offline-only** helpers: they are used by quantization tooling and model builders to produce prepacked weight initializers ahead of time, and are not part of the inference runtime hot path. Because nothing in the runtime imports them, isolating them into a separate, on-demand DLL has no runtime cost and cleanly keeps CUDA out of the base `import onnxruntime` path. **Why not the provider bridge:** An earlier iteration routed the call through the `ProviderInfo_CUDA` virtual interface (`TryGetProviderInfo_CUDA()`). That does not work for the CUDA-EP-as-plugin build (`onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON`): `cuda_provider_factory.cc` is excluded from the plugin sources and there is no provider bridge, so `TryGetProviderInfo_CUDA()` returns `nullptr` and the call throws. The standalone module has no such dependency and works for **both** the legacy in-tree CUDA EP build and the plugin build. ### Key Changes | File | Change | |---|---| | `onnxruntime/python/onnxruntime_pybind_cuda_quant.cc` | **New.** Self-contained `pack_weights_for_cuda_mixed_gemm` (device malloc + transpose/convert + arch permutation) and a `PYBIND11_MODULE(onnxruntime_cuda_quant_preprocess, …)` entry point. | | `cmake/onnxruntime_python.cmake` | Add the `onnxruntime_cuda_quant_preprocess` module target (built when `onnxruntime_USE_CUDA AND NOT WIN32`, compiling the two `fpA_intB` `.cu` files + `CUDA::cudart` + cutlass, hidden visibility) and copy it into `onnxruntime/capi/`. Main pybind module keeps no CUDA sources/links. | | `onnxruntime/python/onnxruntime_pybind_quant.cc` | Remove the `USE_CUDA` `PackWeightsForMixedGemm` and its registration. The CPU-only `pack_fp4_weights_for_cuda_moe_gemm` stays in the main module. | | `onnxruntime/core/providers/cuda/cuda_provider_factory.{h,cc}` | Revert the `PackWeightsForMixedGemm` `ProviderInfo_CUDA` addition (no longer needed; absent in plugin builds). | | `onnxruntime/python/tools/quantization/cuda_quantizer.py` | `_get_pack_weights_for_cuda_mixed_gemm()` now imports `onnxruntime.capi.onnxruntime_cuda_quant_preprocess` lazily; add `has_cuda_weight_prepacking()` capability helper. | | `setup.py` | Package `onnxruntime_cuda_quant_preprocess.so` in the Linux/macOS wheels. | | `onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py` | Point the prepacked-weight parity test and its skip guard at the new module. | | `docs/contrib_ops/cuda/matmul_nbits.md` | Update the offline-packer code snippets to import the new module. | ### Motivation and Context `import onnxruntime` must succeed on CPU-only machines even when the GPU wheel is installed. CUDA dependency errors should surface only when a CUDA provider is explicitly loaded/selected, or when offline CUDA weight prepacking is explicitly requested. This restores the 1.26 behavior where `onnxruntime_pybind11_state.so` had no `NEEDED libcudart.so.*` entry, and — unlike the provider-bridge approach — it also works in the CUDA-EP-as-plugin build. ### Testing Notes - Built both modules in the CUDA build; `readelf -d onnxruntime_pybind11_state.so` shows **no** `libcudart` `NEEDED` entry, while `onnxruntime_cuda_quant_preprocess.so` has `NEEDED libcudart.so.13`. - `import onnxruntime` and lazy loading of `onnxruntime.capi.onnxruntime_cuda_quant_preprocess` both succeed; `has_cuda_weight_prepacking()` returns `True` on a CUDA machine. - `test_op_matmulnbits_prepacked_cuda.py` passes (INT4/INT8 prepacked-vs-runtime parity), confirming the relocated packer produces byte-identical prepacked weights. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Tianlei Wu --- cmake/CMakeLists.txt | 1 + cmake/onnxruntime_python.cmake | 93 +++++-- docs/contrib_ops/cuda/matmul_nbits.md | 6 +- docs/cuda_plugin_ep/QUICK_START.md | 2 +- docs/cuda_plugin_ep/cuda_plugin_ep_design.md | 8 +- .../python/onnxruntime_pybind_cuda_quant.cc | 166 +++++++++++++ .../python/onnxruntime_pybind_mlvalue.cc | 29 --- .../python/onnxruntime_pybind_ortvalue.cc | 9 - .../python/onnxruntime_pybind_quant.cc | 143 +---------- .../tools/quantization/cuda_quantizer.py | 226 ++++++++++++++++-- .../test_op_matmulnbits_prepacked_cuda.py | 44 +++- setup.py | 9 + tools/ci_build/build.py | 4 +- 13 files changed, 519 insertions(+), 221 deletions(-) create mode 100644 onnxruntime/python/onnxruntime_pybind_cuda_quant.cc diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 09e307e124316..ad446214cfc8f 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -77,6 +77,7 @@ cmake_dependent_option(onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS "Build with CUD cmake_dependent_option(onnxruntime_USE_CUDA_NHWC_OPS "Build CUDA with NHWC op support" ON "onnxruntime_USE_CUDA" OFF) cmake_dependent_option(onnxruntime_BUILD_CUDA_EP_AS_PLUGIN "Build CUDA EP as a separate plugin shared library instead of the legacy in-tree provider" OFF "onnxruntime_USE_CUDA" OFF) +option(onnxruntime_BUILD_CUDA_QUANT_PREPROCESS "Build CUDA weight-packing module onnxruntime_cuda_quant_preprocess.so" ON) option(onnxruntime_CUDA_MINIMAL "Build CUDA without any operations apart from memcpy ops. Usefuel for a very minial TRT build" OFF) option(onnxruntime_ENABLE_CUDA_LINE_NUMBER_INFO "When building with CUDA support, generate device code line number information." OFF) option(onnxruntime_USE_OPENVINO "Build with OpenVINO support" OFF) diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 3f6976c3c8955..14b01f1f25473 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -20,6 +20,13 @@ file(GLOB onnxruntime_pybind_srcs CONFIGURE_DEPENDS ${onnxruntime_pybind_srcs_pattern} ) +# onnxruntime_pybind_cuda_quant.cc is compiled into the standalone +# onnxruntime_cuda_quant_preprocess extension module (see below), not into +# onnxruntime_pybind11_state. It includes and links CUDA::cudart, +# so compiling it into the main pybind module would break CPU-only builds and +# re-introduce the hard libcudart dependency this design avoids. +list(REMOVE_ITEM onnxruntime_pybind_srcs ${ONNXRUNTIME_ROOT}/python/onnxruntime_pybind_cuda_quant.cc) + if(onnxruntime_ENABLE_TRAINING) list(REMOVE_ITEM onnxruntime_pybind_srcs ${ONNXRUNTIME_ROOT}/python/onnxruntime_pybind_module.cc) endif() @@ -231,22 +238,11 @@ target_link_libraries(onnxruntime_pybind11_state PRIVATE Python::NumPy ) -# Starting with Python 3.8 on Windows, PATH environment variable are no longer used to resolve DLL dependencies -# for extension modules or libraries loaded via ctypes. -# To avoid package import issues, we do not link pybind module against the CUDA runtime on Windows, instead of -# os.add_dll_directory() to deal with CUDA paths. -if (onnxruntime_USE_CUDA AND NOT WIN32) - target_sources(onnxruntime_pybind11_state PRIVATE - "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.cu" - "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu" - ) - include(cutlass) - target_include_directories(onnxruntime_pybind11_state PRIVATE ${cutlass_SOURCE_DIR}/include) - target_link_libraries(onnxruntime_pybind11_state PRIVATE CUDA::cudart) -endif() -if (onnxruntime_USE_CUDA AND WIN32) - target_compile_definitions(onnxruntime_pybind11_state PRIVATE ORT_NO_CUDA_IN_PYBIND) -endif() +# The CUDA quantization helpers (pack_weights_for_cuda_mixed_gemm) are built into a +# separate extension module (onnxruntime_cuda_quant_preprocess) that is imported on +# demand. Do NOT compile CUDA source files directly into onnxruntime_pybind11_state or +# link CUDA::cudart from it: that would create a hard libcudart.so dependency that +# prevents importing the Python module on CPU-only machines. set(onnxruntime_pybind11_state_dependencies ${onnxruntime_EXTERNAL_DEPENDENCIES} @@ -315,6 +311,71 @@ else() set_target_properties(onnxruntime_pybind11_state PROPERTIES SUFFIX ".so") endif() +# --------------------------------------------------------------------------- +# Standalone CUDA weight-preprocessing extension module. +# +# The CUDA weight-packing kernels (pack_weights_for_cuda_mixed_gemm) are compiled +# into their OWN Python extension module instead of onnxruntime_pybind11_state. +# This keeps the hard libcudart dependency out of the main pybind module so that +# `import onnxruntime` still works on CPU-only machines. +# +# Production weight packing is done in PyTorch (cuda_quantizer.py); this module is +# retained only as a byte-parity oracle for that PyTorch packer. It is gated by +# onnxruntime_BUILD_CUDA_QUANT_PREPROCESS. +# +# It does NOT go through the provider bridge / ProviderInfo_CUDA, so it works for +# both the legacy in-tree CUDA EP build and the CUDA-EP-as-plugin build. +# +# Not built on Windows: matching the previous behavior where CUDA runtime was not +# linked into Python extension modules (DLL search path constraints since +# Python 3.8), so pack_weights_for_cuda_mixed_gemm was unavailable there. +if (onnxruntime_USE_CUDA AND NOT WIN32 AND onnxruntime_BUILD_CUDA_QUANT_PREPROCESS) + onnxruntime_add_shared_library_module(onnxruntime_cuda_quant_preprocess + "${ONNXRUNTIME_ROOT}/python/onnxruntime_pybind_cuda_quant.cc" + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.cu" + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu" + ) + include(cutlass) + onnxruntime_add_include_to_target(onnxruntime_cuda_quant_preprocess Python::Module onnxruntime_common) + target_include_directories(onnxruntime_cuda_quant_preprocess PRIVATE + ${ONNXRUNTIME_ROOT} + ${pybind11_INCLUDE_DIRS} + ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} + ${cutlass_SOURCE_DIR}/include + ${cutlass_SOURCE_DIR}/tools/util/include + ) + target_compile_definitions(onnxruntime_cuda_quant_preprocess PRIVATE USE_CUDA) + target_link_libraries(onnxruntime_cuda_quant_preprocess PRIVATE + onnxruntime_common + Boost::mp11 + safeint_interface + ${ABSEIL_LIBS} + CUDA::cudart + ${pybind11_lib} + Python::NumPy + ) + if (NOT MSVC) + target_compile_options(onnxruntime_cuda_quant_preprocess PRIVATE "-fvisibility=hidden") + endif() + set_target_properties(onnxruntime_cuda_quant_preprocess PROPERTIES PREFIX "" SUFFIX ".so" FOLDER "ONNXRuntime") + if (APPLE) + set_target_properties(onnxruntime_cuda_quant_preprocess PROPERTIES + INSTALL_RPATH "@loader_path" + BUILD_WITH_INSTALL_RPATH TRUE + INSTALL_RPATH_USE_LINK_PATH FALSE) + elseif (NOT CMAKE_SYSTEM_NAME MATCHES "AIX") + target_link_options(onnxruntime_cuda_quant_preprocess PRIVATE "LINKER:-rpath=\$ORIGIN") + endif() + # Place the module next to the main pybind module inside onnxruntime/capi. + add_custom_command( + TARGET onnxruntime_cuda_quant_preprocess POST_BUILD + COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/capi + COMMAND ${CMAKE_COMMAND} -E copy + $ + $/onnxruntime/capi/ + ) +endif() + # Generate version_info.py in Windows build. # Has to be done before onnxruntime_python_srcs is set. if (WIN32) diff --git a/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md index a09a80a4732ed..3f8e1f1d8f616 100644 --- a/docs/contrib_ops/cuda/matmul_nbits.md +++ b/docs/contrib_ops/cuda/matmul_nbits.md @@ -87,9 +87,9 @@ step is **not** performed. The offline CUDA packer exposed through Python produces this layout: ```python -from onnxruntime.capi import _pybind_state as _pybind +from onnxruntime.capi import onnxruntime_cuda_quant_preprocess as _cuda_quant -prepacked_flat = _pybind.pack_weights_for_cuda_mixed_gemm( +prepacked_flat = _cuda_quant.pack_weights_for_cuda_mixed_gemm( q_weight.reshape(N, -1), N, K, bits, 80 ) prepacked_b = np.asarray(prepacked_flat, dtype=np.int8).view(np.uint8).reshape(q_weight.shape) @@ -309,7 +309,7 @@ present. `ComputeInternal` then: GEMV profiling helpers, e.g. `profile_qmoe_gemv.sh`). - CUDA prepacked-weight parity tests: [onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py](../../../onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py). - These use `_pybind_state.pack_weights_for_cuda_mixed_gemm(..., 80)` to produce + These use `onnxruntime_cuda_quant_preprocess.pack_weights_for_cuda_mixed_gemm(..., 80)` to produce `weight_prepacked=1` initializers and compare their outputs against runtime fpA_intB prepacking for int4/int8 and GEMV/GEMM-shaped `M` values. - Constructor failure tests for unsupported prepacked configurations live in diff --git a/docs/cuda_plugin_ep/QUICK_START.md b/docs/cuda_plugin_ep/QUICK_START.md index ab7b4308f13e5..7b971d621083f 100644 --- a/docs/cuda_plugin_ep/QUICK_START.md +++ b/docs/cuda_plugin_ep/QUICK_START.md @@ -146,7 +146,7 @@ sess = ort.InferenceSession( **Python `OrtValue` host/device copies:** -`OrtValue.update_inplace()` and `OrtValue.numpy()` work with CUDA plugin tensors after the plugin has been registered. On Linux, the ONNX Runtime Python binding links the CUDA runtime and can fall back to direct `cudaMemcpy` if the legacy CUDA provider bridge is unavailable. On Windows, the Python binding is built with `ORT_NO_CUDA_IN_PYBIND`, so it cannot call CUDA runtime APIs directly; host/device copies must use the data-transfer implementation registered by the CUDA plugin library. If `OrtValue.update_inplace()` fails with a message about the CUDA provider interface or an unsupported GPU device, verify that the plugin library is registered before creating or updating CUDA `OrtValue` objects. +`OrtValue.update_inplace()` and `OrtValue.numpy()` work with CUDA plugin tensors after the plugin has been registered. The Python binding cannot call CUDA runtime APIs directly; host/device copies must use the data-transfer implementation registered by the CUDA plugin library. If `OrtValue.update_inplace()` fails with a message about the CUDA provider interface or an unsupported GPU device, verify that the plugin library is registered before creating or updating CUDA `OrtValue` objects. ### External GPU Allocator Options diff --git a/docs/cuda_plugin_ep/cuda_plugin_ep_design.md b/docs/cuda_plugin_ep/cuda_plugin_ep_design.md index a06f1de95b99e..7dfa2bd2d8667 100644 --- a/docs/cuda_plugin_ep/cuda_plugin_ep_design.md +++ b/docs/cuda_plugin_ep/cuda_plugin_ep_design.md @@ -344,10 +344,12 @@ This is intentionally conservative and correct for the plugin EP's first sync in The Python `OrtValue` helpers (`update_inplace()` for host-to-device and `numpy()` for device-to-host) historically reached CUDA copies through the legacy provider bridge (`GetProviderInfo_CUDA()`). That bridge requires the provider shared library to export `GetProvider()`, which the CUDA plugin intentionally does not export. -The fallback path is platform-specific: +To keep working when the bridge is absent (as with the plugin EP), the pybind can reach CUDA copies two ways: the legacy provider bridge (`TryGetProviderInfo_CUDA()`) and a plugin-registered `OrtDataTransfer` copy function (`CreateDataTransferMemCpy()`, backed by the plugin EP's `IDataTransfer`). It tries whichever is available and throws if neither is, in which case a CUDA `OrtValue` copy cannot be performed. -- On non-Windows CUDA builds, `onnxruntime_pybind11_state` links `CUDA::cudart`. If `TryGetProviderInfo_CUDA()` fails, pybind can copy directly with `cudaMemcpy`; host-to-device copies synchronize the default stream, matching `ProviderInfo_CUDA::cudaMemcpy_HostToDevice()`. -- On Windows CUDA builds, pybind is compiled with `ORT_NO_CUDA_IN_PYBIND` and does not link CUDA runtime APIs. If `TryGetProviderInfo_CUDA()` fails, pybind must obtain an `OrtDataTransfer` copy function from the registered plugin EP. Without a registered plugin data-transfer implementation, CUDA `OrtValue.update_inplace()` cannot copy host data into the plugin-owned device tensor. +The two code paths differ only in which mechanism they try first, and this does not change the outcome (exactly one applies in a given build): + +- `OrtValue.update_inplace(numpy_array)` / `OrtValue.numpy()` (in `onnxruntime_pybind_ortvalue.cc`) try the provider bridge first, then fall back to the plugin `OrtDataTransfer`. +- `OrtValue.update_inplace(OrtValue)` (`UpdateOrtValueInplace` in `onnxruntime_pybind_mlvalue.cc`) tries the plugin `OrtDataTransfer` first, then falls back to the built-in CUDA provider copy functions. ### 5.2 Handle Access Path diff --git a/onnxruntime/python/onnxruntime_pybind_cuda_quant.cc b/onnxruntime/python/onnxruntime_pybind_cuda_quant.cc new file mode 100644 index 0000000000000..3edadd786a8c9 --- /dev/null +++ b/onnxruntime/python/onnxruntime_pybind_cuda_quant.cc @@ -0,0 +1,166 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Standalone CUDA weight-preprocessing extension module. +// +// This module is intentionally kept SEPARATE from onnxruntime_pybind11_state so +// that `import onnxruntime` never triggers a load-time dependency on the CUDA +// runtime (libcudart). The CUDA weight packing kernels below link CUDA::cudart, +// so this module has a hard libcudart dependency -- but it is imported lazily by +// onnxruntime.python.tools.quantization.cuda_quantizer only when CUDA weight +// prepacking is actually requested. +// +// This approach works for both the legacy in-tree CUDA EP build and the +// CUDA-EP-as-plugin build (onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON), because it +// does not rely on the provider bridge / ProviderInfo_CUDA interface (which is +// not available in plugin builds). + +#include +#include + +#include + +#include +#include +#include + +#include "contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.h" +#include "contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h" + +namespace py = pybind11; + +namespace { + +void ThrowIfCudaError(cudaError_t status, const char* expression) { + if (status != cudaSuccess) { + std::ostringstream oss; + oss << expression << " failed: " << cudaGetErrorString(status); + throw std::runtime_error(oss.str()); + } +} + +struct CudaDeleter { + void operator()(void* p) const { + if (p) cudaFree(p); + } +}; + +using CudaPtr = std::unique_ptr; + +// Preprocess quantized weights for CUDA mixed-precision GEMM kernels (FpA_IntB format). +// +// MatMulNBits/QMoE stores quantized weights in (N, K) layout: +// - N = number of output channels (columns in weight matrix W) +// - K = number of input features (rows in weight matrix W) +// - For 4-bit: shape is (N, K/2) bytes where each byte packs 2 elements +// - For 8-bit: shape is (N, K) bytes +// +// FpA_IntB GEMM kernels expect weights in (K, N) layout (transposed) for efficient +// memory access during matrix multiplication. This function: +// 1. Transposes from (N, K) to (K, N) layout +// 2. Converts unsigned quantized values to signed int8 with zero-point adjustment +// - 4-bit: uint4 -> int8 with zero_point=8 (range [0,15] -> [-8,7]) +// - 8-bit: uint8 -> int8 with zero_point=128 (range [0,255] -> [-128,127]) +// 3. Applies architecture-specific row permutation for optimized tensor core access +// +// Input: q_weights - Quantized weights from MatMulNBits in (N, K) layout +// Output: Preprocessed weights in (K, N) layout ready for fpA_intB GEMM kernels +py::array_t PackWeightsForMixedGemm( + py::array_t q_weights, + int32_t N, + int32_t K, + int32_t bits, + int32_t force_arch = -1) { + py::buffer_info q_weights_buf = q_weights.request(); + + if (bits != 4 && bits != 8) { + throw std::invalid_argument("bits must be 4 or 8"); + } + if (N <= 0 || K <= 0) { + throw std::invalid_argument("N and K must be positive"); + } + if (bits == 4 && K % 2 != 0) { + throw std::invalid_argument("K must be even for 4-bit packed weights"); + } + if (q_weights_buf.ndim != 2 || q_weights_buf.shape[0] != N || q_weights_buf.shape[1] != K / (8 / bits)) { + throw std::invalid_argument("q_weights must have shape (N, K / (8 / bits))"); + } + + int n = static_cast(N); + int k = static_cast(K); + + size_t packed_weight_bytes = static_cast(n) * static_cast(k) / (8 / bits); + py::array_t processed_weights({static_cast(packed_weight_bytes)}); + py::buffer_info processed_weights_buf = processed_weights.request(); + + auto make_cuda_ptr = [](size_t bytes) -> CudaPtr { + void* p = nullptr; + ThrowIfCudaError(cudaMalloc(&p, bytes), "cudaMalloc"); + return CudaPtr(p); + }; + + auto packed_transposed_weight_space = make_cuda_ptr(packed_weight_bytes); + int8_t* packed_transposed_weight = reinterpret_cast(packed_transposed_weight_space.get()); + + auto fpA_intB_weight_buffer_ = make_cuda_ptr(packed_weight_bytes); + int8_t* preprocessed_weight = reinterpret_cast(fpA_intB_weight_buffer_.get()); + + const uint8_t* blob_data_cpu = static_cast(q_weights_buf.ptr); + + auto blob_data_gpu_buf = make_cuda_ptr(packed_weight_bytes); + uint8_t* blob_data_gpu = reinterpret_cast(blob_data_gpu_buf.get()); + + cudaStream_t stream = cudaStreamLegacy; + ThrowIfCudaError(cudaMemcpyAsync(blob_data_gpu, blob_data_cpu, packed_weight_bytes, cudaMemcpyHostToDevice, stream), + "cudaMemcpyAsync host-to-device"); + + if (bits == 4) { + ::onnxruntime::llm::kernels::fpA_intB_gemv::unpack_uint4_transposed_to_int8_direct_cuda( + stream, packed_transposed_weight, blob_data_gpu, n, k); + } else { + // 8 bits + ::onnxruntime::llm::kernels::fpA_intB_gemv::transpose_uint8_matrix_and_convert_to_int8( + stream, packed_transposed_weight, blob_data_gpu, n, k); + } + + using ::onnxruntime::llm::kernels::weight_only::QuantType; + QuantType quant_type = bits == 4 ? QuantType::W4_A16 : QuantType::W8_A16; + + int sm = force_arch; + if (sm < 0) { + int device_id = 0; + ThrowIfCudaError(cudaGetDevice(&device_id), "cudaGetDevice"); + cudaDeviceProp device_prop; + ThrowIfCudaError(cudaGetDeviceProperties(&device_prop, device_id), "cudaGetDeviceProperties"); + sm = device_prop.major * 10 + device_prop.minor; + } + sm = ::onnxruntime::llm::kernels::weight_only::get_arch_for_mixed_gemm_weight_preprocess(sm); + + auto permutation_map_buffer = make_cuda_ptr(32 * sizeof(int32_t)); + + ::onnxruntime::llm::kernels::weight_only::preprocess_weights_for_mixed_gemm_cuda( + stream, + sm, + preprocessed_weight, + packed_transposed_weight, + reinterpret_cast(permutation_map_buffer.get()), + {static_cast(k), static_cast(n)}, + quant_type); + + ThrowIfCudaError(cudaGetLastError(), "preprocess CUDA kernel launch"); + ThrowIfCudaError(cudaMemcpyAsync(processed_weights_buf.ptr, preprocessed_weight, packed_weight_bytes, + cudaMemcpyDeviceToHost, stream), + "cudaMemcpyAsync device-to-host"); + ThrowIfCudaError(cudaStreamSynchronize(stream), "cudaStreamSynchronize"); + + return processed_weights; +} + +} // namespace + +PYBIND11_MODULE(onnxruntime_cuda_quant_preprocess, m) { + m.doc() = "CUDA weight-only quantization preprocessing helpers (loaded on demand)."; + m.def("pack_weights_for_cuda_mixed_gemm", &PackWeightsForMixedGemm, + "Pack quantized weights for CUDA mixed-precision GEMM (FpA_IntB format)", + py::arg("q_weights"), py::arg("N"), py::arg("K"), py::arg("bits"), py::arg("force_arch") = -1); +} diff --git a/onnxruntime/python/onnxruntime_pybind_mlvalue.cc b/onnxruntime/python/onnxruntime_pybind_mlvalue.cc index 10e55259f834e..5ea5f42958926 100644 --- a/onnxruntime/python/onnxruntime_pybind_mlvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_mlvalue.cc @@ -23,10 +23,6 @@ #include "core/framework/kernel_registry.h" #include "core/framework/provider_options_utils.h" -#if defined(USE_CUDA) && !defined(ORT_NO_CUDA_IN_PYBIND) -#include -#endif - #ifdef USE_DML using Microsoft::WRL::ComPtr; @@ -184,23 +180,6 @@ int32_t GetTensorProtoType(const OrtValue& ort_value) { } #ifdef USE_CUDA -namespace { - -#if !defined(ORT_NO_CUDA_IN_PYBIND) -void CudaRuntimeMemCpy(void* dst, const void* src, size_t num_bytes, cudaMemcpyKind kind) { - const auto copy_result = cudaMemcpy(dst, src, num_bytes, kind); - ORT_ENFORCE(copy_result == cudaSuccess, "cudaMemcpy failed: ", cudaGetErrorString(copy_result)); - - if (kind == cudaMemcpyHostToDevice) { - // Match ProviderInfo_CUDA::cudaMemcpy_HostToDevice: cudaMemcpy() uses the default - // stream, and pageable host-to-device copies can return before DMA to device is done. - const auto sync_result = cudaStreamSynchronize(0); - ORT_ENFORCE(sync_result == cudaSuccess, "cudaStreamSynchronize failed: ", cudaGetErrorString(sync_result)); - } -} -#endif - -} // namespace void CpuToCudaMemCpy(void* dst, const void* src, size_t num_bytes) { if (TryGetProviderInfo_CUDA() != nullptr) { @@ -208,11 +187,7 @@ void CpuToCudaMemCpy(void* dst, const void* src, size_t num_bytes) { return; } -#if !defined(ORT_NO_CUDA_IN_PYBIND) - CudaRuntimeMemCpy(dst, src, num_bytes, cudaMemcpyHostToDevice); -#else ORT_THROW("CUDA provider interface is not available for host-to-device copy."); -#endif } void CudaToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { @@ -221,11 +196,7 @@ void CudaToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { return; } -#if !defined(ORT_NO_CUDA_IN_PYBIND) - CudaRuntimeMemCpy(dst, src, num_bytes, cudaMemcpyDeviceToHost); -#else ORT_THROW("CUDA provider interface is not available for device-to-host copy."); -#endif } const std::unordered_map* GetCudaToHostMemCpyFunction(const OrtDevice& device) { diff --git a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc index cf7f86a0b9e41..7bf9325cf2208 100644 --- a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc @@ -205,7 +205,6 @@ void addOrtValueMethods(pybind11::module& m) { #ifdef USE_CUDA if (device.Vendor() == OrtDevice::VendorIds::NVIDIA) { MemCpyFunc cpu_to_device_copy_fn = CpuToCudaMemCpy; -#if defined(ORT_NO_CUDA_IN_PYBIND) if (TryGetProviderInfo_CUDA() != nullptr) { if (!IsCudaDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine."); @@ -217,12 +216,6 @@ void addOrtValueMethods(pybind11::module& m) { "Unsupported GPU device: Cannot find the supported GPU device."); } } -#else - if (TryGetProviderInfo_CUDA() != nullptr && - !IsCudaDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { - throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine."); - } -#endif onnxruntime::python::CopyDataToTensor( py_values, @@ -467,12 +460,10 @@ void addOrtValueMethods(pybind11::module& m) { switch (device.Vendor()) { #ifdef USE_CUDA case OrtDevice::VendorIds::NVIDIA: -#if defined(ORT_NO_CUDA_IN_PYBIND) if (TryGetProviderInfo_CUDA() == nullptr) { return GetPyObjFromTensor(*ml_value, nullptr, nullptr, /*zero_copy_non_owning=*/true); } -#endif return GetPyObjFromTensor(*ml_value, nullptr, GetCudaToHostMemCpyFunction(device), /*zero_copy_non_owning=*/true); #endif diff --git a/onnxruntime/python/onnxruntime_pybind_quant.cc b/onnxruntime/python/onnxruntime_pybind_quant.cc index 7220153b4fa17..b2b6ed77c6296 100644 --- a/onnxruntime/python/onnxruntime_pybind_quant.cc +++ b/onnxruntime/python/onnxruntime_pybind_quant.cc @@ -9,11 +9,6 @@ #include "contrib_ops/cpu/quantization/dequantize_blockwise_bnb4.h" #include "core/util/thread_utils.h" -#if defined(USE_CUDA) && !defined(ORT_NO_CUDA_IN_PYBIND) -#include -#include "contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.h" -#include "contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h" -#endif #include #include #include @@ -147,132 +142,6 @@ void QuantizeMatMulBnb4Blockwise( tp.get()); } -#if defined(USE_CUDA) && !defined(ORT_NO_CUDA_IN_PYBIND) -namespace cuda { -void ThrowIfCudaError(cudaError_t status, const char* expression) { - if (status != cudaSuccess) { - std::ostringstream oss; - oss << expression << " failed: " << cudaGetErrorString(status); - throw std::runtime_error(oss.str()); - } -} - -struct CudaDeleter { - void operator()(void* p) const { - if (p) cudaFree(p); - } -}; - -using CudaPtr = std::unique_ptr; - -// Preprocess quantized weights for CUDA mixed-precision GEMM kernels (FpA_IntB format). -// -// MatMulNBits/QMoE stores quantized weights in (N, K) layout: -// - N = number of output channels (columns in weight matrix W) -// - K = number of input features (rows in weight matrix W) -// - For 4-bit: shape is (N, K/2) bytes where each byte packs 2 elements -// - For 8-bit: shape is (N, K) bytes -// -// FpA_IntB GEMM kernels expect weights in (K, N) layout (transposed) for efficient -// memory access during matrix multiplication. This function: -// 1. Transposes from (N, K) to (K, N) layout -// 2. Converts unsigned quantized values to signed int8 with zero-point adjustment -// - 4-bit: uint4 -> int8 with zero_point=8 (range [0,15] -> [-8,7]) -// - 8-bit: uint8 -> int8 with zero_point=128 (range [0,255] -> [-128,127]) -// 3. Applies architecture-specific row permutation for optimized tensor core access -// -// Input: q_weights - Quantized weights from MatMulNBits in (N, K) layout -// Output: Preprocessed weights in (K, N) layout ready for fpA_intB GEMM kernels -py::array_t PackWeightsForMixedGemm( - py::array_t q_weights, - int32_t N, - int32_t K, - int32_t bits, - int32_t force_arch = -1) { - py::buffer_info q_weights_buf = q_weights.request(); - - if (bits != 4 && bits != 8) { - throw std::invalid_argument("bits must be 4 or 8"); - } - if (N <= 0 || K <= 0) { - throw std::invalid_argument("N and K must be positive"); - } - if (bits == 4 && K % 2 != 0) { - throw std::invalid_argument("K must be even for 4-bit packed weights"); - } - if (q_weights_buf.ndim != 2 || q_weights_buf.shape[0] != N || q_weights_buf.shape[1] != K / (8 / bits)) { - throw std::invalid_argument("q_weights must have shape (N, K / (8 / bits))"); - } - - int n = static_cast(N); - int k = static_cast(K); - - size_t packed_weight_bytes = static_cast(n) * static_cast(k) / (8 / bits); - py::array_t processed_weights({static_cast(packed_weight_bytes)}); - py::buffer_info processed_weights_buf = processed_weights.request(); - - auto make_cuda_ptr = [](size_t bytes) -> CudaPtr { - void* p = nullptr; - ThrowIfCudaError(cudaMalloc(&p, bytes), "cudaMalloc"); - return CudaPtr(p); - }; - - auto packed_transposed_weight_space = make_cuda_ptr(packed_weight_bytes); - int8_t* packed_transposed_weight = reinterpret_cast(packed_transposed_weight_space.get()); - - auto fpA_intB_weight_buffer_ = make_cuda_ptr(packed_weight_bytes); - int8_t* preprocessed_weight = reinterpret_cast(fpA_intB_weight_buffer_.get()); - - const uint8_t* blob_data_cpu = static_cast(q_weights_buf.ptr); - - auto blob_data_gpu_buf = make_cuda_ptr(packed_weight_bytes); - uint8_t* blob_data_gpu = reinterpret_cast(blob_data_gpu_buf.get()); - - cudaStream_t stream = cudaStreamLegacy; - ThrowIfCudaError(cudaMemcpyAsync(blob_data_gpu, blob_data_cpu, packed_weight_bytes, cudaMemcpyHostToDevice, stream), - "cudaMemcpyAsync host-to-device"); - - if (bits == 4) { - ::onnxruntime::llm::kernels::fpA_intB_gemv::unpack_uint4_transposed_to_int8_direct_cuda( - stream, packed_transposed_weight, blob_data_gpu, n, k); - } else { - // 8 bits - ::onnxruntime::llm::kernels::fpA_intB_gemv::transpose_uint8_matrix_and_convert_to_int8( - stream, packed_transposed_weight, blob_data_gpu, n, k); - } - - using ::onnxruntime::llm::kernels::weight_only::QuantType; - QuantType quant_type = bits == 4 ? QuantType::W4_A16 : QuantType::W8_A16; - - int sm = force_arch; - if (sm < 0) { - int device_id = 0; - ThrowIfCudaError(cudaGetDevice(&device_id), "cudaGetDevice"); - cudaDeviceProp device_prop; - ThrowIfCudaError(cudaGetDeviceProperties(&device_prop, device_id), "cudaGetDeviceProperties"); - sm = device_prop.major * 10 + device_prop.minor; - } - sm = ::onnxruntime::llm::kernels::weight_only::get_arch_for_mixed_gemm_weight_preprocess(sm); - - auto permutation_map_buffer = make_cuda_ptr(32 * sizeof(int32_t)); - - ::onnxruntime::llm::kernels::weight_only::preprocess_weights_for_mixed_gemm_cuda( - stream, - sm, - preprocessed_weight, - packed_transposed_weight, - reinterpret_cast(permutation_map_buffer.get()), - {static_cast(k), static_cast(n)}, - quant_type); - - ThrowIfCudaError(cudaGetLastError(), "preprocess CUDA kernel launch"); - ThrowIfCudaError(cudaMemcpyAsync(processed_weights_buf.ptr, preprocessed_weight, packed_weight_bytes, cudaMemcpyDeviceToHost, stream), - "cudaMemcpyAsync device-to-host"); - ThrowIfCudaError(cudaStreamSynchronize(stream), "cudaStreamSynchronize"); - - return processed_weights; -} - // Pack FP4 (MXFP4) weights for MoE GEMM kernels. // // Input: q_weights in [N, K/2] layout (FP4 packed 2 per byte along K dimension, row-major) @@ -280,6 +149,7 @@ py::array_t PackWeightsForMixedGemm( // // Unlike INT4 which requires architecture-specific row permutation and interleaving, // FP4 (SM90+ TMA path) only needs a simple transpose at the nibble level. +// This function is CPU-only and does not require CUDA to be present. py::array_t PackFP4WeightsForMoE( py::array_t q_weights, int32_t N, @@ -300,7 +170,7 @@ py::array_t PackFP4WeightsForMoE( int K_half = K / 2; int N_half = N / 2; size_t out_size = static_cast(K) * static_cast(N_half); - py::array_t output({static_cast(out_size)}); + py::array_t output(static_cast(out_size)); py::buffer_info out_buf = output.request(); uint8_t* dst = static_cast(out_buf.ptr); std::memset(dst, 0, out_size); @@ -327,8 +197,6 @@ py::array_t PackFP4WeightsForMoE( return output; } -} // namespace cuda -#endif void CreateQuantPybindModule(py::module& m) { m.def("quantize_matmul_2bits", &QuantizeMatMulNBitsBlockwise); @@ -343,14 +211,9 @@ void CreateQuantPybindModule(py::module& m) { m.def("quantize_qdq_matmul_2bits", &QuantizeQDQMatMulNBitsBlockwise); m.def("quantize_qdq_matmul_4bits", &QuantizeQDQMatMul4BitsBlockwise); m.def("quantize_qdq_matmul_4bits", &QuantizeQDQMatMul4BitsBlockwise); -#if defined(USE_CUDA) && !defined(ORT_NO_CUDA_IN_PYBIND) - m.def("pack_weights_for_cuda_mixed_gemm", &cuda::PackWeightsForMixedGemm, - "Pack quantized weights for CUDA mixed-precision GEMM (FpA_IntB format)", - py::arg("q_weights"), py::arg("N"), py::arg("K"), py::arg("bits"), py::arg("force_arch") = -1); - m.def("pack_fp4_weights_for_cuda_moe_gemm", &cuda::PackFP4WeightsForMoE, + m.def("pack_fp4_weights_for_cuda_moe_gemm", &PackFP4WeightsForMoE, "Pack FP4 (MXFP4) weights for CUDA MoE GEMM: transpose [N,K/2] to column-major [K,N/2]", py::arg("q_weights"), py::arg("N"), py::arg("K")); -#endif } } // namespace python diff --git a/onnxruntime/python/tools/quantization/cuda_quantizer.py b/onnxruntime/python/tools/quantization/cuda_quantizer.py index 6b03485280d62..2ab532b3d0e58 100644 --- a/onnxruntime/python/tools/quantization/cuda_quantizer.py +++ b/onnxruntime/python/tools/quantization/cuda_quantizer.py @@ -7,9 +7,13 @@ """CUDA weight-only quantization helpers. This module contains small Python utilities for producing the weight layouts -consumed by CUDA weight-only kernels. The helpers deliberately wrap the same C++ -pybind entry points used by runtime prepacking so tests and model builders can -generate byte-identical quantized weights. +consumed by CUDA weight-only kernels. The blockwise quantizers wrap the same C++ +pybind entry points used by runtime prepacking, and the mixed-GEMM weight packer +is a PyTorch reimplementation of the runtime CUDA packing, so tests and model +builders can generate byte-identical quantized weights. The PyTorch packer runs +on CUDA when a device is available and falls back to CPU otherwise, which is the +only option on platforms where the standalone CUDA packer is not built (Windows). +A GPU-gated parity test validates it against that standalone CUDA packer. Two storage families are exposed: @@ -24,10 +28,14 @@ from __future__ import annotations +import functools +import logging from typing import TYPE_CHECKING import numpy as np +_logger = logging.getLogger(__name__) + if TYPE_CHECKING: import torch @@ -43,20 +51,207 @@ def _get_torch(): def _get_pack_weights_for_cuda_mixed_gemm(): - """Return the CUDA mixed-GEMM weight prepacker from the ORT pybind module.""" + """Return the standalone CUDA mixed-GEMM weight packer (parity oracle). + + Production packing uses the PyTorch implementation (``_pack_weights_for_cuda_mixed_gemm``). + This standalone packer lives in ``onnxruntime.capi.onnxruntime_cuda_quant_preprocess``, a + separate extension module that links the CUDA runtime (built only on non-Windows CUDA + builds). It is imported lazily here (never at ``import onnxruntime`` time) and is used by + the parity test to validate the PyTorch packer byte-for-byte. + """ try: - from onnxruntime.capi import _pybind_state as _pybind # noqa: PLC0415 + from onnxruntime.capi import onnxruntime_cuda_quant_preprocess as _cuda_quant # noqa: PLC0415 except ImportError as e: raise ImportError( - "CUDA weight prepacking requires pack_weights_for_cuda_mixed_gemm from an onnxruntime-gpu CUDA build." + "The standalone CUDA weight packer (onnxruntime_cuda_quant_preprocess) is unavailable; " + "it is built only on non-Windows onnxruntime-gpu CUDA builds." ) from e try: - return _pybind.pack_weights_for_cuda_mixed_gemm + return _cuda_quant.pack_weights_for_cuda_mixed_gemm except AttributeError as e: - raise ImportError( - "CUDA weight prepacking requires pack_weights_for_cuda_mixed_gemm from an onnxruntime-gpu CUDA build." - ) from e + raise ImportError("onnxruntime_cuda_quant_preprocess is missing pack_weights_for_cuda_mixed_gemm.") from e + + +def has_cuda_weight_prepacking() -> bool: + """Return True if mixed-GEMM weight prepacking is available. + + Prepacking is implemented with PyTorch (CUDA when available, CPU otherwise), so it is + available whenever torch is importable. Callers use this to skip prepack code paths + (and tests) when torch is unavailable. + """ + try: + _get_torch() + except ImportError: + return False + return True + + +@functools.lru_cache(maxsize=1) +def _warn_cpu_prepack_once() -> None: + _logger.warning( + "CUDA device is not available; packing mixed-GEMM weights on CPU with PyTorch. " + "This is correct but significantly slower for large Mixture-of-Experts models. " + "Pack on a CUDA-enabled machine for best performance." + ) + + +def _prepack_device(): + """Pick the torch device for mixed-GEMM weight packing (CUDA if available, else CPU).""" + torch = _get_torch() + if torch.cuda.is_available(): + return torch.device("cuda") + _warn_cpu_prepack_once() + return torch.device("cpu") + + +def _preprocess_weights_for_mixed_gemm_torch(tensor, bits: int, sm: int): + """PyTorch port of the runtime CUDA ``preprocess_weights_for_mixed_gemm``. + + ``tensor`` is a signed int8 weight in ``(K, N/pack)`` packed row-major layout on any + device. Returns the CUTLASS mixed-GEMM layout with the same shape/dtype/device. This + mirrors ``preprocess_weights_for_mixed_gemm_cuda`` (permute_B_rows -> subbyte_transpose + -> interleave_column_major -> add_bias_and_interleave) so its output is byte-identical + to the standalone CUDA packer, for both the SM80 (Ampere) and SM90 (Hopper) layouts. + """ + torch = _get_torch() + bits_a = 16 # fp16/bf16 activations + bits_b = 4 if bits == 4 else 8 + + if tensor.dim() == 2: + tensor = tensor.unsqueeze(0) + + permutation_map = { + "16_8": [0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15], + "16_4": [ + 0, + 1, + 8, + 9, + 16, + 17, + 24, + 25, + 2, + 3, + 10, + 11, + 18, + 19, + 26, + 27, + 4, + 5, + 12, + 13, + 20, + 21, + 28, + 29, + 6, + 7, + 14, + 15, + 22, + 23, + 30, + 31, + ], + } + mma_shape_n = 8 + b_rows_per_mma = 8 * 16 // bits_b + + num_experts, num_rows, num_cols = tensor.shape[0], tensor.shape[1], tensor.shape[2] + if num_rows % b_rows_per_mma != 0 or num_cols % mma_shape_n != 0: + raise ValueError( + f"weight shape (rows={num_rows}, packed_cols={num_cols}) is incompatible with mixed-GEMM " + f"packing (rows must be a multiple of {b_rows_per_mma}, packed cols a multiple of {mma_shape_n})." + ) + + # permute_B_rows_for_mixed_gemm + if sm < 100: + pmap = permutation_map[f"{bits_a}_{bits_b}"] + row_idx = [(r // b_rows_per_mma) * b_rows_per_mma + pmap[r % b_rows_per_mma] for r in range(num_rows)] + tensor = tensor[:, row_idx, :] + + # subbyte_transpose + original_shape = tensor.shape + if bits_b == 4: + u = tensor.view(torch.uint8) + high = (u >> 4).permute(0, 2, 1).unsqueeze(2) + low = ((u << 4) >> 4).permute(0, 2, 1).unsqueeze(2) + merged = torch.cat([low, high], dim=2).reshape(u.shape[0], -1, u.shape[1]) + merged = merged[:, :, 0::2] + merged[:, :, 1::2] * 16 + tensor = merged.view(torch.int8).reshape(original_shape) + else: + tensor = tensor.permute(0, 2, 1).reshape(original_shape) + + # interleave_column_major_tensor + interleave = bits_a // bits_b + if interleave > 1 and sm < 90: + rows_per_tile = 128 * 8 // bits_a + elts_in_int32 = 32 // bits_b + if num_rows % elts_in_int32 != 0 or num_rows % rows_per_tile != 0: + raise ValueError(f"num_rows ({num_rows}) is incompatible with column-interleave tiling.") + tensor = tensor.reshape( + num_experts, -1, interleave, num_rows // rows_per_tile, rows_per_tile * 4 // elts_in_int32 + ) + tensor = tensor.permute(0, 1, 3, 2, 4).reshape(original_shape) + + # add_bias_and_interleave_quantized_tensor_inplace + if bits_b == 8: + t = tensor.to(torch.int64) # widen so the +128 rebias cannot overflow int8 + t += -256 * (t > 127).to(torch.int64) + 128 + t = t.reshape(-1, 4)[:, [0, 2, 1, 3]].reshape(original_shape) + tensor = t.to(torch.uint8).view(torch.int8) + else: + u = tensor.view(torch.uint8) + high = (u >> 4).unsqueeze(-1) + low = ((u << 4) >> 4).unsqueeze(-1) + merged = torch.cat([low, high], dim=-1).reshape(u.shape[0], u.shape[1], -1) + merged = merged.reshape(-1, 8)[:, [0, 2, 4, 6, 1, 3, 5, 7]].reshape(merged.shape) + merged = merged.to(torch.int16) + merged += -16 * (merged > 7).to(torch.int16) + 8 + merged = merged[:, :, 0::2] + merged[:, :, 1::2] * 16 + tensor = merged.to(torch.uint8).view(torch.int8) + + return tensor.squeeze(0).contiguous() + + +def _pack_weights_for_cuda_mixed_gemm(q_weights, n: int, k: int, bits: int, force_arch: int = 80) -> np.ndarray: + """PyTorch implementation of the CUDA ``pack_weights_for_cuda_mixed_gemm``. + + ``q_weights`` is ORT's unsigned MatMulNBits/QMoE storage ``(N, K/pack)`` (uint8). Returns + a flat ``int8`` numpy array with the CUTLASS mixed-GEMM layout, byte-identical to the + standalone CUDA packer. Runs on CUDA when available, otherwise on CPU. + """ + torch = _get_torch() + bits = int(bits) + force_arch = int(force_arch) + if bits not in (4, 8): + raise ValueError(f"bits must be 4 or 8, got {bits}.") + if force_arch not in (80, 90): + raise ValueError(f"force_arch must be 80 (SM80) or 90 (SM90), got {force_arch}.") + pack = 8 // bits + device = _prepack_device() + + q = torch.as_tensor(np.ascontiguousarray(q_weights)).view(torch.uint8).reshape(n, k // pack).to(device) + + # Front-end adaptor: transpose ORT (N, K) -> (K, N) and convert unsigned -> signed int8. + if bits == 4: + low = (q & 0x0F).to(torch.int16) + high = (q >> 4).to(torch.int16) + unpacked = torch.empty((n, k), dtype=torch.int16, device=device) + unpacked[:, 0::2] = low + unpacked[:, 1::2] = high + signed_t = (unpacked - 8).transpose(0, 1).contiguous() # (K, N), zero point 8 + packed_t = ((signed_t[:, 0::2] & 0x0F) | ((signed_t[:, 1::2] & 0x0F) << 4)).to(torch.uint8).view(torch.int8) + else: + signed_t = (q.to(torch.int16) - 128).transpose(0, 1).contiguous() # (K, N), zero point 128 + packed_t = signed_t.to(torch.uint8).view(torch.int8) + + out = _preprocess_weights_for_mixed_gemm_torch(packed_t.contiguous(), bits, force_arch) + return out.reshape(-1).cpu().numpy() def _get_quantize_matmul_nbits(): @@ -146,8 +341,7 @@ def qmoe_per_channel_quantize( When ``prepack`` is true, returned weights have shape ``[K, N/pack]``. Otherwise, returned weights keep raw per-channel storage ``[N, K/pack]``. - CUDA prepacking requires ``pack_weights_for_cuda_mixed_gemm`` from an - onnxruntime-gpu CUDA build. + Prepacking uses PyTorch (CUDA when available, CPU otherwise). """ torch = _get_torch() @@ -159,14 +353,12 @@ def qmoe_per_channel_quantize( if not prepack: return qweight, scales - pack_weights_for_cuda_mixed_gemm = _get_pack_weights_for_cuda_mixed_gemm() - n, k = weights.shape pack = 8 // int(bits) if n % pack != 0: raise ValueError(f"N ({n}) must be divisible by {pack} for CUDA QMoE prepacked weights.") - packed = pack_weights_for_cuda_mixed_gemm(qweight.numpy(), n, k, int(bits), force_arch) + packed = _pack_weights_for_cuda_mixed_gemm(qweight.numpy(), n, k, int(bits), force_arch) packed = np.asarray(packed).view(np.uint8).reshape(k, n // pack) return torch.from_numpy(np.ascontiguousarray(packed)), scales @@ -330,7 +522,7 @@ def matmulnbits_prepacked_blockwise_quantize( unsigned_full_range=unsigned_full_range, ) - pack_weights_for_cuda_mixed_gemm = _get_pack_weights_for_cuda_mixed_gemm() + pack_weights_for_cuda_mixed_gemm = _pack_weights_for_cuda_mixed_gemm packed = pack_weights_for_cuda_mixed_gemm(qweight.reshape(n, -1).numpy(), n, k, bits, force_arch) packed = np.asarray(packed).view(np.uint8).reshape(qweight.shape) packed = torch.from_numpy(np.ascontiguousarray(packed)) @@ -375,7 +567,7 @@ def qmoe_prepacked_blockwise_quantize( unsigned_full_range=unsigned_full_range, ) - pack_weights_for_cuda_mixed_gemm = _get_pack_weights_for_cuda_mixed_gemm() + pack_weights_for_cuda_mixed_gemm = _pack_weights_for_cuda_mixed_gemm packed = pack_weights_for_cuda_mixed_gemm(qweight.reshape(n, -1).numpy(), n, k, bits, force_arch) packed = np.asarray(packed).view(np.uint8).reshape(k, n // pack) torch = _get_torch() diff --git a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py index 406eee21c059f..9d3a378f6db19 100644 --- a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py +++ b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py @@ -16,6 +16,12 @@ import onnxruntime as ort from onnxruntime.capi import _pybind_state as _pybind +from onnxruntime.quantization.cuda_quantizer import _pack_weights_for_cuda_mixed_gemm + +try: + from onnxruntime.capi import onnxruntime_cuda_quant_preprocess as _cuda_quant +except ImportError: + _cuda_quant = None @contextmanager @@ -32,7 +38,7 @@ def set_env(name: str, value: str): @unittest.skipIf("CUDAExecutionProvider" not in ort.get_available_providers(), "CUDA is not available") -@unittest.skipUnless(hasattr(_pybind, "pack_weights_for_cuda_mixed_gemm"), "fpA_intB weight packer is unavailable") +@unittest.skipUnless(_cuda_quant is not None, "fpA_intB weight packer is unavailable") class TestMatMulNBitsPrepackedCuda(unittest.TestCase): def _quantize_weight(self, weight: np.ndarray, bits: int, block_size: int): k, n = weight.shape @@ -118,7 +124,7 @@ def _check_prepacked_parity( bias = rng.normal(0.0, 1.0, size=(n,)).astype(np.float16) if has_bias else None q_weight, scales = self._quantize_weight(weight, bits, block_size) - prepacked_flat = _pybind.pack_weights_for_cuda_mixed_gemm(q_weight.reshape(n, -1), n, k, bits, force_arch) + prepacked_flat = _cuda_quant.pack_weights_for_cuda_mixed_gemm(q_weight.reshape(n, -1), n, k, bits, force_arch) prepacked_weight = np.asarray(prepacked_flat, dtype=np.int8).view(np.uint8).reshape(q_weight.shape) raw_model = self._make_model((m, k), q_weight, scales, bits, block_size, weight_prepacked=0, bias=bias) @@ -171,5 +177,39 @@ def test_int8_sm90_prepacked_weight_matches_runtime_prepack(self): self._check_sm90_parity(bits=8, block_size=128, m=32) +@unittest.skipIf("CUDAExecutionProvider" not in ort.get_available_providers(), "CUDA is not available") +@unittest.skipUnless(_cuda_quant is not None, "standalone CUDA weight packer (parity oracle) is unavailable") +class TestCudaQuantizerTorchPackerParity(unittest.TestCase): + """Validate the PyTorch mixed-GEMM packer in cuda_quantizer.py against the CUDA oracle. + + ``cuda_quantizer._pack_weights_for_cuda_mixed_gemm`` (PyTorch, used in production, and the + only option on Windows where the standalone module is not built) must be byte-identical to + the standalone ``onnxruntime_cuda_quant_preprocess.pack_weights_for_cuda_mixed_gemm`` (the + CUDA code the runtime prepack uses). This test is the guard against silent drift; it only + runs where the oracle is built (non-Windows CUDA). + """ + + def _check(self, bits: int, force_arch: int, n: int, k: int): + pack = 8 // bits + rng = np.random.default_rng(20260708 + bits * 100 + force_arch + n + k) + q = rng.integers(0, 256, size=(n, k // pack), dtype=np.uint8) + oracle = np.asarray(_cuda_quant.pack_weights_for_cuda_mixed_gemm(q, n, k, bits, force_arch), dtype=np.int8) + torch_out = _pack_weights_for_cuda_mixed_gemm(q, n, k, bits, force_arch).astype(np.int8) + self.assertEqual(oracle.shape, torch_out.shape, f"shape mismatch bits={bits} arch={force_arch} N={n} K={k}") + np.testing.assert_array_equal( + torch_out, oracle, err_msg=f"byte mismatch bits={bits} arch={force_arch} N={n} K={k}" + ) + + def test_torch_packer_matches_cuda_oracle(self): + # Cover both weight bit-widths, both mixed-GEMM layouts (SM80/SM90), and a GPT-OSS-20B + # MoE shape (fused gate+up FC1 [5760, 2880] and down FC2 [2880, 2880]). + shapes = [(256, 256), (512, 256), (256, 512), (5760, 2880), (2880, 2880), (128, 128)] + for bits in (4, 8): + for force_arch in (80, 90): + for n, k in shapes: + with self.subTest(bits=bits, force_arch=force_arch, n=n, k=k): + self._check(bits, force_arch, n, k) + + if __name__ == "__main__": unittest.main() diff --git a/setup.py b/setup.py index 62ced38819f2c..58aadc75b5010 100644 --- a/setup.py +++ b/setup.py @@ -375,6 +375,7 @@ def finalize_options(self): if platform.system() == "Linux" or platform.system() == "AIX": libs = [ "onnxruntime_pybind11_state.so", + "onnxruntime_cuda_quant_preprocess.so", "libdnnl.so.2", "libmklml_intel.so", "libmklml_gnu.so", @@ -388,6 +389,13 @@ def finalize_options(self): dl_libs.append(providers_cann) dl_libs.append(providers_qnn) dl_libs.append("libonnxruntime.so*") + # onnxruntime_cuda_quant_preprocess.so is a standalone CUDA extension module used only as a + # byte-parity oracle for the PyTorch weight packer. It is built (and thus present here) only + # when the CMake option onnxruntime_BUILD_CUDA_QUANT_PREPROCESS is ON. The glob-based filters below + # drop missing files, so listing it here is a no-op when it was not built. It must be listed in + # dl_libs (not just libs) so that manylinux test wheels include it: the manylinux packaging path + # builds "data" from dl_libs only (see the is_manylinux block below). + dl_libs.append("onnxruntime_cuda_quant_preprocess.so") # DNNL, TensorRT, OpenVINO, and QNN EPs are built as shared libs libs.extend(["libonnxruntime_providers_shared.so"]) libs.extend(["libonnxruntime_providers_dnnl.so"]) @@ -422,6 +430,7 @@ def finalize_options(self): elif platform.system() == "Darwin": libs = [ "onnxruntime_pybind11_state.so", + "onnxruntime_cuda_quant_preprocess.so", "libdnnl.2.dylib", "mimalloc.so", "libonnxruntime*.dylib", diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 8317018d33c64..4061e4fafbe7d 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -1822,7 +1822,9 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs): if not args.disable_contrib_ops: run_subprocess( - [sys.executable, "-m", "unittest", "discover", "-s", "quantization"], cwd=cwd, dll_path=dll_path + [sys.executable, "-m", "unittest", "discover", "-s", "quantization", "-v"], + cwd=cwd, + dll_path=dll_path, ) if args.enable_transformers_tool_test and (sys.version_info.major, sys.version_info.minor) < (