diff --git a/docs/contrib_ops/cuda/moe_qmoe.md b/docs/contrib_ops/cuda/moe_qmoe.md index 73682f6a237c6..91259c675e5cb 100644 --- a/docs/contrib_ops/cuda/moe_qmoe.md +++ b/docs/contrib_ops/cuda/moe_qmoe.md @@ -843,6 +843,70 @@ When native CUTLASS is enabled, weights and scales are dual-prepacked (native la raw e8m0 block scales for GEMV), so `ORT_FP4_PREFILL_MIN_TOKENS` selects per call between the prefill and decode kernels with no extra conversion. +### Experimental large-expert INT2 execution + +`ORT_QMOE_INT2_LARGE_EXPERTS=1` extends the packed INT2 grouped-GEMM +execution and matching runtime prepacking from at most 256 to at most 512 +experts. It is disabled by default and requires the existing +`ORT_ENABLE_QMOE_INT2_PREFILL` path, symmetric INT2/INT4 weights with at +least one INT2 projection, raw initializers (`weights_prepacked=0`), fused +SwiGLU, supported block sizes, and aligned projection dimensions. +The packed grouped path also serves decode calls whose expanded row count +or routing is outside the GEMV eligibility limits, including top-10 +routing. It does not expand GEMV eligibility or bypass the dense-fallback +scratch guard. + +FP16 regression coverage compares 512-expert, top-10 calls at one, three, +and 32 rows against the dequantized fallback, with block sizes 64 and 128, +projection biases, and CUDA graph capture/replay. Full-model quality and +throughput must be evaluated separately; this flag is not a claim of +speedup or lossless quantization. + +### Experimental INT4 activation quality simulation for INT2 experts + +For local calibration diagnostics, `ORT_QMOE_CALIBRATION_CAPTURE_DIR` +can name an existing directory to capture NVFP4 QMoE inputs and router +logits. This is disabled when unset, is unavailable in plugin/minimal +builds, and requires CUDA graphs disabled. Each named node writes at +most sixteen invocations with at least sixteen rows, keeping the first +128 rows. Other quantization types are unaffected. Files are never +intentionally overwritten; missing directories, invalid shapes and I/O +failures are explicit errors. File names contain an encoded node name, +process ID, QMoE instance ID, and invocation index so concurrent sessions +do not collide; destination files are created exclusively. The binary +format is five int64 values `[version=1, rows, hidden, experts, +ONNX_activation_dtype]`, followed by +row-major input and router tensors in their original FP16/BF16 storage. +Capture synchronizes the stream and is **not suitable for performance +measurement**. Captured activations are local model data; keep them in +the experiment directory and leave the option unset in production. + +`ORT_QMOE_INT2_FAKE_INT4_ACTIVATIONS=1` enables a default-off numerical +screen for integer QMoE nodes with at least one INT2 projection. It +quantizes/dequantizes both the expanded gate/up input and the post-SwiGLU +down input, independently for each row and contiguous group of 64 values. +The dynamic FP32 scale is `max(abs(group))/7`; codes are rounded to nearest +with ties to even in the symmetric range `[-7,7]`. Dequantized values are +rounded back to the existing FP16/BF16 activation dtype. All-zero groups +remain zero. Scales are recomputed at inference time; no calibration +dataset or model-weight change is required. + +This is **fake quantization**, not packed INT4 activation storage or +integer tensor-core execution. Both GEMMs still consume FP16/BF16. +Additional QDQ launches may slow requests; do not interpret their +throughput or memory as native W2A4 performance. Routing probabilities, +residuals, attention, caches and non-INT2 nodes (including FP8 MTP) are +not quantized by this flag. + +The flag requires `ORT_ENABLE_QMOE_INT2_PREFILL=1`, the supported packed +grouped path, fused interleaved SwiGLU and aligned dimensions; 512-expert +models additionally require `ORT_QMOE_INT2_LARGE_EXPERTS=1`. GEMV is +bypassed for affected nodes to apply QDQ at both projections consistently +in prefill and decode. Unsupported affected nodes fail explicitly rather +than silently executing an unquantized fallback. Flag values other than +zero or one are rejected. With the flag unset/zero existing dispatch and +arithmetic remain unchanged. + ### 9.10 Interleaved GEMV layout + dtype-conditional accumulation **Opt-in (`ORT_FP4_GEMV_INTERLEAVED=1`, default off).** The shipping MXFP4 decode GEMV uses a @@ -1177,6 +1241,12 @@ full-model measurements and do not establish performance on other GPU architectu Disabling GEMV entirely is not equivalent: the dense fallback dequantizes every expert, which can be much slower for large expert counts. +Standalone FP4 GEMV and packed NVFP4 prefill bypass grouped-GEMM tactic profiling +because neither route consumes its tactics or workspace. The legacy +`ep.cuda.qmoe_skip_nvfp4_gemv_profiling` session setting and +`ORT_QMOE_SKIP_NVFP4_GEMV_PROFILING` environment variable remain accepted for +compatibility, but no longer control profiling on standalone FP4 GEMV. + Raw mode ignores `ORT_FP4_GEMV_AUTOTUNE`: the prepacked tiling candidates do not change the raw kernel, so profiling them only adds synchronization and repeated work. With `ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO=1`, the routes are reported as `fp4_gemv_raw` and @@ -1284,6 +1354,150 @@ gate: it applies a looser tolerance **only** to shapes that route native (SM120/ `num_rows >= ORT_FP4_PREFILL_MIN_TOKENS`) and keeps the strict bound for all decode/fallback shapes, so a genuinely broken native kernel (error order ~1.0+) is still caught. +### 9b.6 Experimental packed weight-only prefill + +The runtime automatically selects a memory-efficient W4A16 +prefill route on SM80+ when NVFP4 weights already have the offline +row-major byte layout. `ORT_QMOE_NVFP4_PACKED_PREFILL=0` disables this route; +`1` explicitly selects the default. It requires fused interleaved +SwiGLU (`swiglu_fusion=1`), positive hidden/intermediate dimensions that +are multiples of 64, and supported integer routing sizes. Unsupported +shapes/layouts retain the existing fallback; eligible small decode +continues to use the unchanged FP4 GEMV path. Native block-scaled FP4 +prefill is not replaced. + +The persistent grouped WMMA kernel reads original E2M1 bytes, block-16 +E4M3 scales and per-expert FP32 global scales. It dequantizes only a +64x128x64 shared-memory tile, retaining FP16/BF16 activations and FP32 +accumulation. Eight-element loads reuse the existing `Fp4I2FConverter`; +scaled weights are rounded once to the activation dtype using paired +round-to-nearest conversion. The per-expert global scale is loaded once +per tile. Padded shared +strides and input/epilogue storage reuse reduce tile overhead. No full +FP16/BF16 expert-weight bank, requantization or expert offloading is +introduced. + +The mainloop uses SM80 asynchronous global-to-shared copies to prefetch +the next input and packed-weight tile while computing the current tile. +Only input and packed-weight storage are double buffered; one +dequantized weight tile is reused after the block barrier. This bounds +dynamic shared memory to 44 KiB per block. Copy completion and block +barriers protect both tile reuse and the shared FP32 epilogue. + +Expert sorting, input expansion, biased SwiGLU and routing finalization +reuse the packed integer-prefill pipeline. Its workspace is bounded by +a 384 MiB scratch policy for this NVFP4 route, using internal +row tiles when necessary. Integer prefill retains its 256 MiB bound. +The NVFP4 bound accommodates a full 4,096-row Qwen prefill chunk and +avoids a second read of expert weights. This is **not** a device-wide +memory cap. With +`ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO=1`, the route is +`packed_nvfp4_prefill`. + +Focused H200 tests cover FP16/BF16 reference and dense-path numerical +parity, row/column tails, biases, empty experts, multi-row tiling, +packed CUDA graph capture/replay, unsupported-shape/layout fallback and +unchanged decode routing. Dense comparison sessions are uncaptured; +exact capture/replay equality is asserted for the packed route. +Floating-point accumulation can differ from dense GEMM, so passing +these tolerances does not establish bit-identical full-model outputs. +Full-model Qwen3.8 experiments reached a sampled 79.80 GiB with +1,024-token scheduler chunks. A sequential same-GPU comparison of the +non-pipelined and pipelined builds improved independent E2E throughput +from 83.54 to 90.52 TPS; the 100 TPS target remains unmet in that experiment. +Those measurements are workload-specific, not a production throughput guarantee. + +With a separately capacity-bounded Qwen indexer graph and 4,096-token +scheduling, ten independent requests subsequently reached 99.79 E2E / +133.33 decode TPS and a sampled peak of 79.95 GiB. That graph retains +only the indexer capacity needed by the candidate's existing 40 KV +blocks; it is not a 262K-context replacement. The measured Engine +request ceiling is 9,217 total tokens, and sampled memory headroom is +only about 56 MiB. This does not qualify actual 80 GiB hardware or +guarantee strict 100 TPS. + +#### Optional fused gate/up epilogue + +`ORT_QMOE_NVFP4_FUSED_SWIGLU=1` additionally fuses the eligible packed +NVFP4 gate/up GEMM epilogue with bias and interleaved SwiGLU. Its default +is `0`; it has no effect without the packed-prefill route and does not +change integer prefill, unsupported-layout fallback or decode GEMV. +The FP32 accumulators are still rounded to FP16/BF16 before adding bias +and applying the same alpha/beta/limit semantics. Each paired gate/up +channel belongs to the same tile. + +The epilogue writes the activated down-projection input directly, avoiding +the gate/up output reservation and the separate activation launch. At +4,096 rows, top-10 and intermediate size 640 this removes exactly 100 MiB +from the packed workspace. H200 device-wide sampled peak fell from a +fresh 79.9432 GiB unfused control to 79.7498 GiB; allocator-pool effects +make this device-wide difference larger than the one workspace buffer. +Two five-request fused invocations reached 100.71 and 100.44 E2E TPS, +pooling to 100.58 E2E / 133.60 decode TPS and 1.266 s TTFT. Fresh +unfused controls pooled to 99.54 E2E TPS and 1.299 s TTFT. +All 128 paired quality token streams matched, preserving all 110 +control-correct answers. This is workload-specific evidence, not a +throughput guarantee or qualification on actual 80 GiB hardware; +sampled headroom remains only about 256 MiB. + +#### Optional routing scratch lifetime reuse + +`ORT_QMOE_NVFP4_REUSE_ROUTING_SCRATCH=1` reuses the packed NVFP4 sorting +temporaries (counts, cumulative counts and blocked rows) for the expanded +input and later down-projection output. Default is `0`. The three sorting +kernels finish on the same stream before expansion starts; row maps, +expert offsets and tile offsets remain separate and live through +finalization. The shared reservation is the larger of the aligned sorting +scratch and expanded-input storage, so shapes with more sorting scratch +than input storage are also supported. Integer prefill is unchanged. + +At 4,096 rows, 512 experts, top-10 and hidden size 2,560 this removes +8,404,992 bytes (about 8.02 MiB) of packed workspace. A fresh H200 +flag-off/flag-on memory comparison with fused SwiGLU observed +79.7479 -> 79.7322 GiB, a 16 MiB device-wide sampled reduction. +Reversed-order ten-request pools measured 100.01 -> 100.28 E2E TPS; +that small timing difference does not establish an independent speedup. +Sampled headroom is about 274 MiB, still not actual 80 GiB qualification. + +Direct input gathering into shared memory was also screened. Exact +numerical tests passed, but 4,096-token scheduling showed slower TTFT +(1.267 -> 1.289 seconds) and E2E (100.49 -> 99.95 TPS), with no workspace +reduction because expanded-input storage already aliases down output. +That implementation was rejected and is not present in the retained path. + +#### Optional tiled down-projection output + +`ORT_QMOE_NVFP4_TILED_DOWN_OUTPUT=1` enables a further default-off memory +tradeoff when both supported packed prefill and the fused SwiGLU epilogue +are active. Gate/up gathers original input rows into shared memory using +the unchanged sorted maps instead of materializing expanded input. +Down projection writes at most 1,024 columns of expanded output at a time; +each tile is reduced in original top-k rank order before its storage is +reused on the same stream. Bias addition, activation-dtype GEMM rounding, +FP32 routing accumulation and final rounding match the existing path; +there are no unordered atomic additions. Sorting-lifetime reuse can +still be enabled, with a reservation covering the larger live buffer. + +For 4,096 rows, top-10, hidden size 2,560 this reduces the shared +expanded-input/down-output reservation from 200 to 80 MiB, saving exactly +120 MiB workspace. H200 sampled device-wide peak falls from 79.7322 to +79.5760 GiB, a 160 MiB reduction including allocator-pool effects. The +candidate has about 434 MiB sampled headroom, not a hard memory cap. + +This is **not an established speed optimization**: final paired pools +observe TTFT rising from 1.271 to 1.295 seconds. Two five-request tiled +runs reach 100.02 and 104.46 E2E TPS, but decode variability accounts for +the high second result; another independent tiled screen reaches 99.996 +TPS. Strict at-least-100 throughput and actual 80 GiB hardware remain +unqualified. A smaller 512-column screen saved more workspace but was +slower (99.81 versus 100.31 E2E TPS); it is not selected. + +173 focused/regression tests pass, including exact tiled/untiled +FP16/BF16 output and workspace parity, biases, 512 experts, small hidden +dimensions, column tails, row tiling, empty experts, graph replay and +unaffected integer/decode/fallback routes. Fresh paired quality runs match +all 128 token streams and preserve all 110 control-correct answers. + --- ## 10. FP8 (W8A16) Details diff --git a/docs/contrib_ops/cuda/sparse_paged_attention.md b/docs/contrib_ops/cuda/sparse_paged_attention.md index 927ea4add7386..8bf95f36006c7 100644 --- a/docs/contrib_ops/cuda/sparse_paged_attention.md +++ b/docs/contrib_ops/cuda/sparse_paged_attention.md @@ -38,6 +38,36 @@ The contiguous auxiliary cache has shape | Attention sink | Supported in the joint softmax | | CUDA Graphs | Fixed-capacity, device-resident counts; no host synchronization | -The functional kernel reads selected pages directly and performs an online joint -softmax. Selector fusion, paged auxiliary storage, auxiliary quantization, and -decode-specialized sparse kernels are deferred. +## CUDA dispatch + +Eligible `selected_only`/`main` calls with matching activation/cache types, head +size 256, and query-to-KV head ratios divisible by four use tiled grouped-query +attention. Other geometries and quantized caches retain the general fallback. +All paths preserve suppressed-write prefix-cache reads; INT8 current rows are +read back from the quantized cache. + +The older six-head FP16 grouped-prefill path remains opt-in through session +configuration `ep.cuda.sparse_paged_attention_grouped=1`, or environment variable +`ORT_SPARSE_PAGED_ATTENTION_GROUPED=1` when no session override is present. +`ep.cuda.sparse_paged_attention_grouped_vectorized=1` selects vectorized loads +within that path. Both session options accept only `0` or `1`. + +`ORT_SPARSE_PREFILL_TENSOR_CORE_QK=1` enables experimental tensor-core query/key +scoring for eligible twelve-head prefill tiles. With that option enabled, +`ORT_SPARSE_PREFILL_TENSOR_CORE_PV=1` also permits the FP16 probability/value +path. These environment options default to off and are cached on first use. +Neither option establishes broad model-quality or throughput equivalence. + +The CUDA Python regression `test_sparse_paged_attention_prefill_cuda.py` runs +scalar, QK-only, QK+PV, BF16 QK, and both grouped-load configurations in fresh +processes. Each compares to an independent reference with multiple/partial +tiles, invalid selections, suppressed-write prefix hits, and attention sinks. +Verbose dispatch messages are asserted so an unintended fallback cannot pass. +Run it against a build containing this change on SM80 or newer: + +```bash +python -m pytest onnxruntime/test/python/transformers/test_sparse_paged_attention_prefill_cuda.py +``` + +Selector fusion, paged auxiliary storage, and auxiliary quantization remain +outside this operator's scope. diff --git a/docs/contrib_ops/packed_sparse_attention_indexer.md b/docs/contrib_ops/packed_sparse_attention_indexer.md index bd6877cf32189..0f743f0000ee6 100644 --- a/docs/contrib_ops/packed_sparse_attention_indexer.md +++ b/docs/contrib_ops/packed_sparse_attention_indexer.md @@ -19,6 +19,34 @@ Source: --- +## Experimental CUDA warp-local QSA prefill scoring + +`ORT_PACKED_SPARSE_INDEXER_WARP_HEADS=1` enables a default-off scoring path for +QSA packed prefill with more than 64 tokens, four query heads, head size 128, +and compression ratio four. Set it before session creation; its value is +cached per process. Other geometries and decode keep their existing dispatch. + +Each warp scores one state slot across all four heads, retaining query values +and reusing the loaded key. It removes the shared head-score handoff and +CTA-wide barriers while preserving the per-head warp reduction and ordered +head sum. The path does not change state capacity, causal visibility, top-k +selection, overflow handling, or workspace layout. + +FP16/FP32 flag-off/on regression tests cover random/tied scores, empty and +partial state tiles, and 32768 live slots (128K-token-equivalent state). +Fresh 64K/128K MTP-width-three full-model experiments reduced TTFT from about +11.62 to 10.50 seconds and 27.52 to 22.90 seconds, respectively, in both +control-first and candidate-first orders. These observations are not universal +performance guarantees. The first comparison matched all six measured token +streams; a reverse-order 128K control produced a different stream than both +the earlier control and candidates. Full-model deterministic/quality +equivalence across arbitrary tuned executions is therefore not established. +A further 128K comparison with decoder/MTP GEMM autotuning disabled matched +all three measured streams, with TTFT 27.52 to 22.94 seconds and essentially +unchanged decode throughput. This supports exact parity on that controlled +workload but does not establish the cause of the earlier mismatch or broad +quality equivalence. Keep the option experimental and off by default. + ## 1. Why a separate operator `SparseAttentionIndexer` uses dense `[batch_size, sequence_length, ...]` tensors, an explicit dense @@ -271,16 +299,29 @@ mechanics with no packed equivalent): respective execution provider); - `qsa` state-capacity overflow safety; - active and disabled transactional QSA capture on CUDA and WebGPU; +- the CUDA-only `QsaRadixTopKReferenceParity` test, which re-executes its cases in isolated + processes with `ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK=0` and `=1`, comparing random scores, + threshold ties, an empty row, a partial 32-key tile, and a fully selected radix bucket against + the reference indices and counts. Its eligible geometry selects the hierarchical path, so the + flag-on process must log `QsaPartialTopKKernel` (the flag-off process must log + `QsaMergeTileTopKKernel`) or the test fails. Run it with + `--gtest_filter=PackedSparseAttentionIndexerCudaTest.QsaRadixTopKReferenceParity`; - FP32/FP16 (and CUDA-only BF16) numeric coverage against an in-file reference that mirrors this document's contract. ## 10. Known limitations and follow-ups - The reference CUDA/WebGPU kernels prioritize correctness over throughput (see the top-of-file - comments in the `.cu`/`.cc` implementations). CUDA QSA fuses query rotation into scoring, uses - the shared-memory TopK path for up to 32 selected blocks, bounded partial TopK for 33–512 blocks, - and a hierarchical TopK path for common long-context Qwen shapes. Values above 512 that do not - match the hierarchical specialization retain the correctness-first repeated-scan fallback. + comments in the `.cu`/`.cc` implementations). CUDA QSA fuses query rotation into scoring and + uses a single-read deterministic block TopK for up to 32 selected blocks. Eligible short QSA + calls with four 128-channel heads, compression ratio four and up to 512 selected blocks use + hierarchical scoring/selection. Other bounded cases use radix selection. +- Experimental CUDA flag `ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK=1` replaces the eligible short-QSA + merge tree with live-key radix selection and a bounded sort. It preserves the scored keys, + ascending-index tie ordering, selected counts and state update semantics; no host readback is + introduced. The default is off and preserves hierarchical merging. Set the flag before the + first eligible call in a process; dispatch is cached. Long-context performance is not + established, so this is not a general throughput recommendation. - OgaEngine / Model Builder integration (declaring `past_key_state` etc. as Engine-managed, per-request fixed-size state, analogous to a paged auxiliary cache) is out of scope for this operator definition and is expected in a follow-up to `microsoft/onnxruntime-genai`. diff --git a/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.cc b/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.cc index c03d4b0d2b3fd..17d2641a671cd 100644 --- a/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.cc @@ -6,6 +6,7 @@ #include #include #include +#include "core/platform/env_var_utils.h" #include "contrib_ops/cpu/bert/paged_attention_helper.h" @@ -109,6 +110,19 @@ SparsePagedAttention::SparsePagedAttention(const OpKernelInfo& info) "contiguous", "Only auxiliary_cache_layout='contiguous' is supported."); auxiliary_kv_shared_ = get_boolean_attribute("auxiliary_kv_shared", 0); + const auto grouped = info.GetConfigOptions().GetConfigEntry("ep.cuda.sparse_paged_attention_grouped"); + if (grouped.has_value()) { + ORT_ENFORCE(*grouped == "0" || *grouped == "1", + "ep.cuda.sparse_paged_attention_grouped must be 0 or 1"); + enable_grouped_ = *grouped == "1"; + } else { + enable_grouped_ = ParseEnvironmentVariableWithDefault("ORT_SPARSE_PAGED_ATTENTION_GROUPED", false); + } + const auto vectorized = + info.GetConfigOptions().GetConfigOrDefault("ep.cuda.sparse_paged_attention_grouped_vectorized", "0"); + ORT_ENFORCE(vectorized == "0" || vectorized == "1", + "ep.cuda.sparse_paged_attention_grouped_vectorized must be 0 or 1"); + enable_grouped_vectorized_ = vectorized == "1"; } template @@ -293,9 +307,12 @@ Status SparsePagedAttention::ComputeInternal( const int64_t max_candidate_entries = std::min(static_cast(max_selected_entries) + max_local_entries, std::numeric_limits::max()); + const int heads_per_block = SparsePagedAttentionHeadsPerBlock( + parameters, attention_mode_, selected_kv_source_, std::is_same::value, + device_prop.sharedMemPerBlock); const int num_splits = ComputeSparsePagedAttentionSplits( parameters.token_count, parameters.num_heads, static_cast(max_candidate_entries), - device_prop.multiProcessorCount); + device_prop.multiProcessorCount, heads_per_block); const size_t partial_rows = num_splits > 1 ? static_cast(num_splits) * parameters.token_count * parameters.num_heads : 0; @@ -358,7 +375,7 @@ Status SparsePagedAttention::ComputeInternal( : auxiliary_lengths->Data(), auxiliary_capacity, auxiliary_num_heads, attention_mode_, selected_kv_source_, auxiliary_kv_shared_, partial_out.get(), - partial_max.get(), partial_sum.get(), num_splits); + partial_max.get(), partial_sum.get(), num_splits, enable_grouped_, enable_grouped_vectorized_); } } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.h b/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.h index b8aac466226c6..e5400b9ce19fa 100644 --- a/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.h @@ -32,6 +32,8 @@ class SparsePagedAttention final : public onnxruntime::cuda::CudaKernel { SparseAttentionMode attention_mode_; SelectedKvSource selected_kv_source_; bool auxiliary_kv_shared_; + bool enable_grouped_ = false; + bool enable_grouped_vectorized_ = false; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention_impl.cu index c628fbca9ab0b..2aec91ae050ad 100644 --- a/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention_impl.cu @@ -4,10 +4,14 @@ #include "contrib_ops/cuda/bert/sparse_paged_attention_impl.h" #include +#include +#include #include #include "contrib_ops/cuda/bert/paged_attention_impl.h" +#include "contrib_ops/cuda/bert/gated_delta_net_mma.cuh" #include "core/framework/stream_handles.h" +#include "core/platform/env_var_utils.h" #include "core/providers/cuda/shared_inc/cuda_call.h" namespace onnxruntime { @@ -34,11 +38,62 @@ constexpr int kSparsePagedAttentionTile = 128; constexpr int kSparsePagedAttentionMaxSplits = 32; constexpr int64_t kSparsePagedAttentionDirectCandidate = int64_t{1} << 62; +__device__ __forceinline__ int64_t SparsePagedAttentionResolveMainCandidate( + int logical_position, int batch_id, int query_position, int main_length, int past_length, + const int* cumulative_seqlens_q, const int* block_table, const int* slot_mapping, + int block_size, int num_blocks, int max_num_blocks_per_seq, bool is_causal, + bool quantized_cache = false) { + if (logical_position < 0 || logical_position >= main_length || + (is_causal && logical_position > query_position)) { + return -1; + } + const int logical_block = logical_position / block_size; + const int physical_block = logical_block < max_num_blocks_per_seq + ? block_table[static_cast(batch_id) * max_num_blocks_per_seq + logical_block] + : -1; + if (physical_block < 0 || physical_block >= num_blocks) { + return -1; + } + if (logical_position >= past_length) { + const int current_token = cumulative_seqlens_q[batch_id] + logical_position - past_length; + const int64_t mapped_slot = static_cast(physical_block) * block_size + logical_position % block_size; + const int64_t slot = slot_mapping == nullptr ? mapped_slot : slot_mapping[current_token]; + const int64_t cache_capacity = static_cast(num_blocks) * block_size; + if (slot >= 0 && slot < cache_capacity) { + return quantized_cache ? slot : kSparsePagedAttentionDirectCandidate + current_token; + } + return slot < 0 && mapped_slot >= 0 && mapped_slot < cache_capacity ? mapped_slot : -1; + } + return static_cast(physical_block) * block_size + logical_position % block_size; +} + __host__ __device__ __forceinline__ int SparsePagedAttentionChannelGroups(const int head_size) { return head_size >= kSparsePagedAttentionThreads ? 1 : (kSparsePagedAttentionThreads / head_size); } -template +template +__device__ __forceinline__ float SparsePagedAttentionBlockReduce(float value, float* scratch) { + static_assert(kSparsePagedAttentionThreads == 128); + scratch[threadIdx.x] = value; + __syncthreads(); + const int lane = threadIdx.x % 32; + // Preserve the shared-memory tree's 64, 32, 16, ... reduction order. + const float first = IsMax ? fmaxf(scratch[lane], scratch[lane + 64]) + : scratch[lane] + scratch[lane + 64]; + const float second = IsMax ? fmaxf(scratch[lane + 32], scratch[lane + 96]) + : scratch[lane + 32] + scratch[lane + 96]; + value = IsMax ? fmaxf(first, second) : first + second; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + const float other = __shfl_xor_sync(0xFFFFFFFFU, value, offset); + value = IsMax ? fmaxf(value, other) : value + other; + } + value = __shfl_sync(0xFFFFFFFFU, value, 0); + __syncthreads(); + return value; +} + +template __global__ void SparsePagedAttentionSplitKernel( const T* query, const T* current_key, const T* current_value, const TCACHE* key_cache, const TCACHE* value_cache, @@ -47,7 +102,7 @@ __global__ void SparsePagedAttentionSplitKernel( const int* selected_counts, const T* auxiliary_key, const T* auxiliary_value, const int* auxiliary_lengths, float* partial_out, float* partial_max, float* partial_sum, T* output, const T* head_sink, int batch_size, - int token_count, int num_heads, int kv_num_heads, int head_size, int block_size, + int token_count, int num_heads, int kv_num_heads, int runtime_head_size, int block_size, int num_blocks, int max_num_blocks_per_seq, int max_selected_entries, int auxiliary_capacity, int auxiliary_num_heads, int current_key_stride, int current_value_stride, float scale, float softcap, @@ -59,6 +114,7 @@ __global__ void SparsePagedAttentionSplitKernel( const int token_id = static_cast(token_head_id / num_heads); const int split_id = blockIdx.y; const int tid = threadIdx.x; + const int head_size = StaticHeadSize == 0 ? runtime_head_size : StaticHeadSize; int batch_id = 0; while (batch_id + 1 < batch_size && token_id >= cumulative_seqlens_q[batch_id + 1]) { @@ -147,43 +203,11 @@ __global__ void SparsePagedAttentionSplitKernel( logical_position < auxiliary_capacity) { candidate_ref = -static_cast(logical_position) - 2; } - } else if (logical_position < main_length && (!is_causal || logical_position <= query_position)) { - if (logical_position >= past_seqlens[batch_id]) { - const int current_token = cumulative_seqlens_q[batch_id] + logical_position - past_seqlens[batch_id]; - const int logical_block = logical_position / block_size; - const int physical_block = logical_block < max_num_blocks_per_seq - ? block_table[static_cast(batch_id) * max_num_blocks_per_seq + - logical_block] - : -1; - const int64_t cache_capacity = static_cast(num_blocks) * block_size; - const int64_t mapped_slot = physical_block >= 0 - ? static_cast(physical_block) * block_size + - logical_position % block_size - : -1; - const int64_t slot = slot_mapping == nullptr ? mapped_slot : slot_mapping[current_token]; - if (physical_block >= 0 && physical_block < num_blocks) { - if (slot >= 0 && static_cast(slot) < cache_capacity) { - if constexpr (std::is_same::value) { - // Read back the quantized row so current and past tokens have identical semantics. - candidate_ref = slot; - } else { - candidate_ref = kSparsePagedAttentionDirectCandidate + current_token; - } - } else if (slot < 0 && mapped_slot >= 0 && mapped_slot < cache_capacity) { - // A suppressed write can denote a prefix-cache hit; read its logical cache row. - candidate_ref = mapped_slot; - } - } - } else { - const int logical_block = logical_position / block_size; - if (logical_block < max_num_blocks_per_seq) { - const int physical_block = - block_table[static_cast(batch_id) * max_num_blocks_per_seq + logical_block]; - if (physical_block >= 0 && physical_block < num_blocks) { - candidate_ref = static_cast(physical_block) * block_size + logical_position % block_size; - } - } - } + } else { + candidate_ref = SparsePagedAttentionResolveMainCandidate( + logical_position, batch_id, query_position, main_length, past_seqlens[batch_id], + cumulative_seqlens_q, block_table, slot_mapping, block_size, num_blocks, + max_num_blocks_per_seq, is_causal, std::is_same::value); } } @@ -239,16 +263,7 @@ __global__ void SparsePagedAttentionSplitKernel( for (int index = tid; index < tile_length; index += kSparsePagedAttentionThreads) { tile_max = fmaxf(tile_max, logits[index]); } - reduction[tid] = tile_max; - __syncthreads(); - for (int stride = kSparsePagedAttentionThreads / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - reduction[tid] = fmaxf(reduction[tid], reduction[tid + stride]); - } - __syncthreads(); - } - tile_max = reduction[0]; - __syncthreads(); + tile_max = SparsePagedAttentionBlockReduce(tile_max, reduction); if (tile_max == -FLT_MAX) { continue; } @@ -261,21 +276,17 @@ __global__ void SparsePagedAttentionSplitKernel( logits[index] = weight; local_sum += weight; } - reduction[tid] = local_sum; - __syncthreads(); - for (int stride = kSparsePagedAttentionThreads / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - reduction[tid] += reduction[tid + stride]; - } - __syncthreads(); - } - running_sum = running_sum * old_weight + reduction[0]; + const float tile_sum = SparsePagedAttentionBlockReduce(local_sum, reduction); + running_sum = running_sum * old_weight + tile_sum; running_max = new_max; __syncthreads(); if (channel_groups == 1) { - for (int c = tid; c < head_size; c += kSparsePagedAttentionThreads) { + for (int c = tid; c < head_size; c += 2 * kSparsePagedAttentionThreads) { + const int second_c = c + kSparsePagedAttentionThreads; + const bool has_second_channel = second_c < head_size; float value_sum = accumulator[c] * old_weight; + float second_value_sum = has_second_channel ? accumulator[second_c] * old_weight : 0.0f; for (int index = 0; index < tile_length; ++index) { const int64_t candidate_ref = candidate_refs[index]; if (candidate_ref == -1) { @@ -298,17 +309,30 @@ __global__ void SparsePagedAttentionSplitKernel( const T* auxiliary_row = auxiliary_kv_shared ? auxiliary_key + auxiliary_base : auxiliary_value + auxiliary_base; value_sum += logits[index] * SparseActivationValue(auxiliary_row[c]); + if (has_second_channel) { + second_value_sum += logits[index] * SparseActivationValue(auxiliary_row[second_c]); + } } else if (candidate_is_direct) { const int64_t current_base = candidate_location * current_value_stride + kv_head_id * head_size; value_sum += logits[index] * SparseActivationValue(current_value[current_base + c]); + if (has_second_channel) { + second_value_sum += logits[index] * SparseActivationValue(current_value[current_base + second_c]); + } } else { const int64_t cache_base = (candidate_location * kv_num_heads + kv_head_id) * head_size; value_sum += logits[index] * SparseCacheValue(value_cache[cache_base + c], v_scale, kv_head_id * head_size + c, v_per_channel); + if (has_second_channel) { + second_value_sum += logits[index] * SparseCacheValue(value_cache[cache_base + second_c], v_scale, + kv_head_id * head_size + second_c, v_per_channel); + } } } accumulator[c] = value_sum; + if (has_second_channel) { + accumulator[second_c] = second_value_sum; + } } } else if (tid < accumulator_elements) { const int group = tid / head_size; @@ -378,6 +402,604 @@ __global__ void SparsePagedAttentionSplitKernel( } } +template +__device__ void SparsePagedAttentionStageRow(const T* row, const TCACHE* cached_row, + float* staged, int lane_id) { + if constexpr (kVectorized && std::is_same::value && std::is_same::value) { + const half* source = row != nullptr ? row : cached_row; + if (source != nullptr && reinterpret_cast(source) % alignof(half2) == 0) { +#pragma unroll + for (int channel = 2 * lane_id; channel < 256; channel += 64) { + const float2 values = __half22float2(*reinterpret_cast(source + channel)); + staged[channel] = values.x; + staged[channel + 1] = values.y; + } + return; + } + } +#pragma unroll + for (int channel = lane_id; channel < 256; channel += 32) { + staged[channel] = row != nullptr + ? SparseActivationValue(row[channel]) + : (cached_row != nullptr ? SparseActivationValue(cached_row[channel]) : 0.0f); + } +} + +template +__global__ void SparsePagedAttentionGroupedKernel( + const T* query, const T* current_key, const T* current_value, + const TCACHE* key_cache, const TCACHE* value_cache, + const int* cumulative_seqlens_q, const int* past_seqlens, + const int* block_table, const int* slot_mapping, const int* selected_indices, + const int* selected_counts, const T* auxiliary_key, const T* auxiliary_value, + const int* auxiliary_lengths, float* partial_out, float* partial_max, + float* partial_sum, T* output, const T* head_sink, int batch_size, + int token_count, int num_heads, int kv_num_heads, int block_size, + int num_blocks, int max_num_blocks_per_seq, int max_selected_entries, + int auxiliary_capacity, int auxiliary_num_heads, int current_key_stride, + int current_value_stride, float scale, float softcap, + int local_window_size, bool is_causal, bool local_plus_selected, + bool selected_from_auxiliary, bool auxiliary_kv_shared, int num_splits) { + constexpr int kHeads = 6; + constexpr int kChannels = 256; + constexpr int kTile = 8; + constexpr int kScoreTile = kSparsePagedAttentionTile; + __shared__ float staged_key[kTile][kChannels]; + __shared__ float staged_value[kTile][kChannels]; + __shared__ float logits[kHeads][kScoreTile]; + __shared__ const T* value_rows[kScoreTile]; + __shared__ const TCACHE* cached_value_rows[kScoreTile]; + __shared__ bool valid[kScoreTile]; + const int warp_id = threadIdx.x / 32; + const int lane_id = threadIdx.x % 32; + const int head_group = blockIdx.x % (num_heads / kHeads); + const int head_id = head_group * kHeads + warp_id; + const int kv_head_id = head_id / (num_heads / kv_num_heads); + const int token_id = blockIdx.x / (num_heads / kHeads); + const int split_id = blockIdx.y; + int batch_id = 0; + while (batch_id + 1 < batch_size && token_id >= cumulative_seqlens_q[batch_id + 1]) { + ++batch_id; + } + const int query_index = token_id - cumulative_seqlens_q[batch_id]; + const int query_position = past_seqlens[batch_id] + query_index; + const int main_length = past_seqlens[batch_id] + + cumulative_seqlens_q[batch_id + 1] - cumulative_seqlens_q[batch_id]; + const int local_end = local_plus_selected ? (is_causal ? query_position : main_length - 1) : -1; + const int local_begin = local_window_size > 0 ? max(0, query_position - local_window_size + 1) : 0; + const int local_count = max(0, local_end - local_begin + 1); + const int candidate_count = local_count + min(max(selected_counts[token_id], 0), max_selected_entries); + const int candidates_per_split = (candidate_count + num_splits - 1) / num_splits; + const int candidate_begin = split_id * candidates_per_split; + const int candidate_end = min(candidate_count, candidate_begin + candidates_per_split); + const int64_t query_base = (static_cast(token_id) * num_heads + head_id) * kChannels; + float query_values[kChannels / 32]; + float accumulator[kChannels / 32] = {}; +#pragma unroll + for (int channel = 0; channel < kChannels / 32; ++channel) { + query_values[channel] = SparseActivationValue(query[query_base + lane_id + channel * 32]); + } + float running_max = -FLT_MAX; + float running_sum = 0.0f; + for (int tile_begin = candidate_begin; tile_begin < candidate_end; tile_begin += kScoreTile) { + const int tile_length = min(kScoreTile, candidate_end - tile_begin); + for (int chunk_begin = 0; chunk_begin < tile_length; chunk_begin += kTile) { + const int chunk_end = min(tile_length, chunk_begin + kTile); + for (int tile_offset = chunk_begin + warp_id; tile_offset < chunk_end; tile_offset += kHeads) { + const int candidate = tile_begin + tile_offset; + const bool is_local = candidate < local_count; + const int logical_position = is_local + ? local_begin + candidate + : selected_indices[static_cast(token_id) * max_selected_entries + + candidate - local_count]; + const bool use_auxiliary = !is_local && selected_from_auxiliary; + const T* key_row = nullptr; + const T* value_row = nullptr; + const TCACHE* cached_key_row = nullptr; + const TCACHE* cached_value_row = nullptr; + if (logical_position >= 0 && + !(!is_local && !use_auxiliary && local_plus_selected && + logical_position >= local_begin && logical_position <= local_end)) { + if (use_auxiliary) { + if (logical_position < auxiliary_lengths[batch_id] && logical_position < auxiliary_capacity) { + const int auxiliary_head = auxiliary_num_heads == 1 ? 0 : kv_head_id; + const int64_t base = ((static_cast(batch_id) * auxiliary_capacity + logical_position) * + auxiliary_num_heads + + auxiliary_head) * + kChannels; + key_row = auxiliary_key + base; + value_row = auxiliary_kv_shared ? key_row : auxiliary_value + base; + } + } else { + const int64_t candidate_ref = SparsePagedAttentionResolveMainCandidate( + logical_position, batch_id, query_position, main_length, past_seqlens[batch_id], + cumulative_seqlens_q, block_table, slot_mapping, block_size, num_blocks, + max_num_blocks_per_seq, is_causal); + if (candidate_ref >= kSparsePagedAttentionDirectCandidate) { + const int64_t current_token = candidate_ref - kSparsePagedAttentionDirectCandidate; + key_row = current_key + current_token * current_key_stride + kv_head_id * kChannels; + value_row = current_value + current_token * current_value_stride + kv_head_id * kChannels; + } else if (candidate_ref >= 0) { + const int64_t base = (candidate_ref * kv_num_heads + kv_head_id) * kChannels; + cached_key_row = key_cache + base; + cached_value_row = value_cache + base; + } + } + } + const bool candidate_valid = key_row != nullptr || cached_key_row != nullptr; + if (lane_id == 0) { + valid[tile_offset] = candidate_valid; + value_rows[tile_offset] = value_row; + cached_value_rows[tile_offset] = cached_value_row; + } + SparsePagedAttentionStageRow( + key_row, cached_key_row, staged_key[tile_offset % kTile], lane_id); + } + __syncthreads(); + for (int tile_offset = chunk_begin; tile_offset < chunk_end; ++tile_offset) { + if (!valid[tile_offset]) { + if (lane_id == 0) { + logits[warp_id][tile_offset] = -FLT_MAX; + } + continue; + } + float dot = 0.0f; +#pragma unroll + for (int channel = 0; channel < kChannels / 32; ++channel) { + dot += query_values[channel] * staged_key[tile_offset % kTile][lane_id + channel * 32]; + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + dot += __shfl_xor_sync(0xFFFFFFFFU, dot, offset); + } + if (lane_id == 0) { + dot *= scale; + logits[warp_id][tile_offset] = softcap > 0.0f ? tanhf(dot / softcap) * softcap : dot; + } + } + __syncthreads(); + } + float tile_max = -FLT_MAX; + for (int tile_offset = lane_id; tile_offset < tile_length; tile_offset += 32) { + tile_max = fmaxf(tile_max, logits[warp_id][tile_offset]); + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + tile_max = fmaxf(tile_max, __shfl_xor_sync(0xFFFFFFFFU, tile_max, offset)); + } + const bool has_values = tile_max != -FLT_MAX; + const float new_max = has_values ? fmaxf(running_max, tile_max) : running_max; + const float old_weight = has_values ? __expf(running_max - new_max) : 1.0f; + float weights[kScoreTile / 32]; +#pragma unroll + for (int part = 0; part < kScoreTile / 32; ++part) { + const int tile_offset = lane_id + part * 32; + weights[part] = has_values && tile_offset < tile_length + ? __expf(logits[warp_id][tile_offset] - new_max) + : 0.0f; + logits[warp_id][tile_offset] = weights[part]; + } + float tile_sum = (weights[0] + weights[2]) + (weights[1] + weights[3]); +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + tile_sum += __shfl_down_sync(0xFFFFFFFFU, tile_sum, offset); + } + running_sum = running_sum * old_weight + __shfl_sync(0xFFFFFFFFU, tile_sum, 0); + running_max = new_max; +#pragma unroll + for (int channel = 0; channel < kChannels / 32; ++channel) { + accumulator[channel] *= old_weight; + } + for (int chunk_begin = 0; chunk_begin < tile_length; chunk_begin += kTile) { + const int chunk_end = min(tile_length, chunk_begin + kTile); + for (int tile_offset = chunk_begin + warp_id; tile_offset < chunk_end; tile_offset += kHeads) { + SparsePagedAttentionStageRow( + value_rows[tile_offset], cached_value_rows[tile_offset], + staged_value[tile_offset % kTile], lane_id); + } + __syncthreads(); + for (int tile_offset = chunk_begin; tile_offset < chunk_end; ++tile_offset) { + if (!valid[tile_offset]) { + continue; + } +#pragma unroll + for (int channel = 0; channel < kChannels / 32; ++channel) { + accumulator[channel] += logits[warp_id][tile_offset] * + staged_value[tile_offset % kTile][lane_id + channel * 32]; + } + } + __syncthreads(); + } + __syncthreads(); + } + const int64_t partial_index = (static_cast(split_id) * token_count + token_id) * num_heads + head_id; + float inverse_sum = 1.0f; + if (num_splits == 1) { + const float final_max = head_sink == nullptr ? running_max + : fmaxf(running_max, SparseActivationValue(head_sink[head_id])); + const float numerator_scale = __expf(running_max - final_max); + const float final_sum = running_sum * numerator_scale + + (head_sink == nullptr ? 0.0f + : __expf(SparseActivationValue(head_sink[head_id]) - final_max)); + inverse_sum = final_sum > 0.0f ? numerator_scale / final_sum : 0.0f; + } +#pragma unroll + for (int channel = 0; channel < kChannels / 32; ++channel) { + const int64_t base = num_splits == 1 ? query_base : partial_index * kChannels; + if (num_splits == 1) { + output[base + lane_id + channel * 32] = static_cast(accumulator[channel] * inverse_sum); + } else { + partial_out[base + lane_id + channel * 32] = accumulator[channel]; + } + } + if (num_splits > 1 && lane_id == 0) { + partial_max[partial_index] = running_max; + partial_sum[partial_index] = running_sum; + } +} + +constexpr int kSparsePagedAttentionGqaHeads = 4; +constexpr int kSparsePagedAttentionGqaTile = 64; +constexpr int kSparsePagedAttentionGqaHeadSize = 256; + +template +__device__ __forceinline__ void SparsePagedAttentionStageKv( + T* tile, const T* current, const T* cache, const int64_t* candidate_refs, + int tile_length, int load_rows, int kv_head, int kv_heads, int current_stride) { + static_assert(sizeof(T) == 2 && SharedStride % 8 == 0); + for (int i = threadIdx.x; i < load_rows * 32; i += blockDim.x) { + const int row = i / 32; + const int c = (i % 32) * 8; + const int64_t ref = row < tile_length ? candidate_refs[row] : -1; + const T* source = nullptr; + if (ref >= kSparsePagedAttentionDirectCandidate) { + source = current + (ref - kSparsePagedAttentionDirectCandidate) * current_stride + + kv_head * kSparsePagedAttentionGqaHeadSize + c; + } else if (ref >= 0) { + source = cache + (ref * kv_heads + kv_head) * kSparsePagedAttentionGqaHeadSize + c; + } + T* destination = tile + row * SharedStride + c; + if (source == nullptr) { + *reinterpret_cast(destination) = uint4{}; + } else if (reinterpret_cast(source) % sizeof(uint4) == 0) { + *reinterpret_cast(destination) = *reinterpret_cast(source); + } else { +#pragma unroll + for (int j = 0; j < 8; ++j) { + destination[j] = source[j]; + } + } + } +} + +template +__global__ void SparsePagedAttentionGqaTiledKernel( + const T* query, const T* current_key, const T* current_value, + const T* key_cache, const T* value_cache, const int* cumulative_seqlens_q, + const int* past_seqlens, const int* block_table, const int* slot_mapping, + const int* selected_indices, const int* selected_counts, T* output, const T* head_sink, + float* partial_out, float* partial_max, float* partial_sum, + int batch_size, int token_count, int num_heads, int kv_num_heads, int block_size, + int num_blocks, int max_num_blocks_per_seq, int max_selected_entries, + int current_key_stride, int current_value_stride, float scale, float softcap, bool is_causal, int num_splits) { + constexpr int head_size = kSparsePagedAttentionGqaHeadSize; + constexpr int tile_size = kSparsePagedAttentionGqaTile; + constexpr int channels_per_lane = head_size / 32; + constexpr int heads_per_warp = TensorCorePv ? 3 : 1; + constexpr int head_stride = TensorCorePv ? 4 : 1; + constexpr int shared_head_stride = head_size + (TensorCoreQk ? 8 : 0); + constexpr int shared_logit_stride = tile_size + (TensorCoreQk ? 4 : 0); + __shared__ __align__(32) T kv_tile[tile_size * shared_head_stride]; + __shared__ int64_t candidate_refs[tile_size]; + __shared__ __align__(32) float logits[TensorCoreQk ? 16 : HeadsPerBlock][shared_logit_stride]; + __shared__ __align__(32) T query_tile[TensorCoreQk ? 16 * shared_head_stride : 1]; + T* probabilities = query_tile; + __shared__ float rescale[TensorCorePv ? 16 : 1]; + const int tid = threadIdx.x; + const int warp = tid / 32; + const int lane = tid % 32; +#if __CUDA_ARCH__ >= 800 + namespace wmma = nvcuda::wmma; + using MmaT = std::conditional_t::value, half, __nv_bfloat16>; +#endif + float pv[TensorCorePv ? 8 : 1][4] = {}; + static_assert(!TensorCorePv || (TensorCoreQk && std::is_same::value)); + const int head_group = blockIdx.x % (num_heads / HeadsPerBlock); + const int head_id = head_group * HeadsPerBlock + warp; + const int token_id = blockIdx.x / (num_heads / HeadsPerBlock); + const int kv_head_id = head_id / (num_heads / kv_num_heads); + int batch_id = 0; + while (batch_id + 1 < batch_size && token_id >= cumulative_seqlens_q[batch_id + 1]) { + ++batch_id; + } + const int query_index = token_id - cumulative_seqlens_q[batch_id]; + const int query_length = cumulative_seqlens_q[batch_id + 1] - cumulative_seqlens_q[batch_id]; + const int query_position = past_seqlens[batch_id] + query_index; + const int main_length = past_seqlens[batch_id] + query_length; + const int selected_count = min(max(selected_counts[token_id], 0), max_selected_entries); + const int candidates_per_split = (selected_count + num_splits - 1) / num_splits; + const int candidate_begin = blockIdx.y * candidates_per_split; + const int candidate_end = min(selected_count, candidate_begin + candidates_per_split); + const int64_t partial_head_index = + (static_cast(blockIdx.y) * token_count + token_id) * num_heads + head_id; + const int64_t query_base = (static_cast(token_id) * num_heads + head_id) * head_size; + float query_values[channels_per_lane]; + float accumulator[channels_per_lane] = {}; +#pragma unroll + for (int i = 0; i < channels_per_lane; ++i) { + query_values[i] = SparseActivationValue(query[query_base + lane + i * 32]); + } + float running_max[heads_per_warp]; + float running_sum[heads_per_warp] = {}; +#pragma unroll + for (int i = 0; i < heads_per_warp; ++i) { + running_max[i] = -FLT_MAX; + } + + for (int tile_begin = candidate_begin; tile_begin < candidate_end; tile_begin += tile_size) { + const int tile_length = min(tile_size, candidate_end - tile_begin); + if constexpr (TensorCoreQk) { + static_assert(HeadsPerBlock == 12); + for (int i = tid; i < 16 * head_size; i += blockDim.x) { + const int h = i / head_size; + const int c = i % head_size; + query_tile[h * shared_head_stride + c] = + h < HeadsPerBlock + ? query[(static_cast(token_id) * num_heads + head_group * HeadsPerBlock + h) * head_size + c] + : static_cast(0.0f); + } + } + if (tid < tile_length) { + const int logical_position = + selected_indices[static_cast(token_id) * max_selected_entries + tile_begin + tid]; + candidate_refs[tid] = SparsePagedAttentionResolveMainCandidate( + logical_position, batch_id, query_position, main_length, past_seqlens[batch_id], + cumulative_seqlens_q, block_table, slot_mapping, block_size, num_blocks, + max_num_blocks_per_seq, is_causal); + } + __syncthreads(); + // Each K/V row is fetched once for query heads sharing the same KV head. + if constexpr (HeadsPerBlock == 12) { + SparsePagedAttentionStageKv( + kv_tile, current_key, key_cache, candidate_refs, tile_length, TensorCoreQk ? tile_size : tile_length, + kv_head_id, kv_num_heads, current_key_stride); + } else { + for (int i = tid; i < (TensorCoreQk ? tile_size : tile_length) * head_size; i += blockDim.x) { + const int row = i / head_size; + const int c = i % head_size; + const int64_t ref = row < tile_length ? candidate_refs[row] : -1; + T value = static_cast(0.0f); + if (ref >= kSparsePagedAttentionDirectCandidate) { + value = current_key[(ref - kSparsePagedAttentionDirectCandidate) * current_key_stride + + kv_head_id * head_size + c]; + } else if (ref >= 0) { + value = key_cache[(ref * kv_num_heads + kv_head_id) * head_size + c]; + } + kv_tile[row * shared_head_stride + c] = value; + } + } + __syncthreads(); +#if __CUDA_ARCH__ >= 800 + if constexpr (TensorCoreQk) { + namespace wmma = nvcuda::wmma; + using MmaT = std::conditional_t::value, half, __nv_bfloat16>; + static_assert(sizeof(MmaT) == sizeof(T)); + if (warp < 4) { + wmma::fragment dots; + wmma::fill_fragment(dots, 0.0f); + for (int k = 0; k < head_size; k += 16) { + wmma::fragment q; + wmma::fragment key; + wmma::load_matrix_sync(q, reinterpret_cast(query_tile + k), shared_head_stride); + wmma::load_matrix_sync(key, reinterpret_cast(kv_tile + warp * 16 * shared_head_stride + k), shared_head_stride); + wmma::mma_sync(dots, q, key, dots); + } + wmma::store_matrix_sync(&logits[0][warp * 16], dots, shared_logit_stride, wmma::mem_row_major); + } + __syncthreads(); + for (int slot = 0; slot < heads_per_warp; ++slot) { + const int h = warp + slot * head_stride; + for (int row = lane; row < tile_length; row += 32) { + const float scaled_dot = logits[h][row] * scale; + logits[h][row] = candidate_refs[row] == -1 + ? -FLT_MAX + : (softcap > 0.0f ? tanhf(scaled_dot / softcap) * softcap : scaled_dot); + } + } + } else +#endif + { + // Interleave independent rows without changing each dot product's reduction order. + for (int row_begin = 0; row_begin < tile_length; row_begin += DotRows) { + float dots[DotRows] = {}; +#pragma unroll + for (int i = 0; i < channels_per_lane; ++i) { +#pragma unroll + for (int r = 0; r < DotRows; ++r) { + if (row_begin + r < tile_length) { + dots[r] += query_values[i] * SparseActivationValue(kv_tile[(row_begin + r) * shared_head_stride + lane + i * 32]); + } + } + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { +#pragma unroll + for (int r = 0; r < DotRows; ++r) { + dots[r] += __shfl_xor_sync(0xFFFFFFFFU, dots[r], offset); + } + } + if (lane == 0) { +#pragma unroll + for (int r = 0; r < DotRows; ++r) { + const int row = row_begin + r; + if (row < tile_length) { + const float scaled_dot = dots[r] * scale; + logits[warp][row] = candidate_refs[row] == -1 + ? -FLT_MAX + : (softcap > 0.0f ? tanhf(scaled_dot / softcap) * softcap : scaled_dot); + } + } + } + } + } + __syncwarp(); + float old_weight[heads_per_warp]; + for (int slot = 0; slot < heads_per_warp; ++slot) { + const int h = warp + slot * head_stride; + float tile_max = -FLT_MAX; + for (int row = lane; row < tile_length; row += 32) { + tile_max = fmaxf(tile_max, logits[h][row]); + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + tile_max = fmaxf(tile_max, __shfl_xor_sync(0xFFFFFFFFU, tile_max, offset)); + } + const float new_max = fmaxf(running_max[slot], tile_max); + old_weight[slot] = __expf(running_max[slot] - new_max); + float tile_sum = 0.0f; + for (int row = lane; row < tile_length; row += 32) { + const float weight = candidate_refs[row] == -1 ? 0.0f : __expf(logits[h][row] - new_max); + logits[h][row] = weight; + tile_sum += weight; + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + tile_sum += __shfl_xor_sync(0xFFFFFFFFU, tile_sum, offset); + } + running_sum[slot] = running_sum[slot] * old_weight[slot] + tile_sum; + running_max[slot] = new_max; + if constexpr (TensorCorePv) { + if (lane == 0) { + rescale[h] = old_weight[slot]; + } + } + } + __syncthreads(); + if constexpr (TensorCorePv) { + if (tid >= HeadsPerBlock && tid < 16) { + rescale[tid] = 1.0f; + } + for (int i = tid; i < 16 * tile_size; i += blockDim.x) { + const int h = i / tile_size; + const int row = i % tile_size; + probabilities[i] = h < HeadsPerBlock && row < tile_length + ? static_cast(logits[h][row]) + : static_cast(0.0f); + } + } + // Reuse the key tile's storage after all warps have finished their dot products. + if constexpr (HeadsPerBlock == 12) { + SparsePagedAttentionStageKv( + kv_tile, current_value, value_cache, candidate_refs, tile_length, TensorCorePv ? tile_size : tile_length, + kv_head_id, kv_num_heads, current_value_stride); + } else { + for (int i = tid; i < (TensorCorePv ? tile_size : tile_length) * head_size; i += blockDim.x) { + const int row = i / head_size; + const int c = i % head_size; + const int64_t ref = row < tile_length ? candidate_refs[row] : -1; + T value = static_cast(0.0f); + if (ref >= kSparsePagedAttentionDirectCandidate) { + value = current_value[(ref - kSparsePagedAttentionDirectCandidate) * current_value_stride + + kv_head_id * head_size + c]; + } else if (ref >= 0) { + value = value_cache[(ref * kv_num_heads + kv_head_id) * head_size + c]; + } + kv_tile[row * shared_head_stride + c] = value; + } + } + __syncthreads(); +#if __CUDA_ARCH__ >= 800 + if constexpr (TensorCorePv) { + const int h = lane / 4; +#pragma unroll + for (int n = 0; n < 8; ++n) { + pv[n][0] *= rescale[h]; + pv[n][1] *= rescale[h]; + pv[n][2] *= rescale[h + 8]; + pv[n][3] *= rescale[h + 8]; + } + for (int k = 0; k < tile_size; k += 16) { + uint32_t p[4]; + gated_delta_net::LoadFragA(p, probabilities, tile_size, 0, k, lane); +#pragma unroll + for (int n = 0; n < 8; ++n) { + uint32_t value[2]; + gated_delta_net::LoadFragB(value, kv_tile, shared_head_stride, k, warp * 64 + n * 8, lane); + gated_delta_net::MmaM16N8K16(pv[n], p, value); + } + } + } else +#endif + { +#pragma unroll + for (int i = 0; i < channels_per_lane; ++i) { + accumulator[i] *= old_weight[0]; + } + for (int row = 0; row < tile_length; ++row) { + const float weight = logits[warp][row]; +#pragma unroll + for (int i = 0; i < channels_per_lane; ++i) { + accumulator[i] += weight * SparseActivationValue(kv_tile[row * shared_head_stride + lane + i * 32]); + } + } + } + __syncthreads(); + } + + float inverse_sum[heads_per_warp]; + for (int slot = 0; slot < heads_per_warp; ++slot) { + inverse_sum[slot] = 1.0f; + const int h = head_id + slot * head_stride; + if (num_splits == 1) { + const float final_max = head_sink == nullptr + ? running_max[slot] + : fmaxf(running_max[slot], SparseActivationValue(head_sink[h])); + const float final_sum = running_sum[slot] * __expf(running_max[slot] - final_max) + + (head_sink == nullptr ? 0.0f : __expf(SparseActivationValue(head_sink[h]) - final_max)); + inverse_sum[slot] = final_sum > 0.0f ? __expf(running_max[slot] - final_max) / final_sum : 0.0f; + } + } +#if __CUDA_ARCH__ >= 800 + if constexpr (TensorCorePv) { + if (lane == 0) { + for (int slot = 0; slot < heads_per_warp; ++slot) { + rescale[warp + slot * head_stride] = inverse_sum[slot]; + } + } + __syncthreads(); + const int h = lane / 4; +#pragma unroll + for (int n = 0; n < 8; ++n) { + const int c = warp * 64 + n * 8 + (lane % 4) * 2; +#pragma unroll + for (int row = 0; row < 2; ++row) { + const int head = h + row * 8; + if (head < HeadsPerBlock) { + const int64_t base = (static_cast(token_id) * num_heads + head_group * HeadsPerBlock + head) * head_size; + output[base + c] = static_cast(pv[n][row * 2] * rescale[head]); + output[base + c + 1] = static_cast(pv[n][row * 2 + 1] * rescale[head]); + } + } + } + } else +#endif + { +#pragma unroll + for (int i = 0; i < channels_per_lane; ++i) { + if (num_splits == 1) { + output[query_base + lane + i * 32] = static_cast(accumulator[i] * inverse_sum[0]); + } else { + partial_out[partial_head_index * head_size + lane + i * 32] = accumulator[i]; + } + } + if (num_splits > 1 && lane == 0) { + partial_max[partial_head_index] = running_max[0]; + partial_sum[partial_head_index] = running_sum[0]; + } + } +} + template __global__ void SparsePagedAttentionReduceKernel( T* output, const float* partial_out, const float* partial_max, @@ -430,15 +1052,37 @@ __global__ void SparsePagedAttentionReduceKernel( } } +int SparsePagedAttentionHeadsPerBlock( + const contrib::PagedAttentionParameters& parameters, + SparseAttentionMode attention_mode, SelectedKvSource selected_kv_source, + bool matching_cache_type, size_t shared_memory_per_block) { + constexpr size_t tiled_shared_bytes = + kSparsePagedAttentionGqaTile * kSparsePagedAttentionGqaHeadSize * sizeof(half) + + kSparsePagedAttentionGqaTile * sizeof(int64_t) + + kSparsePagedAttentionGqaHeads * kSparsePagedAttentionGqaTile * sizeof(float); + return matching_cache_type && parameters.head_size == kSparsePagedAttentionGqaHeadSize && + (parameters.num_heads / parameters.kv_num_heads) % kSparsePagedAttentionGqaHeads == 0 && + attention_mode == SparseAttentionMode::kSelectedOnly && selected_kv_source == SelectedKvSource::kMain && + tiled_shared_bytes <= shared_memory_per_block + ? kSparsePagedAttentionGqaHeads + : 1; +} + int ComputeSparsePagedAttentionSplits( const int token_count, const int num_heads, const int max_candidate_count, - const int multi_processor_count) { - const int64_t base_blocks = static_cast(token_count) * num_heads; + const int multi_processor_count, const int heads_per_block) { + const int64_t base_blocks = static_cast(token_count) * (num_heads / heads_per_block); if (base_blocks <= 0 || base_blocks >= 2 * multi_processor_count) { return 1; } - const int by_occupancy = static_cast((2 * multi_processor_count + base_blocks - 1) / base_blocks); - const int by_length = (max_candidate_count + kSparsePagedAttentionTile - 1) / kSparsePagedAttentionTile; + const int target_waves = + heads_per_block == kSparsePagedAttentionGqaHeads && token_count <= 8 ? 4 : 2; + const int by_occupancy = + static_cast((target_waves * multi_processor_count + base_blocks - 1) / base_blocks); + const int tile_size = heads_per_block == kSparsePagedAttentionGqaHeads + ? kSparsePagedAttentionGqaTile + : kSparsePagedAttentionTile; + const int by_length = (max_candidate_count + tile_size - 1) / tile_size; return max(1, min(min(by_occupancy, by_length), kSparsePagedAttentionMaxSplits)); } @@ -458,7 +1102,8 @@ Status SparseQkvToContext( const T* auxiliary_key, const T* auxiliary_value, const int* auxiliary_lengths, int auxiliary_capacity, int auxiliary_num_heads, SparseAttentionMode attention_mode, SelectedKvSource selected_kv_source, bool auxiliary_kv_shared, - float* partial_out, float* partial_max, float* partial_sum, int num_splits) { + float* partial_out, float* partial_max, float* partial_sum, int num_splits, + bool enable_grouped, bool enable_grouped_vectorized) { const size_t shared_memory_bytes = GetSparsePagedAttentionSharedMemoryBytes(parameters.head_size); if (shared_memory_bytes > static_cast(device_prop.sharedMemPerBlock)) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, @@ -487,8 +1132,113 @@ Status SparseQkvToContext( const float attention_scale = parameters.scale == 0.0f ? 1.0f / sqrtf(static_cast(parameters.head_size)) : parameters.scale; - SparsePagedAttentionSplitKernel<<(stream->GetHandle())>>>( + const bool use_grouped = std::is_same::value && std::is_same::value && + parameters.head_size == 256 && parameters.num_heads / parameters.kv_num_heads == 12 && + parameters.token_count >= 32 && num_splits == 1 && enable_grouped && + device_prop.sharedMemPerBlock >= 2 * 8 * 256 * sizeof(float) + + 6 * kSparsePagedAttentionTile * sizeof(float) + + kSparsePagedAttentionTile * (sizeof(const T*) + + sizeof(const TCACHE*) + sizeof(bool)); + if (use_grouped) { + LOGS_DEFAULT(VERBOSE) << "SparsePagedAttention dispatch: grouped, vectorized=" << enable_grouped_vectorized; + const dim3 grouped_grid(static_cast(token_head_count / 6), num_splits); + const auto launch_grouped = [&](auto vectorized_tag) { + SparsePagedAttentionGroupedKernel<<(stream->GetHandle())>>>( + prepared_query, prepared_key, prepared_value, data.key_cache, data.value_cache, + data.cumulative_seqlens_q, data.past_seqlens, data.block_table, data.slot_mapping, + selected_indices, selected_counts, auxiliary_key, auxiliary_value, auxiliary_lengths, + partial_out, partial_max, partial_sum, data.output, data.head_sink, + parameters.batch_size, parameters.token_count, parameters.num_heads, parameters.kv_num_heads, + parameters.block_size, parameters.num_blocks, parameters.max_num_blocks_per_seq, + max_selected_entries, auxiliary_capacity, auxiliary_num_heads, prepared_key_stride, prepared_value_stride, + attention_scale, parameters.softcap, parameters.local_window_size, parameters.is_causal, + attention_mode == SparseAttentionMode::kLocalPlusSelected, + selected_kv_source == SelectedKvSource::kAuxiliary, auxiliary_kv_shared, num_splits); + }; + if (enable_grouped_vectorized) { + launch_grouped(std::true_type{}); + } else { + launch_grouped(std::false_type{}); + } + + return CUDA_CALL(cudaGetLastError()); + } + if constexpr (std::is_same::value) { + if (SparsePagedAttentionHeadsPerBlock(parameters, attention_mode, selected_kv_source, true, + device_prop.sharedMemPerBlock) == kSparsePagedAttentionGqaHeads) { + constexpr size_t twelve_head_shared_bytes = + kSparsePagedAttentionGqaTile * kSparsePagedAttentionGqaHeadSize * sizeof(T) + + kSparsePagedAttentionGqaTile * sizeof(int64_t) + + 12 * kSparsePagedAttentionGqaTile * sizeof(float); + constexpr size_t six_head_shared_bytes = + kSparsePagedAttentionGqaTile * kSparsePagedAttentionGqaHeadSize * sizeof(T) + + kSparsePagedAttentionGqaTile * sizeof(int64_t) + + 6 * kSparsePagedAttentionGqaTile * sizeof(float); + const bool use_twelve_heads = parameters.token_count > 64 && num_splits == 1 && + (parameters.num_heads / parameters.kv_num_heads) % 12 == 0 && + twelve_head_shared_bytes <= device_prop.sharedMemPerBlock; + const bool use_six_heads = parameters.token_count > 64 && num_splits == 1 && + (parameters.num_heads / parameters.kv_num_heads) % 6 == 0 && + six_head_shared_bytes <= device_prop.sharedMemPerBlock; + const int heads_per_block = use_twelve_heads ? 12 : (use_six_heads ? 6 : kSparsePagedAttentionGqaHeads); + const dim3 tiled_grid(static_cast(token_head_count / heads_per_block), num_splits); + static const bool enable_tensor_core_qk = + ParseEnvironmentVariableWithDefault("ORT_SPARSE_PREFILL_TENSOR_CORE_QK", false); + static const bool enable_tensor_core_pv = + ParseEnvironmentVariableWithDefault("ORT_SPARSE_PREFILL_TENSOR_CORE_PV", false); + const auto tensor_core_pv_kernel = [] { + if constexpr (std::is_same::value) { + return SparsePagedAttentionGqaTiledKernel; + } else { + return SparsePagedAttentionGqaTiledKernel; + } + }(); + bool use_tensor_core_qk = false; + bool use_tensor_core_pv = false; + if (enable_tensor_core_qk && use_twelve_heads && device_prop.major >= 8) { + cudaFuncAttributes attributes{}; + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaFuncGetAttributes(&attributes, SparsePagedAttentionGqaTiledKernel))); + use_tensor_core_qk = attributes.ptxVersion >= 80 && + attributes.sharedSizeBytes <= device_prop.sharedMemPerBlock; + if (use_tensor_core_qk && enable_tensor_core_pv && std::is_same::value) { + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaFuncGetAttributes(&attributes, tensor_core_pv_kernel))); + use_tensor_core_pv = attributes.sharedSizeBytes <= device_prop.sharedMemPerBlock; + } + } + const auto tiled_kernel = use_twelve_heads + ? (use_tensor_core_pv ? tensor_core_pv_kernel + : use_tensor_core_qk ? SparsePagedAttentionGqaTiledKernel + : SparsePagedAttentionGqaTiledKernel) + : (use_six_heads ? SparsePagedAttentionGqaTiledKernel + : SparsePagedAttentionGqaTiledKernel); + LOGS_DEFAULT(VERBOSE) << "SparsePagedAttention dispatch: tiled, heads=" << heads_per_block + << ", qk=" << use_tensor_core_qk << ", pv=" << use_tensor_core_pv; + tiled_kernel<<(stream->GetHandle())>>>( + prepared_query, prepared_key, prepared_value, data.key_cache, data.value_cache, + data.cumulative_seqlens_q, data.past_seqlens, data.block_table, data.slot_mapping, + selected_indices, selected_counts, data.output, data.head_sink, + partial_out, partial_max, partial_sum, + parameters.batch_size, parameters.token_count, parameters.num_heads, parameters.kv_num_heads, parameters.block_size, + parameters.num_blocks, parameters.max_num_blocks_per_seq, max_selected_entries, + prepared_key_stride, prepared_value_stride, attention_scale, parameters.softcap, parameters.is_causal, num_splits); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + if (num_splits > 1) { + const dim3 reduce_grid(static_cast(token_head_count)); + SparsePagedAttentionReduceKernel<<(stream->GetHandle())>>>( + data.output, partial_out, partial_max, partial_sum, data.head_sink, + parameters.token_count, parameters.num_heads, parameters.head_size, num_splits); + } + return CUDA_CALL(cudaGetLastError()); + } + } + const auto kernel = parameters.head_size == 256 + ? SparsePagedAttentionSplitKernel + : SparsePagedAttentionSplitKernel; + kernel<<(stream->GetHandle())>>>( prepared_query, prepared_key, prepared_value, data.key_cache, data.value_cache, data.k_scale, data.v_scale, data.cumulative_seqlens_q, data.past_seqlens, data.block_table, data.slot_mapping, selected_indices, selected_counts, auxiliary_key, auxiliary_value, auxiliary_lengths, partial_out, @@ -519,7 +1269,7 @@ Status SparseQkvToContext( const cudaDeviceProp&, Stream*, contrib::PagedAttentionParameters&, \ PagedAttentionData&, const int*, const int*, int, const T*, \ const T*, const int*, int, int, SparseAttentionMode, SelectedKvSource, \ - bool, float*, float*, float*, int); + bool, float*, float*, float*, int, bool, bool); INSTANTIATE_SPARSE_PAGED_ATTENTION(half, half) INSTANTIATE_SPARSE_PAGED_ATTENTION(BFloat16, BFloat16) diff --git a/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention_impl.h index e9ae012d82c30..49a563e831bd0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention_impl.h @@ -22,11 +22,17 @@ enum class SelectedKvSource { kAuxiliary, }; +int SparsePagedAttentionHeadsPerBlock( + const contrib::PagedAttentionParameters& parameters, + SparseAttentionMode attention_mode, SelectedKvSource selected_kv_source, + bool matching_cache_type, size_t shared_memory_per_block); + int ComputeSparsePagedAttentionSplits( int token_count, int num_heads, int max_candidate_count, - int multi_processor_count); + int multi_processor_count, + int heads_per_block = 1); template Status SparseQkvToContext( @@ -48,7 +54,7 @@ Status SparseQkvToContext( float* partial_out, float* partial_max, float* partial_sum, - int num_splits); + int num_splits, bool enable_grouped = false, bool enable_grouped_vectorized = false); } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_int2_prefill.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_int2_prefill.cu index e9b8fe1f7e50c..bdfe80aba6184 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_int2_prefill.cu +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_int2_prefill.cu @@ -8,11 +8,289 @@ #include "contrib_ops/cuda/llm/moe_gemm/moe_kernels.h" #include "contrib_ops/cuda/llm/moe_gemm/moe_util_kernels.h" #include "contrib_ops/cuda/llm/moe_gemm/moe_gemm_activation_kernels.cuh" +#include "contrib_ops/cuda/llm/fpA_intB_gemv/details.h" #include "core/common/safeint.h" +#include "core/providers/cuda/cuda_common.h" + +#include +#include +#include +#include +#include +#include "cutlass/float8.h" namespace onnxruntime::llm::kernels::cutlass_kernels { namespace { +constexpr int kNvfp4TileK = 64; +constexpr int kNvfp4TileM = 64; +constexpr int kNvfp4TileN = 128; +constexpr int kNvfp4DownColumns = 1024; +constexpr int kNvfp4TileStride = kNvfp4TileK + 8; +constexpr int kNvfp4AElements = kNvfp4TileM * kNvfp4TileStride; +constexpr int kNvfp4PackedTileBytes = kNvfp4TileN * kNvfp4TileK / 2; +constexpr int kNvfp4PackedOffset = (2 * kNvfp4TileM + kNvfp4TileN) * kNvfp4TileStride * sizeof(half); +constexpr int kNvfp4SharedBytes = kNvfp4PackedOffset + 2 * kNvfp4PackedTileBytes; +static_assert(kNvfp4SharedBytes >= kNvfp4TileM * kNvfp4TileN * sizeof(float)); + +template +__device__ void LoadNvfp4Stage(T* a, uint8_t* packed, const T* input, const uint8_t* weights, + int64_t first_row, int64_t end_row, int expert, + int first_column, int columns, int reduction, int k, + const int* source_rows) { +#if __CUDA_ARCH__ >= 800 + constexpr int vectors = kNvfp4TileK / 8; + for (int i = threadIdx.x; i < kNvfp4TileM * vectors; i += blockDim.x) { + const int64_t row = first_row + i / vectors; + const int kk = (i % vectors) * 8; + if (row < end_row) { + const int64_t source_row = GatherInput ? source_rows[i / vectors] : row; + __pipeline_memcpy_async(a + (i / vectors) * kNvfp4TileStride + kk, + input + source_row * reduction + k + kk, sizeof(uint4)); + } else { + *reinterpret_cast(a + (i / vectors) * kNvfp4TileStride + kk) = uint4{}; + } + } + for (int i = threadIdx.x; i < kNvfp4TileN * vectors; i += blockDim.x) { + const int col = first_column + i / vectors; + const int kk = k + (i % vectors) * 8; + if (col < columns) { + const int64_t index = (static_cast(expert) * columns + col) * (reduction / 2) + kk / 2; + __pipeline_memcpy_async(packed + i * sizeof(uint32_t), weights + index, sizeof(uint32_t)); + } else { + reinterpret_cast(packed)[i] = 0; + } + } + __pipeline_commit(); +#endif +} + +__global__ void BuildNvfp4TileOffsets(const int64_t* offsets, int* tiles, int experts) { + if (threadIdx.x == 0) { + tiles[0] = 0; + for (int expert = 0; expert < experts; ++expert) { + tiles[expert + 1] = tiles[expert] + + static_cast((offsets[expert + 1] - offsets[expert] + kNvfp4TileM - 1) / kNvfp4TileM); + } + } +} + +template +__global__ void Nvfp4GroupedGemm(const T* input, const uint8_t* weights, + const uint8_t* scales, const float* global_scales, + const int64_t* offsets, const int* tiles, T* output, + int columns, int reduction, int experts, + const T* bias, float alpha, float beta, float limit, + int column_offset, int output_columns, + const int* permuted_rows, int64_t num_rows) { +#if __CUDA_ARCH__ >= 800 + using namespace nvcuda; + constexpr int stride = kNvfp4TileStride; + // Input tiles share this storage with the FP32 epilogue, after the final K barrier. + extern __shared__ __align__(32) uint8_t storage[]; + float* c = reinterpret_cast(storage); + __shared__ int expert, row_tile; + __shared__ int source_rows[GatherInput ? kNvfp4TileM : 1]; + const int warp = threadIdx.x / 32; + const int m = (warp / 2) * 16; + const int n = (warp % 2) * (kNvfp4TileN / 2); + const int n_tiles = (output_columns - 1) / kNvfp4TileN + 1; + const int64_t total = static_cast(tiles[experts]) * n_tiles; + for (int64_t tile = blockIdx.x; tile < total; tile += gridDim.x) { + if (threadIdx.x == 0) { + const int target = static_cast(tile / n_tiles); + int lo = 0, hi = experts; + while (lo + 1 < hi) { + const int mid = (lo + hi) / 2; + if (tiles[mid] <= target) + lo = mid; + else + hi = mid; + } + expert = lo; + row_tile = target - tiles[lo]; + } + __syncthreads(); + const int64_t first_row = offsets[expert] + static_cast(row_tile) * kNvfp4TileM; + const int64_t end_row = offsets[expert + 1]; + const int first_column = column_offset + static_cast(tile % n_tiles) * kNvfp4TileN; + const float global_scale = global_scales[expert]; + if constexpr (GatherInput) { + if (threadIdx.x < kNvfp4TileM) { + const int64_t row = first_row + threadIdx.x; + source_rows[threadIdx.x] = row < end_row ? permuted_rows[row] % num_rows : 0; + } + __syncthreads(); + } + constexpr int fragments = kNvfp4TileN / 32; + wmma::fragment acc[fragments]; + for (int j = 0; j < fragments; ++j) wmma::fill_fragment(acc[j], 0.0f); + LoadNvfp4Stage(reinterpret_cast(storage), storage + kNvfp4PackedOffset, + input, weights, first_row, end_row, expert, first_column, columns, reduction, 0, source_rows); + for (int k = 0; k < reduction; k += kNvfp4TileK) { + const int buffer = (k / kNvfp4TileK) % 2; + T* a = reinterpret_cast(storage) + buffer * kNvfp4AElements; + T* b = reinterpret_cast(storage) + 2 * kNvfp4AElements; + const uint32_t* packed_tile = reinterpret_cast( + storage + kNvfp4PackedOffset + buffer * kNvfp4PackedTileBytes); + __pipeline_wait_prior(0); + __syncthreads(); + // Only packed weights and input rows need double buffering; the dequantized B tile is reused. + if (k + kNvfp4TileK < reduction) { + T* next_a = reinterpret_cast(storage) + (1 - buffer) * kNvfp4AElements; + uint8_t* next_packed = storage + kNvfp4PackedOffset + (1 - buffer) * kNvfp4PackedTileBytes; + LoadNvfp4Stage(next_a, next_packed, + input, weights, first_row, end_row, expert, first_column, columns, reduction, k + kNvfp4TileK, + source_rows); + } + constexpr int vectors = kNvfp4TileK / 8; + for (int i = threadIdx.x; i < kNvfp4TileN * vectors; i += blockDim.x) { + const int col = first_column + i / vectors; + const int kk = k + (i % vectors) * 8; + alignas(16) T values[8]{}; + if (col >= columns) { + *reinterpret_cast(b + (i / vectors) * stride + (i % vectors) * 8) = + *reinterpret_cast(values); + continue; + } + uint32_t packed = packed_tile[i]; + fpA_intB_gemv::Fp4I2FConverter::template convert<8>(&packed, values); + const int64_t scale_index = (static_cast(expert) * columns + col) * (reduction / 16) + kk / 16; + const float scale = static_cast(reinterpret_cast(scales)[scale_index]) * + global_scale; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float left = static_cast(values[j * 2]) * scale; + const float right = static_cast(values[j * 2 + 1]) * scale; + if constexpr (std::is_same_v) { + reinterpret_cast(values)[j] = __floats2half2_rn(left, right); + } else { + reinterpret_cast<__nv_bfloat162*>(values)[j] = __floats2bfloat162_rn(left, right); + } + } + *reinterpret_cast(b + (i / vectors) * stride + (i % vectors) * 8) = + *reinterpret_cast(values); + } + __syncthreads(); + for (int step = 0; step < kNvfp4TileK; step += 16) { + wmma::fragment af; + wmma::load_matrix_sync(af, a + m * stride + step, stride); + for (int j = 0; j < fragments; ++j) { + wmma::fragment bf; + wmma::load_matrix_sync(bf, b + (n + j * 16) * stride + step, stride); + wmma::mma_sync(acc[j], af, bf, acc[j]); + } + } + } + // The next iteration's copy-wait barrier also protects B reuse; only the final epilogue needs this barrier. + __syncthreads(); + for (int j = 0; j < fragments; ++j) { + wmma::store_matrix_sync(c + m * kNvfp4TileN + n + j * 16, acc[j], kNvfp4TileN, wmma::mem_row_major); + } + __syncthreads(); + if constexpr (FuseSwiglu) { + for (int i = threadIdx.x; i < kNvfp4TileM * kNvfp4TileN / 2; i += blockDim.x) { + const int local_row = i / (kNvfp4TileN / 2); + const int local_col = (i % (kNvfp4TileN / 2)) * 2; + const int64_t row = first_row + local_row; + const int col = first_column + local_col; + if (row < end_row && col < columns) { + // Preserve the standalone GEMM's activation-dtype rounding before bias and SwiGLU. + float gate = static_cast(T(c[local_row * kNvfp4TileN + local_col])); + float linear = static_cast(T(c[local_row * kNvfp4TileN + local_col + 1])); + if (bias) { + gate += static_cast(bias[static_cast(expert) * columns + col]); + linear += static_cast(bias[static_cast(expert) * columns + col + 1]); + } + float activated; + if (alpha != 1.0f || beta != 0.0f || isfinite(limit)) { + if (isfinite(limit)) { + gate = fminf(gate, limit); + linear = fminf(fmaxf(linear, -limit), limit); + } + const float sigmoid = 1.0f / (1.0f + expf(-alpha * gate)); + linear += beta; + activated = gate * sigmoid * linear; + } else { + activated = cutlass::epilogue::thread::SiLu{}(gate)*linear; + } + output[row * (columns / 2) + col / 2] = T(activated); + } + } + } else { + for (int i = threadIdx.x; i < kNvfp4TileM * kNvfp4TileN; i += blockDim.x) { + const int64_t row = first_row + i / kNvfp4TileN; + const int col = first_column + i % kNvfp4TileN; + if (row < end_row && col < column_offset + output_columns) { + output[row * output_columns + col - column_offset] = T(c[i]); + } + } + } + __syncthreads(); + } +#endif +} + +template +__global__ void FinalizeNvfp4Columns(const ElementType* expanded, ElementType* output, + const ElementType* bias, const float* scales, + const int* reverse_map, const int* experts, int64_t rows, + int hidden, int columns, int column_offset, int top_k) { + constexpr int elements = 8; + using Values = cutlass::Array; + using Compute = cutlass::Array; + const int64_t row = blockIdx.x; + for (int vector = threadIdx.x; vector < columns / elements; vector += blockDim.x) { + Compute result; + result.fill(0); + for (int rank = 0; rank < top_k; ++rank) { + const int64_t route = row * top_k + rank; + const int permuted = reverse_map[row + rank * rows]; + const auto* input = reinterpret_cast(expanded + static_cast(permuted) * columns); + Compute value = arrayConvert(input[vector]); + if (bias) { + const auto* expert_bias = reinterpret_cast( + bias + static_cast(experts[route]) * hidden + column_offset); + value = value + arrayConvert(expert_bias[vector]); + } + result = result + scales[route] * value; + } + auto* target = reinterpret_cast(output + row * hidden + column_offset); + target[vector] = arrayConvert(result); + } +} + +template +__global__ void FakeInt4Activations(ElementType* values, int64_t groups) { + __shared__ float magnitudes[64]; + for (int64_t group = blockIdx.x; group < groups; group += gridDim.x) { + const int lane = threadIdx.x; + const int64_t index = group * 64 + lane; + const float value = static_cast(values[index]); + magnitudes[lane] = fabsf(value); + __syncthreads(); + for (int stride = 32; stride > 0; stride /= 2) { + if (lane < stride) { + magnitudes[lane] = fmaxf(magnitudes[lane], magnitudes[lane + stride]); + } + __syncthreads(); + } + const float scale = magnitudes[0] / 7.0f; + const float quantized = scale == 0.0f ? value : nearbyintf(value / scale) * scale; + values[index] = static_cast(quantized); + __syncthreads(); + } +} + +template +void SimulateInt4Activations(ElementType* values, int64_t rows, int columns, cudaStream_t stream) { + ORT_ENFORCE(columns > 0 && columns % 64 == 0, "INT4 activation simulation requires groups of 64."); + const int64_t groups = SafeInt(rows) * (columns / 64); + const int blocks = static_cast(std::min(groups, 65535)); + FakeInt4Activations<<>>(values, groups); + CUDA_CALL_THROW(cudaGetLastError()); +} + struct PrefillWorkspace { size_t bytes = 0; size_t permuted_rows; @@ -25,6 +303,7 @@ struct PrefillWorkspace { size_t fc1_output; size_t activated_output; size_t fc2_output; + size_t tile_offsets; size_t Reserve(size_t count, size_t element_size) { const size_t offset = bytes; @@ -38,16 +317,34 @@ struct PrefillWorkspace { const int64_t block_rows = computeNumTokensPerBlock(params.num_rows, params.num_experts); const size_t blocks = (params.num_rows + block_rows - 1) / block_rows; const size_t count_elements = SafeInt(params.num_experts) * blocks; + const int input_columns = params.is_nvfp4 && params.tile_nvfp4_down_output + ? std::min(params.hidden_size, kNvfp4DownColumns) + : params.hidden_size; permuted_rows = Reserve(expanded, sizeof(int)); permuted_experts = Reserve(expanded, sizeof(int)); expert_offsets = Reserve(SafeInt(params.num_experts) + 1, sizeof(int64_t)); - counts = Reserve(count_elements, sizeof(int)); - cumulative_counts = Reserve(count_elements, sizeof(int)); - blocked_rows = Reserve(SafeInt(params.num_rows) * params.num_experts, sizeof(int)); - expanded_input = Reserve(SafeInt(expanded) * params.hidden_size, sizeof(half)); - fc1_output = Reserve(SafeInt(expanded) * params.inter_size * 2, sizeof(half)); + if (params.is_nvfp4 && params.reuse_nvfp4_routing_scratch) { + // Sorting completes before expansion; retain maps/offsets but reuse its temporary counters. + const size_t scratch_start = bytes; + counts = Reserve(count_elements, sizeof(int)); + cumulative_counts = Reserve(count_elements, sizeof(int)); + blocked_rows = Reserve(SafeInt(params.num_rows) * params.num_experts, sizeof(int)); + const size_t routing_end = bytes; + bytes = scratch_start; + expanded_input = Reserve(SafeInt(expanded) * input_columns, sizeof(half)); + bytes = std::max(bytes, routing_end); + } else { + counts = Reserve(count_elements, sizeof(int)); + cumulative_counts = Reserve(count_elements, sizeof(int)); + blocked_rows = Reserve(SafeInt(params.num_rows) * params.num_experts, sizeof(int)); + expanded_input = Reserve(SafeInt(expanded) * input_columns, sizeof(half)); + } + fc1_output = params.is_nvfp4 && params.fuse_nvfp4_swiglu + ? 0 + : Reserve(SafeInt(expanded) * params.inter_size * 2, sizeof(half)); activated_output = Reserve(SafeInt(expanded) * params.inter_size, sizeof(half)); fc2_output = expanded_input; + tile_offsets = params.is_nvfp4 ? Reserve(SafeInt(params.num_experts) + 1, sizeof(int)) : 0; } }; @@ -55,7 +352,23 @@ template void RunPackedGroupedGemm(const Int2MoePrefillParams& params, const ElementType* activations, const uint8_t* weights, const void* scales, const int64_t* offsets, ElementType* output, - int64_t num_rows, int num_columns, int reduction_size, int weight_bits) { + int64_t num_rows, int num_columns, int reduction_size, int weight_bits, + const int* tile_offsets, const float* global_scales, bool fuse_swiglu = false, + const int* gathered_rows = nullptr, int column_offset = 0, int output_columns = 0) { + if (params.is_nvfp4) { + auto kernel = gathered_rows + ? (fuse_swiglu ? Nvfp4GroupedGemm : Nvfp4GroupedGemm) + : (fuse_swiglu ? Nvfp4GroupedGemm : Nvfp4GroupedGemm); + CUDA_CALL_THROW(cudaFuncSetAttribute(kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, kNvfp4SharedBytes)); + kernel<<>>( + activations, weights, static_cast(scales), global_scales, + offsets, tile_offsets, output, num_columns, reduction_size, params.num_experts, + static_cast(params.fc1_bias), params.alpha, params.beta, params.limit, + column_offset, output_columns ? output_columns : num_columns, gathered_rows, params.num_rows); + CUDA_CALL_THROW(cudaGetLastError()); + return; + } if (weight_bits == 2) { Int2GroupedGemmParamsT gemm; gemm.activations = activations; @@ -96,35 +409,70 @@ void RunInt2MoePrefillImpl(const Int2MoePrefillParams& params, void* workspace) auto* activated_output = reinterpret_cast(storage + layout.activated_output); auto* fc2_output = reinterpret_cast(storage + layout.fc2_output); const int64_t expanded = params.num_rows * params.top_k; + auto* tile_offsets = reinterpret_cast(storage + layout.tile_offsets); threeStepBuildExpertMapsSortFirstToken( params.selected_experts, permuted_experts, permuted_rows, params.unpermuted_to_permuted, offsets, reinterpret_cast(storage + layout.counts), reinterpret_cast(storage + layout.cumulative_counts), reinterpret_cast(storage + layout.blocked_rows), params.num_rows, params.num_experts, params.top_k, 0, params.stream); + if (params.is_nvfp4) { + BuildNvfp4TileOffsets<<<1, 32, 0, params.stream>>>(offsets, tile_offsets, params.num_experts); + CUDA_CALL_THROW(cudaGetLastError()); + } const QuantParams quant_params{}; - expandInputRowsKernelLauncher( - static_cast(params.input), expanded_input, nullptr, nullptr, permuted_rows, - params.num_rows, params.hidden_size, - params.top_k, params.num_experts, quant_params, false, offsets, nullptr, nullptr, nullptr, params.stream); - - RunPackedGroupedGemm(params, expanded_input, params.fc1_weights, params.fc1_scales, - offsets, fc1_output, expanded, params.inter_size * 2, - params.hidden_size, params.fc1_weight_bits); - - ActivationParams activation(ActivationType::Swiglu); - activation.alpha = params.alpha; - activation.beta = params.beta; - activation.limit = params.limit; - activation.swiglu_fusion = 1; - doActivation( - activated_output, fc1_output, nullptr, static_cast(params.fc1_bias), true, offsets, - params.num_experts, params.inter_size, expanded, ActivationType::Swiglu, - quant_params, false, nullptr, params.stream, activation); + const bool tiled_down = params.is_nvfp4 && params.tile_nvfp4_down_output; + if (!tiled_down) { + expandInputRowsKernelLauncher( + static_cast(params.input), expanded_input, nullptr, nullptr, permuted_rows, + params.num_rows, params.hidden_size, + params.top_k, params.num_experts, quant_params, false, offsets, nullptr, nullptr, nullptr, params.stream); + } + + if (params.fake_int4_activations) { + SimulateInt4Activations(expanded_input, expanded, params.hidden_size, params.stream); + } + const bool fused = params.is_nvfp4 && params.fuse_nvfp4_swiglu; + RunPackedGroupedGemm(params, tiled_down ? static_cast(params.input) : expanded_input, + params.fc1_weights, params.fc1_scales, + offsets, fused ? activated_output : fc1_output, expanded, params.inter_size * 2, + params.hidden_size, params.fc1_weight_bits, tile_offsets, params.fc1_global_scales, fused, + tiled_down ? permuted_rows : nullptr); + + if (!fused) { + ActivationParams activation(ActivationType::Swiglu); + activation.alpha = params.alpha; + activation.beta = params.beta; + activation.limit = params.limit; + activation.swiglu_fusion = 1; + doActivation( + activated_output, fc1_output, nullptr, static_cast(params.fc1_bias), true, offsets, + params.num_experts, params.inter_size, expanded, ActivationType::Swiglu, + quant_params, false, nullptr, params.stream, activation); + } + if (params.fake_int4_activations) { + SimulateInt4Activations(activated_output, expanded, params.inter_size, params.stream); + } + if (tiled_down) { + for (int column = 0; column < params.hidden_size;) { + const int columns = std::min(kNvfp4DownColumns, params.hidden_size - column); + RunPackedGroupedGemm(params, activated_output, params.fc2_weights, params.fc2_scales, + offsets, fc2_output, expanded, params.hidden_size, + params.inter_size, params.fc2_weight_bits, tile_offsets, params.fc2_global_scales, + false, nullptr, column, columns); + FinalizeNvfp4Columns<<>>( + fc2_output, static_cast(params.output), static_cast(params.fc2_bias), + params.routing_weights, params.unpermuted_to_permuted, params.selected_experts, params.num_rows, + params.hidden_size, columns, column, params.top_k); + CUDA_CALL_THROW(cudaGetLastError()); + column += columns; + } + return; + } RunPackedGroupedGemm(params, activated_output, params.fc2_weights, params.fc2_scales, offsets, fc2_output, expanded, params.hidden_size, - params.inter_size, params.fc2_weight_bits); + params.inter_size, params.fc2_weight_bits, tile_offsets, params.fc2_global_scales); finalizeMoeRoutingKernelLauncher( fc2_output, static_cast(params.output), static_cast(params.fc2_bias), params.routing_weights, params.unpermuted_to_permuted, @@ -139,6 +487,23 @@ size_t GetInt2MoePrefillWorkspaceSize(const Int2MoePrefillParams& params) { } void RunInt2MoePrefill(const Int2MoePrefillParams& params, void* workspace) { + if (params.is_nvfp4) { + ORT_ENFORCE(!params.tile_nvfp4_down_output || params.fuse_nvfp4_swiglu, + "Tiled NVFP4 down output requires the fused gate/up epilogue."); + ORT_ENFORCE(params.sm >= 80 && params.multiprocessor_count > 0 && + params.num_rows > 0 && params.hidden_size > 0 && params.hidden_size % 64 == 0 && + params.inter_size > 0 && params.inter_size <= std::numeric_limits::max() / 2 && + params.inter_size % 64 == 0 && params.num_experts > 0 && + params.top_k > 0 && params.top_k <= params.num_experts && + SafeInt(params.num_rows) * params.top_k <= + std::numeric_limits::max() - params.num_experts, + "Packed NVFP4 prefill requires supported tensor-core shapes and routing sizes."); + ORT_ENFORCE(workspace && params.input && params.output && params.selected_experts && + params.routing_weights && params.unpermuted_to_permuted && + params.fc1_weights && params.fc2_weights && params.fc1_scales && params.fc2_scales && + params.fc1_global_scales && params.fc2_global_scales, + "Packed NVFP4 prefill requires input, output, routing, weight and scale buffers."); + } if (params.is_bf16) { #ifdef ENABLE_BF16 RunInt2MoePrefillImpl<__nv_bfloat16>(params, workspace); diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_int2_prefill.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_int2_prefill.h index 54e98c316a5fc..f5ff72dcfb1f0 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_int2_prefill.h +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_int2_prefill.h @@ -23,6 +23,13 @@ struct Int2MoePrefillParams { int* unpermuted_to_permuted = nullptr; void* output = nullptr; bool is_bf16 = false; + bool fake_int4_activations = false; + bool is_nvfp4 = false; + bool fuse_nvfp4_swiglu = false; + bool reuse_nvfp4_routing_scratch = false; + bool tile_nvfp4_down_output = false; + const float* fc1_global_scales = nullptr; + const float* fc2_global_scales = nullptr; int64_t num_rows = 0; int hidden_size = 0; int inter_size = 0; diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc index a48b2b9ca8e3c..9bfa2722b0c77 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc @@ -38,12 +38,27 @@ #include #include +#include +#include +#include #include +#include #include #include #include #include +#ifdef _WIN32 +#include +#include +#include +#include +#else +#include +#include +#include +#endif + using namespace onnxruntime::cuda; using namespace ::onnxruntime::common; using namespace ONNX_NAMESPACE; @@ -51,6 +66,73 @@ using namespace ONNX_NAMESPACE; namespace { constexpr size_t kMinH200GlobalMemoryBytes = size_t{120} << 30; +#if !defined(BUILD_CUDA_EP_AS_PLUGIN) && !defined(ORT_MINIMAL_BUILD) +std::string EncodeCalibrationCaptureName(const std::string& name) { + constexpr char kHex[] = "0123456789ABCDEF"; + std::string encoded; + encoded.reserve(name.size()); + for (const unsigned char character : name) { + if ((character >= 'a' && character <= 'z') || + (character >= 'A' && character <= 'Z') || + (character >= '0' && character <= '9')) { + encoded.push_back(static_cast(character)); + } else { + encoded.push_back('_'); + encoded.push_back(kHex[character >> 4]); + encoded.push_back(kHex[character & 0x0F]); + } + } + return encoded; +} + +Status WriteCalibrationCapture( + const std::filesystem::path& path, + const std::array& header, + const std::vector& inputs, + const std::vector& routers) { + int fd = -1; +#ifdef _WIN32 + const errno_t open_error = _wsopen_s( + &fd, path.c_str(), _O_CREAT | _O_EXCL | _O_BINARY | _O_WRONLY, + _SH_DENYRW, _S_IREAD | _S_IWRITE); + if (open_error != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Cannot exclusively create calibration capture: ", + path.string(), " (error ", open_error, ")"); + } + FILE* output = _fdopen(fd, "wb"); +#else + fd = open(path.c_str(), O_CREAT | O_EXCL | O_WRONLY, 0600); + if (fd < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Cannot exclusively create calibration capture: ", + path.string(), " (error ", errno, ")"); + } + FILE* output = fdopen(fd, "wb"); +#endif + if (output == nullptr) { +#ifdef _WIN32 + _close(fd); +#else + close(fd); +#endif + std::error_code remove_error; + std::filesystem::remove(path, remove_error); + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Cannot open calibration capture stream: ", path.string()); + } + + const bool write_succeeded = + std::fwrite(header.data(), sizeof(header), 1, output) == 1 && + std::fwrite(inputs.data(), 1, inputs.size(), output) == inputs.size() && + std::fwrite(routers.data(), 1, routers.size(), output) == routers.size(); + const bool close_succeeded = std::fclose(output) == 0; + if (!write_succeeded || !close_succeeded) { + std::error_code remove_error; + std::filesystem::remove(path, remove_error); + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to write calibration capture: ", path.string()); + } + return Status::OK(); +} +#endif + void LogQMoESwigluFusionRemapOnce() { static std::once_flag log_warning; std::call_once(log_warning, []() { @@ -392,12 +474,41 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE bool is_fp16 = input_type == ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT16; #endif is_fp16_ = is_fp16; +#if !defined(BUILD_CUDA_EP_AS_PLUGIN) && !defined(ORT_MINIMAL_BUILD) + calibration_capture_directory_ = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_QMOE_CALIBRATION_CAPTURE_DIR", ""); + if (!calibration_capture_directory_.empty() && quant_type_ == "nvfp4") { + ORT_ENFORCE(std::filesystem::is_directory(calibration_capture_directory_), + "QMoE calibration capture directory must already exist."); + const std::string& node_name = op_kernel_info.node().Name(); + ORT_ENFORCE(!node_name.empty(), "QMoE calibration capture requires named nodes."); + static std::atomic next_capture_instance{0}; + calibration_capture_name_ = + EncodeCalibrationCaptureName(node_name) + "_p" + + std::to_string(logging::GetProcessId()) + "_i" + + std::to_string(next_capture_instance.fetch_add(1, std::memory_order_relaxed)); + LOGS_DEFAULT(WARNING) << "QMoE calibration capture enabled for " << op_kernel_info.node().Name() + << "; synchronous diagnostic capture, not suitable for performance measurements."; + } else { + calibration_capture_directory_.clear(); + } +#endif enable_int2_prefill_ = quant_type_ == "int" && sm_ >= 80 && (block_size_ == 32 || block_size_ == 64 || block_size_ == 128) && (fc1_expert_weight_bits_ == 2 || fc1_expert_weight_bits_ == 4) && (fc2_expert_weight_bits_ == 2 || fc2_expert_weight_bits_ == 4) && (fc1_expert_weight_bits_ == 2 || fc2_expert_weight_bits_ == 2) && onnxruntime::ParseEnvironmentVariableWithDefault("ORT_ENABLE_QMOE_INT2_PREFILL", 1) != 0; + enable_int2_large_experts_ = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_QMOE_INT2_LARGE_EXPERTS", 0) == 1; + const int fake_int4_activations = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_QMOE_INT2_FAKE_INT4_ACTIVATIONS", 0); + ORT_ENFORCE(fake_int4_activations == 0 || fake_int4_activations == 1, + "ORT_QMOE_INT2_FAKE_INT4_ACTIVATIONS must be 0 or 1."); + int2_fake_int4_activations_ = fake_int4_activations == 1 && quant_type_ == "int" && + (fc1_expert_weight_bits_ == 2 || fc2_expert_weight_bits_ == 2); + ORT_ENFORCE(!int2_fake_int4_activations_ || enable_int2_prefill_, + "INT4 activation simulation requires ORT_ENABLE_QMOE_INT2_PREFILL=1 and a supported INT2/mixed-width configuration."); fp4_deep_gemm_num_experts_ = StaticFp4DeepGemmNumExperts(op_kernel_info); enable_fp4_deep_gemm_ = @@ -719,6 +830,43 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { const bool uses_global_weight_scales = is_fp4_family || (is_fp8 && !is_block_fp8) || is_wfp4afp8; const Tensor* input = context->Input(0); const Tensor* router_probs = context->Input(1); +#if !defined(BUILD_CUDA_EP_AS_PLUGIN) && !defined(ORT_MINIMAL_BUILD) + if (!calibration_capture_directory_.empty()) { + std::lock_guard lock(calibration_capture_mutex_); + ORT_RETURN_IF_NOT(input && router_probs && input->Shape().NumDimensions() > 0 && + router_probs->Shape().NumDimensions() > 0, + "QMoE calibration capture requires nonempty input and router shapes."); + const int64_t hidden = input->Shape().GetDims().back(); + ORT_RETURN_IF_NOT(hidden > 0, "QMoE calibration capture requires positive hidden size."); + const int64_t rows = input->Shape().Size() / hidden; + if (rows >= 16 && calibration_capture_count_ < 16) { + cudaStream_t capture_stream = Stream(context); + cudaStreamCaptureStatus capture_status; + CUDA_CALL_THROW(cudaStreamIsCapturing(capture_stream, &capture_status)); + ORT_RETURN_IF_NOT(capture_status == cudaStreamCaptureStatusNone, + "QMoE calibration capture requires CUDA graphs disabled."); + ORT_RETURN_IF_NOT(input->DataType()->Size() == 2 && router_probs->DataType()->Size() == 2, + "QMoE calibration capture requires FP16/BF16 input and router."); + const int64_t experts = router_probs->Shape().GetDims().back(); + const int64_t captured_rows = std::min(rows, 128); + ORT_RETURN_IF_NOT(experts > 0 && router_probs->Shape().Size() >= captured_rows * experts, + "QMoE calibration router shape does not cover the captured rows."); + const size_t input_bytes = SafeInt(captured_rows) * hidden * 2; + const size_t router_bytes = SafeInt(captured_rows) * experts * 2; + std::vector inputs(input_bytes), routers(router_bytes); + CUDA_CALL_THROW(cudaMemcpyAsync(inputs.data(), input->DataRaw(), input_bytes, + cudaMemcpyDeviceToHost, capture_stream)); + CUDA_CALL_THROW(cudaMemcpyAsync(routers.data(), router_probs->DataRaw(), router_bytes, + cudaMemcpyDeviceToHost, capture_stream)); + CUDA_CALL_THROW(cudaStreamSynchronize(capture_stream)); + const auto path = std::filesystem::path(calibration_capture_directory_) / + (calibration_capture_name_ + "_" + std::to_string(calibration_capture_count_) + ".bin"); + const std::array header{1, captured_rows, hidden, experts, is_fp16_ ? 10 : 16}; + ORT_RETURN_IF_ERROR(WriteCalibrationCapture(path, header, inputs, routers)); + ++calibration_capture_count_; + } + } +#endif // When PrePack consumed the int4/int8 expert-weight initializers // (``weights_prepacked == false`` opt-in path), the original tensors // were freed; ``context->Input(2)/(5)`` would return nothing. @@ -1190,6 +1338,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { gemv_group_size, gemv::MoeGemvConfig::kDefault, fc2_gemv_sm80_layout, use_raw_nvfp4_gemv, nvfp4_weights_row_major_); } + const bool run_fp4_gemv = !use_fp4_deep_gemm && use_fp4_gemv; bool use_packed_int_gemv = false; const int64_t packed_int_expanded = moe_params.num_rows * static_cast(k_); @@ -1202,7 +1351,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { const bool packed_int_routing_supported = moe_params.num_rows <= 256 && moe_params.num_experts <= 256 && (k_ == 1 || k_ == 2 || k_ == 4 || k_ == 6 || k_ == 8); - if (enable_int2_gemv_ && !has_any_zero_point && is_fused_swiglu && swiglu_fusion == 1 && + if (enable_int2_gemv_ && !int2_fake_int4_activations_ && !has_any_zero_point && is_fused_swiglu && swiglu_fusion == 1 && (block_size_ == 64 || block_size_ == 128) && packed_fc1_weights_ != nullptr && packed_fc2_weights_ != nullptr) { use_packed_int_gemv = packed_int_fc1_shape_supported && packed_int_fc2_shape_supported && @@ -1229,22 +1378,45 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { moe_params.hidden_size > 0 && moe_params.hidden_size <= std::numeric_limits::max() && moe_params.hidden_size % 64 == 0 && moe_params.inter_size > 0 && moe_params.inter_size <= std::numeric_limits::max() / 2 && moe_params.inter_size % 64 == 0 && - moe_params.num_experts > 0 && moe_params.num_experts <= 256; + moe_params.num_experts > 0 && + moe_params.num_experts <= (enable_int2_large_experts_ ? 512 : 256); + ORT_RETURN_IF_NOT(!int2_fake_int4_activations_ || use_packed_int_prefill, + "INT4 activation simulation requires the packed INT2 grouped path; refusing an unquantized fallback."); + const bool use_nvfp4_packed_prefill = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_QMOE_NVFP4_PACKED_PREFILL", 1) == 1 && + is_nvfp4 && nvfp4_weights_row_major_ && use_fp4_dequant_fallback_ && !run_fp4_gemv && + is_fused_swiglu && swiglu_fusion == 1 && sm_ >= 80 && + moe_params.hidden_size > 0 && moe_params.hidden_size <= std::numeric_limits::max() && + moe_params.hidden_size % 64 == 0 && moe_params.inter_size > 0 && + moe_params.inter_size <= std::numeric_limits::max() / 2 && moe_params.inter_size % 64 == 0 && + moe_params.num_experts > 0 && moe_params.num_experts <= std::numeric_limits::max() && + moe_params.num_rows > 0 && SafeInt(moe_params.num_rows) * k_ <= std::numeric_limits::max() - moe_params.num_experts; int64_t effective_row_tile_size = row_tile_size_; + const bool fuse_nvfp4_swiglu = use_nvfp4_packed_prefill && + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_QMOE_NVFP4_FUSED_SWIGLU", 0) == 1; + const bool reuse_nvfp4_routing_scratch = use_nvfp4_packed_prefill && + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_QMOE_NVFP4_REUSE_ROUTING_SCRATCH", 0) == 1; + const bool tile_nvfp4_down_output = use_nvfp4_packed_prefill && fuse_nvfp4_swiglu && + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_QMOE_NVFP4_TILED_DOWN_OUTPUT", 0) == 1; size_t packed_prefill_workspace_bytes = 0; - if (use_packed_int_prefill) { + if (use_packed_int_prefill || use_nvfp4_packed_prefill) { namespace ck = onnxruntime::llm::kernels::cutlass_kernels; - constexpr size_t kMaxPackedPrefillScratchBytes = 256 * 1024 * 1024; + const size_t kMaxPackedPrefillScratchBytes = + (use_nvfp4_packed_prefill ? 384 : 256) * 1024 * 1024; ck::Int2MoePrefillParams sizing_params; + sizing_params.is_nvfp4 = use_nvfp4_packed_prefill; + sizing_params.fuse_nvfp4_swiglu = fuse_nvfp4_swiglu; + sizing_params.reuse_nvfp4_routing_scratch = reuse_nvfp4_routing_scratch; + sizing_params.tile_nvfp4_down_output = tile_nvfp4_down_output; sizing_params.hidden_size = static_cast(moe_params.hidden_size); sizing_params.inter_size = static_cast(moe_params.inter_size); sizing_params.num_experts = static_cast(moe_params.num_experts); sizing_params.top_k = static_cast(k_); size_t scale_scratch_bytes = 0; - if (!gemv_int_fc1_scales_ && fc1_scales) { + if (!use_nvfp4_packed_prefill && !gemv_int_fc1_scales_ && fc1_scales) { scale_scratch_bytes = SafeInt(scale_scratch_bytes) + fc1_scales->SizeInBytes(); } - if (!gemv_int_fc2_scales_ && fc2_scales) { + if (!use_nvfp4_packed_prefill && !gemv_int_fc2_scales_ && fc2_scales) { scale_scratch_bytes = SafeInt(scale_scratch_bytes) + fc2_scales->SizeInBytes(); } effective_row_tile_size = row_tile_size_ == qmoe::kDisabledRowTileSize @@ -1313,11 +1485,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { std::array runner_tile_configs{}; size_t runner_tile_config_count = 0; size_t workspace_size = 0; - // The packed INT GEMV path allocates its own scratch below and returns before reaching the - // dense grouped-GEMM tile loop, so profiling/sizing the dense runner here would be pure - // overhead (mutex, two dense-tactic profiling launches, and an unused large workspace - // allocation). Skip it entirely for that path; workspace_size stays 0. - if (!use_packed_int && !use_fp8_fused) { + // Standalone FP4 GEMV returns before runMoe and does not consume grouped-GEMM tactics or workspace. + if (!use_fp8_fused && !use_packed_int && !use_nvfp4_packed_prefill && !run_fp4_gemv) { std::lock_guard profiler_lock(mGemmProfilerMutex); // Profiling launches grouped-GEMM kernels, records/synchronizes CUDA events, and @@ -1876,7 +2045,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { } }; - if (!use_int_dequant_fallback || use_packed_int) { + if ((!use_int_dequant_fallback || use_packed_int) && !use_nvfp4_packed_prefill) { const auto& fc1_prepared_scales = use_packed_int ? gemv_int_fc1_scales_ : packed_fc1_scales_; const auto& fc2_prepared_scales = use_packed_int ? gemv_int_fc2_scales_ : packed_fc2_scales_; prepare_scale_zp(fc1_scales, fc1_zeros, fc1_prepared_scales, packed_fc1_bias_, @@ -2041,14 +2210,23 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { } #endif - if (use_packed_int_prefill) { + if (use_packed_int_prefill || use_nvfp4_packed_prefill) { namespace ck = onnxruntime::llm::kernels::cutlass_kernels; ck::Int2MoePrefillParams params; params.is_bf16 = !is_fp16_; - params.fc1_weights = static_cast(packed_fc1_weights_.get()); - params.fc2_weights = static_cast(packed_fc2_weights_.get()); - params.fc1_scales = p_fc1_scales; - params.fc2_scales = p_fc2_scales; + params.fake_int4_activations = int2_fake_int4_activations_; + params.is_nvfp4 = use_nvfp4_packed_prefill; + params.fuse_nvfp4_swiglu = fuse_nvfp4_swiglu; + params.reuse_nvfp4_routing_scratch = reuse_nvfp4_routing_scratch; + params.tile_nvfp4_down_output = tile_nvfp4_down_output; + params.fc1_weights = static_cast( + use_nvfp4_packed_prefill ? fc1_experts_weights->DataRaw() : packed_fc1_weights_.get()); + params.fc2_weights = static_cast( + use_nvfp4_packed_prefill ? fc2_experts_weights->DataRaw() : packed_fc2_weights_.get()); + params.fc1_scales = use_nvfp4_packed_prefill ? fp4_fc1_block_scales->DataRaw() : p_fc1_scales; + params.fc2_scales = use_nvfp4_packed_prefill ? fp4_fc2_block_scales->DataRaw() : p_fc2_scales; + params.fc1_global_scales = use_nvfp4_packed_prefill ? fc1_global_scale->Data() : nullptr; + params.fc2_global_scales = use_nvfp4_packed_prefill ? fc2_global_scale->Data() : nullptr; params.fc1_bias = fc1_experts_bias_optional ? fc1_experts_bias_optional->DataRaw() : nullptr; params.fc2_bias = fc2_experts_bias_optional ? fc2_experts_bias_optional->DataRaw() : nullptr; params.selected_experts = expert_indices; @@ -2095,7 +2273,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { #endif } if (enable_kernel_debug_info_) { - PrintQMoEKernelDebugInfo("packed_int_prefill", moe_params.num_rows, row_tile_plan.rows_per_tile, + PrintQMoEKernelDebugInfo(use_nvfp4_packed_prefill ? "packed_nvfp4_prefill" : "packed_int_prefill", + moe_params.num_rows, row_tile_plan.rows_per_tile, final_tile_rows, row_tile_plan.rows_per_tile * k_, final_tile_rows * k_, prefill_bytes, SafeInt(total_scratch_bytes) + prefill_bytes); } @@ -2204,7 +2383,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { } // The FP4 DeepGEMM path, when eligible, outranks the GEMV for the same shapes. - if (!use_fp4_deep_gemm && use_fp4_gemv) { + if (run_fp4_gemv) { namespace gemv = onnxruntime::llm::kernels::moe_gemv; namespace ck = onnxruntime::llm::kernels::cutlass_kernels; const int num_experts = static_cast(moe_params.num_experts); @@ -2922,7 +3101,8 @@ Status QMoE::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, } const int64_t reduction_size = shape[2] * pack_factor; - if (enable_int2_prefill_ && shape[0] > 0 && shape[0] <= 256 && + if (enable_int2_prefill_ && shape[0] > 0 && + shape[0] <= (enable_int2_large_experts_ ? 512 : 256) && shape[1] > 0 && shape[1] <= std::numeric_limits::max() && shape[1] % 64 == 0 && reduction_size > 0 && reduction_size <= std::numeric_limits::max() && reduction_size % 64 == 0) { return true; diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h index c77bf8a27c389..4f26c07b35455 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h @@ -115,6 +115,14 @@ class QMoE final : public CudaKernel, public MoEBase { bool enable_kernel_debug_info_ = false; bool enable_int2_gemv_ = false; bool enable_int2_prefill_ = false; + bool enable_int2_large_experts_ = false; + bool int2_fake_int4_activations_ = false; +#if !defined(BUILD_CUDA_EP_AS_PLUGIN) && !defined(ORT_MINIMAL_BUILD) + std::string calibration_capture_directory_; + std::string calibration_capture_name_; + mutable std::mutex calibration_capture_mutex_; + mutable size_t calibration_capture_count_ = 0; +#endif int64_t row_tile_size_ = qmoe::kDisabledRowTileSize; int64_t int_dequant_max_scratch_bytes_ = int64_t{1} << 30; diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu index 331351a6ac33b..63e1bcf2871e4 100644 --- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu +++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu @@ -1036,26 +1036,16 @@ __device__ __forceinline__ float DecodeFloat8E4M3FN(uint8_t code) { return e4m3_to_float(code); } #endif - // ONNX float8e4m3fn has no infinities. The only NaN payloads are 0x7F/0xFF; - // finite values, including the max finite code 0x7E, use the normal E4M3 formula. - const int sign = code & 0x80; - const int exponent = (code >> 3) & 0x0F; - const int mantissa = code & 0x07; - - if ((code & 0x7F) == 0) { - return sign ? -0.0f : 0.0f; - } - if (exponent == 0x0F && mantissa == 0x07) { + const uint32_t sign = static_cast(code & 0x80) << 24; + const uint32_t magnitude = code & 0x7F; + if (magnitude == 0x7F) { return __int_as_float(0x7fffffff); } - - float value = 0.0f; - if (exponent == 0) { - value = ldexpf(static_cast(mantissa), -9); - } else { - value = ldexpf(1.0f + static_cast(mantissa) * 0.125f, exponent - 7); + if (magnitude < 8) { + return __uint_as_float(sign | __float_as_uint(static_cast(magnitude) * (1.0f / 512.0f))); } - return sign ? -value : value; + // Every normal E4M3FN value has an exact FP32 exponent/mantissa representation. + return __uint_as_float(sign | (((magnitude >> 3) + 120) << 23) | ((magnitude & 7) << 20)); } // Tile shape for QMoEDequantizeFp4WeightsVecKernel. kTileN = 64 rows is exactly 32 packed bytes, @@ -1555,8 +1545,15 @@ void LaunchQMoECompactExperts( CUDA_CALL_THROW(cudaGetLastError()); } -template -__device__ __forceinline__ T QMoEFp8Weight(const QMoEFp8ProjectionParams& params, int expert, int row, int col) { +struct QMoEFp8WeightRow { + const uint8_t* weights; + const void* scales; + int scale_type; + int64_t scale_offset; +}; + +__device__ __forceinline__ QMoEFp8WeightRow ResolveQMoEFp8WeightRow( + const QMoEFp8ProjectionParams& params, int expert, int row) { const uint8_t* weights = params.weights; const void* scales = params.scales; int scale_type = params.scale_type; @@ -1574,13 +1571,23 @@ __device__ __forceinline__ T QMoEFp8Weight(const QMoEFp8ProjectionParams& params } const int scale_n = (source_n + params.block_size - 1) / params.block_size; const int scale_k = (params.k + params.block_size - 1) / params.block_size; - const int64_t scale_index = (static_cast(expert) * scale_n + row / params.block_size) * scale_k + - col / params.block_size; - const int64_t weight_index = (static_cast(expert) * source_n + row) * params.k + col; - const float scale = scale_type == 0 ? static_cast(scales)[scale_index] - : scale_type == 1 ? __half2float(static_cast(scales)[scale_index]) - : __bfloat162float(static_cast(scales)[scale_index]); - return static_cast(DecodeFloat8E4M3FN(weights[weight_index]) * scale); + const int64_t scale_offset = (static_cast(expert) * scale_n + row / params.block_size) * scale_k; + const int64_t weight_offset = (static_cast(expert) * source_n + row) * params.k; + return {weights + weight_offset, scales, scale_type, scale_offset}; +} + +__device__ __forceinline__ float QMoEFp8RowScale(const QMoEFp8WeightRow& row, int block) { + const int64_t index = row.scale_offset + block; + return row.scale_type == 0 ? static_cast(row.scales)[index] + : row.scale_type == 1 ? __half2float(static_cast(row.scales)[index]) + : __bfloat162float(static_cast(row.scales)[index]); +} + +template +__device__ __forceinline__ T QMoEFp8Weight(const QMoEFp8ProjectionParams& params, int expert, int row, int col) { + const auto weight_row = ResolveQMoEFp8WeightRow(params, expert, row); + const float scale = QMoEFp8RowScale(weight_row, col / params.block_size); + return static_cast(DecodeFloat8E4M3FN(weight_row.weights[col]) * scale); } template @@ -1592,10 +1599,27 @@ __global__ void QMoEFp8GemvKernel(QMoEFp8ProjectionParams params, const T* input return; } const int input_row = params.row_to_unpermuted ? params.row_to_unpermuted[row] % params.num_rows : row; + const auto weight_row = ResolveQMoEFp8WeightRow(params, expert, output_col); + const T* activation = input + static_cast(input_row) * params.k; float sum = 0.0f; - for (int col = threadIdx.x % 32; col < params.k; col += 32) { - sum = fmaf(static_cast(input[static_cast(input_row) * params.k + col]), - static_cast(QMoEFp8Weight(params, expert, output_col, col)), sum); + if (params.block_size == 128) { + for (int block_begin = 0; block_begin < params.k; block_begin += 128) { + const float scale = QMoEFp8RowScale(weight_row, block_begin / 128); +#pragma unroll + for (int offset = 0; offset < 128; offset += 32) { + const int col = block_begin + threadIdx.x % 32 + offset; + if (col < params.k) { + const T weight = static_cast(DecodeFloat8E4M3FN(weight_row.weights[col]) * scale); + sum = fmaf(static_cast(activation[col]), static_cast(weight), sum); + } + } + } + } else { + for (int col = threadIdx.x % 32; col < params.k; col += 32) { + const T weight = static_cast(DecodeFloat8E4M3FN(weight_row.weights[col]) * + QMoEFp8RowScale(weight_row, col / params.block_size)); + sum = fmaf(static_cast(activation[col]), static_cast(weight), sum); + } } sum = WarpReduceSum(sum); if (threadIdx.x % 32 == 0) { @@ -1936,6 +1960,44 @@ __global__ void QMoEDequantizeNvfp4WeightsKernel( } } +template +__global__ void QMoEDequantizeNvfp4RowMajorWeightsVecKernel( + const uint8_t* __restrict__ packed_weights, + const uint8_t* __restrict__ block_scales, + const float* __restrict__ global_scales, + T* __restrict__ output, + int n, int k, const int* compact_to_expert, + const T* bias, T* output_bias) { + constexpr int kVecK = kQMoEDequantizeFp4VecK; + const int row = blockIdx.x; + const int k_base = (static_cast(blockIdx.y) * blockDim.x + threadIdx.x) * kVecK; + if (k_base >= k) { + return; + } + const int output_expert = blockIdx.z; + const int expert = compact_to_expert ? compact_to_expert[output_expert] : output_expert; + if (expert < 0) { + return; + } + const int64_t source_row = static_cast(expert) * n + row; + const uint32_t packed = *reinterpret_cast(packed_weights + source_row * (k / 2) + k_base / 2); + const float block_scale = DecodeFloat8E4M3FN(block_scales[source_row * (k / 16) + k_base / 16]); + const float global_scale = global_scales[expert]; + uint4 staged; + T* values = reinterpret_cast(&staged); + static_assert(kVecK * sizeof(T) == sizeof(uint4)); +#pragma unroll + for (int j = 0; j < kVecK; ++j) { + // Keep the scalar path's multiplication order and final FP16/BF16 rounding. + values[j] = static_cast(DecodeFp4E2M1(static_cast((packed >> (4 * j)) & 0x0F)) * + block_scale * global_scale); + } + *reinterpret_cast(output + (static_cast(output_expert) * n + row) * k + k_base) = staged; + if (bias && k_base == 0) { + output_bias[static_cast(output_expert) * n + row] = bias[source_row]; + } +} + template void LaunchQMoEDequantizeNvfp4WeightsImpl( const uint8_t* packed_weights, @@ -1952,6 +2014,16 @@ void LaunchQMoEDequantizeNvfp4WeightsImpl( bool weights_row_major) { ORT_ENFORCE(bias == nullptr || output_bias != nullptr, "QMoE NVFP4 bias gathering requires an output buffer."); constexpr int block = 256; + const int64_t row_major_grid_y = + (static_cast(k) + block * kQMoEDequantizeFp4VecK - 1) / + (block * kQMoEDequantizeFp4VecK); + if (weights_row_major && k % 16 == 0 && row_major_grid_y <= 65535 && num_experts <= 65535) { + const dim3 grid(n, onnxruntime::narrow(row_major_grid_y), num_experts); + QMoEDequantizeNvfp4RowMajorWeightsVecKernel<<>>( + packed_weights, block_scales, global_scales, output, n, k, compact_to_expert, bias, output_bias); + CUDA_CALL_THROW(cudaGetLastError()); + return; + } if (!weights_row_major && !compact_to_expert && !bias && QMoEDequantizeFp4VecApplies<16>(num_experts, n, k)) { const dim3 tile_block(kQMoEDequantizeFp4TileK / kQMoEDequantizeFp4VecK, kQMoEDequantizeFp4TileN); const dim3 tile_grid((n + kQMoEDequantizeFp4TileN - 1) / kQMoEDequantizeFp4TileN, diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu index 3304a7839dd3d..9d72330156740 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu @@ -36,6 +36,7 @@ #include "contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h" #include "contrib_ops/cuda/sparse/sparse_attention_indexer_device_math.cuh" +#include "core/platform/env_var_utils.h" #include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" #include "core/providers/cuda/cu_inc/topk_warp_sort.cuh" @@ -56,7 +57,9 @@ constexpr int kBoundedTopKItemsPerThread = kBoundedTopKMax / kThreads; constexpr int kDistributedTopKMaxRows = 64; constexpr int kDistributedTopKMinBlocks = 2048; constexpr int kHierarchicalTileBlocks = 8; +constexpr int kHierarchicalTopKTileBlocks = 32; constexpr int kHierarchicalScoreThreads = 1024; +constexpr int kHierarchicalScoreMaxBlocksPerRow = 128; constexpr int kHierarchicalMergeThreads = 256; constexpr int kQwenHeadSize = 128; constexpr int kQwenCompressRatio = 4; @@ -89,7 +92,7 @@ bool UseRaggedQsaPrefill(const PackedSparseAttentionIndexerParams& params) { } int GetHierarchicalTileCount(const PackedSparseAttentionIndexerParams& params) { - return (params.state_capacity + kHierarchicalTileBlocks - 1) / kHierarchicalTileBlocks; + return (params.state_capacity + kHierarchicalTopKTileBlocks - 1) / kHierarchicalTopKTileBlocks; } template @@ -400,7 +403,8 @@ __global__ void QsaBlockScoreKernel(const T* present_key_state, const float* que // One CTA per packed query row. Eight state blocks are scored concurrently, with one warp per // (state block, query head). Unlike the generic rectangular launch, each CTA stops at its row's // live causal block count and never launches or writes work for unused fixed-capacity slots. -template +// WarpHeads instead keeps all heads in one warp to reuse keys without CTA-wide barriers. +template __global__ void QsaRaggedPrefillScoreKernel(const T* present_key_state, const float* query_rotated, const int32_t* cumulative_sequence_lengths, const int32_t* past_sequence_lengths, const int64_t* position_ids, @@ -410,7 +414,7 @@ __global__ void QsaRaggedPrefillScoreKernel(const T* present_key_state, const fl const int warp = static_cast(threadIdx.x) / kWarpSize; const int lane = static_cast(threadIdx.x) % kWarpSize; - const int block_slot = warp / kQwenNumHeads; + const int block_slot = WarpHeads ? warp : warp / kQwenNumHeads; const int head = warp % kQwenNumHeads; for (int token = static_cast(blockIdx.x); token < params.total_tokens; token += static_cast(gridDim.x)) { const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token); @@ -423,36 +427,64 @@ __global__ void QsaRaggedPrefillScoreKernel(const T* present_key_state, const fl const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast(key_len_after); const int block_count = overflow_flags[batch] == 0 ? static_cast(visible_64 < 0 ? 0 : visible_64) : 0; - for (int tile_start = 0; tile_start < block_count; tile_start += kHierarchicalTileBlocks) { - const int block_index = tile_start + block_slot; - float score = 0.0f; - if (block_index < block_count) { - const int d = lane * 4; - const int64_t query_base = - (static_cast(token) * kQwenNumHeads + head) * kQwenHeadSize; + if constexpr (WarpHeads) { + float4 queries[kQwenNumHeads]; +#pragma unroll + for (int query_head = 0; query_head < kQwenNumHeads; ++query_head) { + queries[query_head] = LoadQsaVector4( + query_rotated + (static_cast(token) * kQwenNumHeads + query_head) * kQwenHeadSize + lane * 4); + } + for (int block_index = block_slot; block_index < block_count; block_index += kHierarchicalTileBlocks) { const int64_t key_base = (static_cast(batch) * params.state_capacity + block_index) * kQwenHeadSize; - const float4 query_value = LoadQsaVector4(query_rotated + query_base + d); - const float4 key_value = LoadQsaVector4(present_key_state + key_base + d); - score = query_value.x * key_value.x + query_value.y * key_value.y + - query_value.z * key_value.z + query_value.w * key_value.w; - score = topk::WarpReduceSum(score); - } - if (lane == 0) { - head_scores[block_slot][head] = fmaxf(score, 0.0f); - } - __syncthreads(); - if (threadIdx.x < kHierarchicalTileBlocks && tile_start + threadIdx.x < block_count) { + const float4 key_value = LoadQsaVector4(present_key_state + key_base + lane * 4); float total_score = 0.0f; #pragma unroll for (int query_head = 0; query_head < kQwenNumHeads; ++query_head) { - total_score += head_scores[threadIdx.x][query_head]; + const float4 query_value = queries[query_head]; + float score = query_value.x * key_value.x + query_value.y * key_value.y + + query_value.z * key_value.z + query_value.w * key_value.w; + score = topk::WarpReduceSum(score); + total_score += fmaxf(score, 0.0f); + } + if (lane == 0) { + const float scaled_score = total_score * params.scale; + block_scores[static_cast(token) * params.state_capacity + block_index] = + scaled_score == 0.0f ? 0.0f : scaled_score; } - const float scaled_score = total_score * params.scale; - block_scores[static_cast(token) * params.state_capacity + tile_start + threadIdx.x] = - scaled_score == 0.0f ? 0.0f : scaled_score; } - __syncthreads(); + } else { + for (int tile_start = 0; tile_start < block_count; tile_start += kHierarchicalTileBlocks) { + const int block_index = tile_start + block_slot; + float score = 0.0f; + if (block_index < block_count) { + const int d = lane * 4; + const int64_t query_base = + (static_cast(token) * kQwenNumHeads + head) * kQwenHeadSize; + const int64_t key_base = + (static_cast(batch) * params.state_capacity + block_index) * kQwenHeadSize; + const float4 query_value = LoadQsaVector4(query_rotated + query_base + d); + const float4 key_value = LoadQsaVector4(present_key_state + key_base + d); + score = query_value.x * key_value.x + query_value.y * key_value.y + + query_value.z * key_value.z + query_value.w * key_value.w; + score = topk::WarpReduceSum(score); + } + if (lane == 0) { + head_scores[block_slot][head] = fmaxf(score, 0.0f); + } + __syncthreads(); + if (threadIdx.x < kHierarchicalTileBlocks && tile_start + threadIdx.x < block_count) { + float total_score = 0.0f; +#pragma unroll + for (int query_head = 0; query_head < kQwenNumHeads; ++query_head) { + total_score += head_scores[threadIdx.x][query_head]; + } + const float scaled_score = total_score * params.scale; + block_scores[static_cast(token) * params.state_capacity + tile_start + threadIdx.x] = + scaled_score == 0.0f ? 0.0f : scaled_score; + } + __syncthreads(); + } } } } @@ -464,14 +496,16 @@ __global__ void QsaScoreTileTopKKernel(const T* present_key_state, const float* const int32_t* present_state_lengths, const int32_t* overflow_flags, uint64_t* tile_keys, int tile_count, PackedSparseAttentionIndexerParams params) { - __shared__ uint64_t keys[kHierarchicalTileBlocks]; - __shared__ float head_scores[kHierarchicalTileBlocks][kQwenNumHeads]; + __shared__ uint64_t keys[kHierarchicalTopKTileBlocks]; + __shared__ float head_scores[kHierarchicalTopKTileBlocks][kQwenNumHeads]; const int warp = static_cast(threadIdx.x) / kWarpSize; const int lane = static_cast(threadIdx.x) % kWarpSize; const int block_slot = warp / kQwenNumHeads; const int head = warp % kQwenNumHeads; - const int token = static_cast(blockIdx.y); + const int blocks_per_row = min(tile_count, kHierarchicalScoreMaxBlocksPerRow); + const int first_tile = static_cast(blockIdx.x) % blocks_per_row; + const int token = static_cast(blockIdx.x) / blocks_per_row; const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token); const int key_len_after = present_state_lengths[batch * 2 + psai::kKeyStateLength]; const int64_t abs_position = @@ -481,35 +515,42 @@ __global__ void QsaScoreTileTopKKernel(const T* present_key_state, const float* const int64_t causal_count = SaiCausalThreshold(abs_position, kQwenCompressRatio); const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast(key_len_after); const int block_count = overflow_flags[batch] == 0 ? static_cast(visible_64 < 0 ? 0 : visible_64) : 0; - const int active_tiles = (block_count + kHierarchicalTileBlocks - 1) / kHierarchicalTileBlocks; - const int key_stride = tile_count * kHierarchicalTileBlocks; + + const int live_tiles = (block_count + kHierarchicalTopKTileBlocks - 1) / kHierarchicalTopKTileBlocks; + const int key_stride = tile_count * kHierarchicalTopKTileBlocks; const int64_t row_base = static_cast(token) * key_stride; - for (int index = active_tiles * kHierarchicalTileBlocks + blockIdx.x * blockDim.x + threadIdx.x; - index < key_stride; index += gridDim.x * blockDim.x) { + // Initialize only the fully padded suffix; scoring writes every key in live tiles. + for (int index = live_tiles * kHierarchicalTopKTileBlocks + first_tile * blockDim.x + threadIdx.x; + index < key_stride; index += blocks_per_row * blockDim.x) { tile_keys[row_base + index] = topk::kPaddingSortKey; } - - for (int tile = static_cast(blockIdx.x); tile < active_tiles; tile += static_cast(gridDim.x)) { - const int block_index = tile * kHierarchicalTileBlocks + block_slot; - float score = 0.0f; - if (block_index < block_count) { - const int d = lane * 4; - const int64_t query_base = - (static_cast(token) * kQwenNumHeads + head) * kQwenHeadSize; - const int64_t key_base = - (static_cast(batch) * params.state_capacity + block_index) * kQwenHeadSize; - const float4 query_value = LoadQsaVector4(query_rotated + query_base + d); - const float4 key_value = LoadQsaVector4(present_key_state + key_base + d); - score = query_value.x * key_value.x + query_value.y * key_value.y + - query_value.z * key_value.z + query_value.w * key_value.w; - score = topk::WarpReduceSum(score); - } - if (lane == 0) { - head_scores[block_slot][head] = fmaxf(score, 0.0f); + if (first_tile >= live_tiles) { + return; + } + const int d = lane * 4; + const int64_t query_base = + (static_cast(token) * kQwenNumHeads + head) * kQwenHeadSize; + const float4 query_value = LoadQsaVector4(query_rotated + query_base + d); + for (int tile = first_tile; tile < live_tiles; tile += blocks_per_row) { + const int64_t output_base = row_base + tile * kHierarchicalTopKTileBlocks; + for (int slot = block_slot; slot < kHierarchicalTopKTileBlocks; slot += kHierarchicalTileBlocks) { + const int block_index = tile * kHierarchicalTopKTileBlocks + slot; + float score = 0.0f; + if (block_index < block_count) { + const int64_t key_base = + (static_cast(batch) * params.state_capacity + block_index) * kQwenHeadSize; + const float4 key_value = LoadQsaVector4(present_key_state + key_base + d); + score = query_value.x * key_value.x + query_value.y * key_value.y + + query_value.z * key_value.z + query_value.w * key_value.w; + score = topk::WarpReduceSum(score); + } + if (lane == 0) { + head_scores[slot][head] = fmaxf(score, 0.0f); + } } __syncthreads(); - if (threadIdx.x < kHierarchicalTileBlocks) { - const int tile_block = tile * kHierarchicalTileBlocks + static_cast(threadIdx.x); + if (threadIdx.x < kHierarchicalTopKTileBlocks) { + const int tile_block = tile * kHierarchicalTopKTileBlocks + static_cast(threadIdx.x); uint64_t key = topk::kPaddingSortKey; if (tile_block < block_count) { float total_score = 0.0f; @@ -523,20 +564,21 @@ __global__ void QsaScoreTileTopKKernel(const T* present_key_state, const float* keys[threadIdx.x] = key; } __syncthreads(); - if (threadIdx.x == 0) { - for (int index = 1; index < kHierarchicalTileBlocks; ++index) { - const uint64_t key = keys[index]; - int insertion = index; - while (insertion > 0 && key > keys[insertion - 1]) { - keys[insertion] = keys[insertion - 1]; - --insertion; + if (warp == 0) { + uint64_t key = keys[lane]; +#pragma unroll + for (int width = 2; width <= kHierarchicalTopKTileBlocks; width *= 2) { +#pragma unroll + for (int stride = width / 2; stride > 0; stride /= 2) { + const uint64_t other = __shfl_xor_sync(0xffffffffU, key, stride); + const bool take_maximum = ((lane & width) == 0) == ((lane & stride) == 0); + key = take_maximum ? max(key, other) : min(key, other); } - keys[insertion] = key; } + keys[lane] = key; } __syncthreads(); - const int64_t output_base = row_base + tile * kHierarchicalTileBlocks; - if (threadIdx.x < kHierarchicalTileBlocks) { + if (threadIdx.x < kHierarchicalTopKTileBlocks) { tile_keys[output_base + threadIdx.x] = keys[threadIdx.x]; } __syncthreads(); @@ -719,10 +761,20 @@ __global__ void QsaMergeTileTopKKernel(const uint64_t* input_keys, uint64_t* out const int right_count = min(min(list_width, kBoundedTopKMax), max(0, key_stride - right_start)); const int output_count = min(left_count + right_count, kBoundedTopKMax); const int64_t row_base = row * key_stride; + const uint64_t* left = input_keys + row_base + left_start; + const uint64_t* right = input_keys + row_base + right_start; + const bool left_empty = left_count == 0 || left[0] == topk::kPaddingSortKey; + const bool right_empty = right_count == 0 || right[0] == topk::kPaddingSortKey; for (int rank = threadIdx.x; rank < output_count; rank += blockDim.x) { - output_keys[row_base + left_start + rank] = - QsaMergeRank(input_keys + row_base + left_start, left_count, - input_keys + row_base + right_start, right_count, rank); + uint64_t key = topk::kPaddingSortKey; + if (left_empty) { + key = rank < right_count ? right[rank] : topk::kPaddingSortKey; + } else if (right_empty) { + key = rank < left_count ? left[rank] : topk::kPaddingSortKey; + } else { + key = QsaMergeRank(left, left_count, right, right_count, rank); + } + output_keys[row_base + left_start + rank] = key; } } } @@ -881,8 +933,18 @@ __global__ void QsaSelectKernel(const float* block_scores, const int32_t* topk_i } } +template +__device__ __forceinline__ uint64_t QsaSelectionKey(const float* scores, int index) { + if constexpr (PackedKeys) { + return reinterpret_cast(scores)[index]; + } else { + return topk::PackStableSortKey(scores[index], index); + } +} + // One block per packed query token. Radix-selects a bounded candidate set, then sorts only that // set. The output aliases the corresponding score row after all scores for that row have been read. +template __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulative_sequence_lengths, const int32_t* past_sequence_lengths, const int64_t* position_ids, const int32_t* present_state_lengths, const int32_t* overflow_flags, @@ -916,7 +978,10 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat if (selected == 0) { continue; } - float* scores = block_scores + static_cast(token) * params.state_capacity; + const int64_t score_stride = PackedKeys + ? 2 * ((static_cast(params.state_capacity) + kHierarchicalTopKTileBlocks - 1) / kHierarchicalTopKTileBlocks) * kHierarchicalTopKTileBlocks + : params.state_capacity; + float* scores = block_scores + static_cast(token) * score_stride; if (threadIdx.x == 0) { prefix = 0; @@ -931,7 +996,7 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat __syncthreads(); const uint32_t current_prefix = prefix; for (int index = threadIdx.x; index < block_count; index += blockDim.x) { - const uint32_t score_key = static_cast(topk::PackStableSortKey(scores[index], 0) >> 32); + const uint32_t score_key = static_cast(QsaSelectionKey(scores, index) >> 32); if (shift == 24 || (score_key >> (shift + 8)) == (current_prefix >> (shift + 8))) { atomicAdd(&temp.histogram[(score_key >> shift) & 0xffu], 1u); } @@ -946,11 +1011,23 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat } else { prefix |= static_cast(bucket) << shift; remaining = rank; + if constexpr (PackedKeys) { + // A fully selected radix bucket needs no finer threshold. + if (rank == count && prefix != 0) { + --prefix; + remaining = 0; + } + } break; } } } __syncthreads(); + if constexpr (PackedKeys) { + if (remaining == 0) { + break; + } + } } if (threadIdx.x == 0) { @@ -959,7 +1036,7 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat __syncthreads(); const uint32_t threshold = prefix; for (int index = threadIdx.x; index < block_count; index += blockDim.x) { - const uint64_t key = topk::PackStableSortKey(scores[index], index); + const uint64_t key = QsaSelectionKey(scores, index); if (static_cast(key >> 32) > threshold) { const int slot = atomicAdd(&gathered, 1); selected_keys[slot] = key; @@ -971,10 +1048,11 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat equal_seen = 0; } __syncthreads(); + // Sorted 32-key tiles still encounter equal scores in ascending original-index order. for (int base = 0; base < block_count && equal_seen < remaining; base += blockDim.x) { const int index = base + static_cast(threadIdx.x); const bool equal = index < block_count && - static_cast(topk::PackStableSortKey(scores[index], 0) >> 32) == threshold; + static_cast(QsaSelectionKey(scores, index) >> 32) == threshold; const unsigned int warp_mask = __ballot_sync(0xffffffffu, equal); const int lane = static_cast(threadIdx.x) % kWarpSize; const int warp = static_cast(threadIdx.x) / kWarpSize; @@ -995,7 +1073,7 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat __syncthreads(); const int equal_rank = equal_seen + equal_warp_offsets[warp] + lane_rank; if (equal && equal_rank < remaining) { - selected_keys[gathered + equal_rank] = topk::PackStableSortKey(scores[index], index); + selected_keys[gathered + equal_rank] = QsaSelectionKey(scores, index); } __syncthreads(); if (threadIdx.x == 0) { @@ -1016,7 +1094,11 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat for (int item = 0; item < kBoundedTopKItemsPerThread; ++item) { const int rank = threadIdx.x * kBoundedTopKItemsPerThread + item; if (rank < selected) { - topk_row[rank] = topk::UnpackStableSortIndex(keys[item]); + if constexpr (PackedKeys) { + reinterpret_cast(scores)[rank] = keys[item]; + } else { + topk_row[rank] = topk::UnpackStableSortIndex(keys[item]); + } } } __syncthreads(); @@ -1406,7 +1488,7 @@ size_t GetQsaPackedWorkspaceFloatCount(const PackedSparseAttentionIndexerParams& const size_t rows = static_cast(params.total_tokens); const size_t query_count = rows * params.num_heads * params.head_size; if (UseHierarchicalQsaTopK(params)) { - const size_t key_count = rows * GetHierarchicalTileCount(params) * kHierarchicalTileBlocks; + const size_t key_count = rows * GetHierarchicalTileCount(params) * kHierarchicalTopKTileBlocks; return query_count + 4 * key_count; } return query_count + rows * static_cast(std::max(params.state_capacity, 1)); @@ -1484,49 +1566,49 @@ Status LaunchQsaPackedSparseAttentionIndexer( const int token_blocks = static_cast(std::min(params.total_tokens, kSaiMaxGridDimX)); if (use_hierarchical_topk) { const int tile_count = GetHierarchicalTileCount(params); - const int key_stride = tile_count * kHierarchicalTileBlocks; + const int key_stride = tile_count * kHierarchicalTopKTileBlocks; const size_t key_count = static_cast(params.total_tokens) * key_stride; uint64_t* merge_input = reinterpret_cast(selection_workspace); uint64_t* merge_output = merge_input + key_count; - if (params.batch_size == 1 && params.total_tokens >= 4) { - const dim3 score_grid(static_cast(std::min(tile_count, 128)), - static_cast((params.total_tokens + 1) / 2)); - QsaPairedScoreTileTopKKernel<<>>( - present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, - present_state_lengths, overflow_flags, merge_input, tile_count, params); + const int64_t tile_work = + static_cast(params.total_tokens) * std::min(tile_count, kHierarchicalScoreMaxBlocksPerRow); + QsaScoreTileTopKKernel<<(std::min(tile_work, kSaiMaxGridDimX)), + kHierarchicalScoreThreads, 0, stream>>>( + present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, + present_state_lengths, overflow_flags, merge_input, tile_count, params); + static const bool use_radix_topk = + ParseEnvironmentVariableWithDefault("ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK", false); + if (use_radix_topk) { + LOGS_DEFAULT(VERBOSE) << "[PackedSparseAttentionIndexer] dispatching QsaPartialTopKKernel"; + QsaPartialTopKKernel<<>>( + selection_workspace, cumulative_sequence_lengths, past_sequence_lengths, position_ids, + present_state_lengths, overflow_flags, params); } else { - const dim3 score_grid(static_cast(std::min(tile_count, 128)), - static_cast(params.total_tokens)); - QsaScoreTileTopKKernel<<>>( - present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, - present_state_lengths, overflow_flags, merge_input, tile_count, params); - } - int compact_stride = key_stride; - int first_list_width = kHierarchicalTileBlocks; - if (params.total_tokens >= 4) { - const int chunks = (key_stride + 2047) / 2048; - compact_stride = chunks * kBoundedTopKMax; - QsaCompactTileTopKKernel<<(chunks), static_cast(params.total_tokens)), + LOGS_DEFAULT(VERBOSE) << "[PackedSparseAttentionIndexer] dispatching QsaMergeTileTopKKernel"; + for (int list_width = kHierarchicalTopKTileBlocks; list_width < key_stride; list_width *= 2) { + const int pairs_per_row = (key_stride + 2 * list_width - 1) / (2 * list_width); + const int64_t merge_work = static_cast(params.total_tokens) * pairs_per_row; + QsaMergeTileTopKKernel<<(std::min(merge_work, kSaiMaxGridDimX)), kHierarchicalMergeThreads, 0, stream>>>( - merge_input, merge_output, key_stride, compact_stride); - std::swap(merge_input, merge_output); - first_list_width = kBoundedTopKMax; - } - for (int list_width = first_list_width; list_width < compact_stride; list_width *= 2) { - const int pairs_per_row = (compact_stride + 2 * list_width - 1) / (2 * list_width); - const int64_t merge_work = static_cast(params.total_tokens) * pairs_per_row; - QsaMergeTileTopKKernel<<(std::min(merge_work, kSaiMaxGridDimX)), - kHierarchicalMergeThreads, 0, stream>>>( - merge_input, merge_output, compact_stride, list_width, params.total_tokens); - std::swap(merge_input, merge_output); + merge_input, merge_output, key_stride, list_width, params.total_tokens); + std::swap(merge_input, merge_output); + } } QsaEmitHierarchicalTopKKernel<<>>( merge_input, cumulative_sequence_lengths, past_sequence_lengths, position_ids, present_state_lengths, - overflow_flags, compact_stride, selected_indices, selected_counts, params); + overflow_flags, key_stride, selected_indices, selected_counts, params); } else if (use_ragged_prefill) { - QsaRaggedPrefillScoreKernel<<>>( - present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, - present_state_lengths, overflow_flags, block_scores, params); + static const bool warp_heads = + ParseEnvironmentVariableWithDefault("ORT_PACKED_SPARSE_INDEXER_WARP_HEADS", false); + if (warp_heads) { + QsaRaggedPrefillScoreKernel<<>>( + present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, + present_state_lengths, overflow_flags, block_scores, params); + } else { + QsaRaggedPrefillScoreKernel<<>>( + present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, + present_state_lengths, overflow_flags, block_scores, params); + } } else if (params.state_capacity > 0) { const int64_t score_work = static_cast(params.total_tokens) * params.state_capacity; const int score_blocks = static_cast(std::min(score_work, kSaiMaxGridDimX)); @@ -1539,7 +1621,7 @@ Status LaunchQsaPackedSparseAttentionIndexer( !use_hierarchical_topk && params.block_topk > kSaiFastTopKMax && params.block_topk <= kBoundedTopKMax; int32_t* topk_indices = use_bounded_topk ? reinterpret_cast(block_scores) : nullptr; if (use_bounded_topk) { - QsaPartialTopKKernel<<>>( + QsaPartialTopKKernel<<>>( block_scores, cumulative_sequence_lengths, past_sequence_lengths, position_ids, present_state_lengths, overflow_flags, params); } diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc index c47c029a58168..d21b2ef93667e 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc @@ -5,6 +5,7 @@ #include #include #include +#include #include #include @@ -443,6 +444,262 @@ TEST(CUDA_EP_Unittest, QMoENvfp4GemvNativeScale) { } #endif +template +void TestAllFp8Codes() { + std::array codes{}; + for (int i = 0; i < 256; ++i) { + codes[i] = static_cast(i); + } + const float scale = 1.0f; + CudaBuffer weights(codes.size()); + CudaBuffer scales(sizeof(scale)); + CudaBuffer output(codes.size() * sizeof(T)); + weights.Upload(codes.data()); + scales.Upload(&scale); + qmoe::LaunchQMoEDequantizeFp8Weights(weights.As(), scales.As(), output.As(), + 1, 1, 256, nullptr); + std::array actual{}; + output.Download(actual.data()); + for (int code = 0; code < 256; ++code) { + SCOPED_TRACE(code); + const float value = static_cast(actual[code]); + if ((code & 127) == 127) { + EXPECT_TRUE(std::isnan(value)); + continue; + } + const int exponent = (code >> 3) & 15; + const int mantissa = code & 7; + const float magnitude = exponent == 0 ? std::ldexp(static_cast(mantissa), -9) + : std::ldexp(1.0f + static_cast(mantissa) / 8.0f, exponent - 7); + const float expected = code & 128 ? -magnitude : magnitude; + EXPECT_EQ(value, expected); + EXPECT_EQ(std::signbit(value), std::signbit(expected)); + } +} + +TEST(CUDA_EP_Unittest, QMoEFp8AllCodesPreserveFiniteValuesSignedZerosAndNaNs) { + TestAllFp8Codes(); + TestAllFp8Codes<__nv_bfloat16>(); +} + +template +void TestFp8ProjectionScaleBlocks(int fusion, int block_size, int k, bool graph, bool gemm) { + constexpr int n = 6; + constexpr int rows = 3; + constexpr int num_experts = 2; + const bool split = fusion == 1; + const int source_n = split ? n / 2 : n; + const int scale_n = (source_n + block_size - 1) / block_size; + const int scale_k = (k + block_size - 1) / block_size; + const std::array experts{0, 0, 1}; + const std::array row_map{5, 0, 4}; + const std::array expert_offsets{0, 2, 3}; + std::vector weights(num_experts * source_n * k); + std::vector up_weights(weights.size()); + std::vector scales(num_experts * scale_n * scale_k); + std::vector up_scales(scales.size()); + std::vector input(rows * k); + std::vector output(rows * n); + for (size_t i = 0; i < weights.size(); ++i) { + weights[i] = static_cast(i % 3 == 0 ? 0xb8 : (i % 3 == 1 ? 0x38 : 0x40)); + up_weights[i] = 0x30; + } + for (size_t i = 0; i < scales.size(); ++i) { + scales[i] = ScaleT(static_cast(i % 5 + 1) / 32.0f); + up_scales[i] = ScaleT(static_cast(i % 5 + 2) / 32.0f); + } + for (size_t i = 0; i < input.size(); ++i) { + input[i] = T(static_cast(static_cast(i % 7) - 3) / 16.0f); + } + CudaBuffer d_weights(weights.size()), d_up_weights(up_weights.size()); + CudaBuffer d_scales(scales.size() * sizeof(ScaleT)), d_up_scales(up_scales.size() * sizeof(ScaleT)); + CudaBuffer d_input(input.size() * sizeof(T)), d_output(output.size() * sizeof(OutputT)); + CudaBuffer d_experts(sizeof(experts)), d_rows(sizeof(row_map)), d_offsets(sizeof(expert_offsets)); + CudaBuffer d_tiles((num_experts + 1) * sizeof(int)); + d_weights.Upload(weights.data()); + d_up_weights.Upload(up_weights.data()); + d_scales.Upload(scales.data()); + d_up_scales.Upload(up_scales.data()); + d_input.Upload(input.data()); + d_experts.Upload(experts.data()); + d_rows.Upload(row_map.data()); + d_offsets.Upload(expert_offsets.data()); + qmoe::QMoEFp8ProjectionParams params; + params.weights = d_weights.As(); + params.scales = d_scales.data; + params.scale_type = std::is_same_v ? 0 : (std::is_same_v ? 1 : 2); + params.up_scale_type = params.scale_type; + params.up_weights = split ? d_up_weights.As() : nullptr; + params.up_scales = split ? d_up_scales.data : nullptr; + params.experts = d_experts.As(); + params.row_to_unpermuted = d_rows.As(); + params.expert_offsets = d_offsets.As(); + params.num_experts = num_experts; + params.num_rows = rows; + params.expanded_rows = rows; + params.n = n; + params.k = k; + params.block_size = block_size; + params.fusion = fusion; + if (gemm) { + qmoe::LaunchQMoEFp8ExpertTiles(params.expert_offsets, d_tiles.As(), num_experts, nullptr); + params.tile_offsets = d_tiles.As(); + CUDA_CALL_THROW(cudaDeviceSynchronize()); + } + cudaStream_t stream; + CUDA_CALL_THROW(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + if (graph) { + CUDA_CALL_THROW(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal)); + } + qmoe::LaunchQMoEFp8Projection(params, d_input.As(), d_output.As(), stream); + if (graph) { + cudaGraph_t captured; + cudaGraphExec_t executable; + CUDA_CALL_THROW(cudaStreamEndCapture(stream, &captured)); + CUDA_CALL_THROW(cudaGraphInstantiate(&executable, captured, nullptr, nullptr, 0)); + for (int replay = 0; replay < 3; ++replay) { + CUDA_CALL_THROW(cudaGraphLaunch(executable, stream)); + } + CUDA_CALL_THROW(cudaStreamSynchronize(stream)); + CUDA_CALL_THROW(cudaGraphExecDestroy(executable)); + CUDA_CALL_THROW(cudaGraphDestroy(captured)); + } + CUDA_CALL_THROW(cudaStreamSynchronize(stream)); + CUDA_CALL_THROW(cudaStreamDestroy(stream)); + d_output.Download(output.data()); + for (int row = 0; row < rows; ++row) { + for (int col = 0; col < n; ++col) { + const bool up = split && col % 2; + const int source_col = split ? col / 2 : (fusion == 2 ? col / 2 + (col % 2) * (n / 2) : col); + float expected = 0.0f; + for (int inner = 0; inner < k; ++inner) { + const int weight_index = (experts[row] * source_n + source_col) * k + inner; + const int scale_index = (experts[row] * scale_n + source_col / block_size) * scale_k + + inner / block_size; + const uint8_t code = up ? up_weights[weight_index] : weights[weight_index]; + const float value = code == 0xb8 ? -1.0f : (code == 0x38 ? 1.0f : (code == 0x40 ? 2.0f : 0.5f)); + const T weight(value * static_cast(up ? up_scales[scale_index] : scales[scale_index])); + expected += static_cast(input[(row_map[row] % rows) * k + inner]) * static_cast(weight); + } + ASSERT_EQ(static_cast(output[row * n + col]), static_cast(OutputT(expected))) + << "row=" << row << " col=" << col; + } + } +} + +TEST(CUDA_EP_Unittest, QMoEFp8ProjectionPreservesScaleBlocksLayoutsAndGraphReplay) { + int device = 0; + cudaDeviceProp properties{}; + ASSERT_EQ(cudaGetDevice(&device), cudaSuccess); + ASSERT_EQ(cudaGetDeviceProperties(&properties, device), cudaSuccess); + for (int fusion : {0, 1, 2}) { + for (int block_size : {64, 128}) { + for (int k : {127, 128, 257}) { + for (bool graph : {false, true}) { + for (bool gemm : {false, true}) { + if (gemm && properties.major < 8) { + continue; + } + SCOPED_TRACE(::testing::Message() << "fusion=" << fusion << " block=" << block_size + << " k=" << k << " graph=" << graph << " gemm=" << gemm); + TestFp8ProjectionScaleBlocks(fusion, block_size, k, graph, gemm); + TestFp8ProjectionScaleBlocks(fusion, block_size, k, graph, gemm); + TestFp8ProjectionScaleBlocks(fusion, block_size, k, graph, gemm); + TestFp8ProjectionScaleBlocks<__nv_bfloat16, float>(fusion, block_size, k, graph, gemm); + TestFp8ProjectionScaleBlocks<__nv_bfloat16, half>(fusion, block_size, k, graph, gemm); + TestFp8ProjectionScaleBlocks<__nv_bfloat16, __nv_bfloat16>(fusion, block_size, k, graph, gemm); + TestFp8ProjectionScaleBlocks(fusion, block_size, k, graph, gemm); + TestFp8ProjectionScaleBlocks<__nv_bfloat16, float, float>(fusion, block_size, k, graph, gemm); + } + } + } + } + } +} + +template +void TestNvfp4RowMajorDequantization(int k, bool compact, bool with_bias) { + constexpr int source_experts = 5; + constexpr int n = 3; + const std::array selected{4, 1, -1}; + const int capacity = compact ? static_cast(selected.size()) : source_experts; + constexpr std::array magnitudes{0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f}; + constexpr std::array scale_codes{0x00, 0x01, 0x20, 0x38, 0x41, 0x7e, 0x80, 0xb8}; + constexpr std::array scale_values{0.0f, 1.0f / 512.0f, 0.125f, 1.0f, 2.25f, 448.0f, -0.0f, -1.0f}; + std::vector weights(source_experts * n * k / 2); + std::vector scales(source_experts * n * k / 16); + std::vector globals(source_experts); + std::vector bias(source_experts * n); + for (int expert = 0; expert < source_experts; ++expert) { + globals[expert] = 0.37f * (expert + 1); + for (int row = 0; row < n; ++row) { + const int source_row = expert * n + row; + bias[source_row] = T(static_cast(source_row)); + for (int col = 0; col < k; col += 2) { + const int low = (col + row + expert) % 16; + const int high = (low + 1) % 16; + weights[source_row * (k / 2) + col / 2] = static_cast(low | (high << 4)); + } + for (int group = 0; group < k / 16; ++group) { + scales[source_row * (k / 16) + group] = scale_codes[(group + source_row) % scale_codes.size()]; + } + } + } + std::vector output(capacity * n * k, T(-123.0f)); + std::vector output_bias(capacity * n, T(-123.0f)); + CudaBuffer device_weights(weights.size()); + CudaBuffer device_scales(scales.size()); + CudaBuffer device_globals(globals.size() * sizeof(float)); + CudaBuffer device_bias(bias.size() * sizeof(T)); + CudaBuffer device_output(output.size() * sizeof(T)); + CudaBuffer device_output_bias(output_bias.size() * sizeof(T)); + CudaBuffer device_selected(selected.size() * sizeof(int)); + device_weights.Upload(weights.data()); + device_scales.Upload(scales.data()); + device_globals.Upload(globals.data()); + device_bias.Upload(bias.data()); + device_output.Upload(output.data()); + device_output_bias.Upload(output_bias.data()); + device_selected.Upload(selected.data()); + qmoe::LaunchQMoEDequantizeNvfp4Weights( + device_weights.As(), device_scales.As(), device_globals.As(), + device_output.As(), capacity, n, k, nullptr, + compact ? device_selected.As() : nullptr, + with_bias ? device_bias.As() : nullptr, + with_bias ? device_output_bias.As() : nullptr, true); + device_output.Download(output.data()); + device_output_bias.Download(output_bias.data()); + for (int slot = 0; slot < capacity; ++slot) { + const int expert = compact ? selected[slot] : slot; + for (int row = 0; row < n; ++row) { + const T expected_bias = with_bias && expert >= 0 ? bias[expert * n + row] : T(-123.0f); + EXPECT_EQ(std::memcmp(&output_bias[slot * n + row], &expected_bias, sizeof(T)), 0); + for (int col = 0; col < k; ++col) { + T expected(-123.0f); + if (expert >= 0) { + const int code = (col + row + expert) % 16; + const float value = code < 8 ? magnitudes[code] : -magnitudes[code - 8]; + expected = T(value * scale_values[(col / 16 + expert * n + row) % scale_values.size()] * globals[expert]); + } + ASSERT_EQ(std::memcmp(&output[(slot * n + row) * k + col], &expected, sizeof(T)), 0) + << "slot=" << slot << " row=" << row << " col=" << col; + } + } + } +} + +TEST(CUDA_EP_Unittest, QMoENvfp4RowMajorVectorDequantizationPreservesValuesAndBias) { + for (int k : {16, 48, 1280, 2560}) { + for (bool compact : {false, true}) { + for (bool with_bias : {false, true}) { + SCOPED_TRACE(::testing::Message() << "k=" << k << " compact=" << compact << " bias=" << with_bias); + TestNvfp4RowMajorDequantization(k, compact, with_bias); + TestNvfp4RowMajorDequantization<__nv_bfloat16>(k, compact, with_bias); + } + } + } +} + } // namespace } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index 3ccf50716367c..84e7a2d5416a2 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc @@ -16,24 +16,33 @@ #include #include +#include +#include #include #include #include #include #include #include +#include #include #include #include #include "gtest/gtest.h" +#if defined(GTEST_HAS_ABSL) && !defined(GTEST_NO_ABSL_FLAGS) +#include "absl/flags/reflection.h" +#endif #include "contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h" +#include "core/common/logging/logging.h" #include "core/graph/constants.h" #include "core/graph/model.h" +#include "core/platform/env.h" #include "test/common/tensor_op_test_utils.h" #include "test/providers/provider_test_utils.h" #include "test/test_environment.h" +#include "test/util/include/scoped_env_vars.h" #include "test/unittest_util/graph_transform_test_builder.h" #include "test/util/include/asserts.h" #include "test/util/include/default_providers.h" @@ -1072,6 +1081,86 @@ void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedP } } +#if GTEST_HAS_DEATH_TEST +QsaPackedProblem MakeRadixTopKProblem(int visible_blocks, int block_topk, int token_count = 1) { + QsaPackedProblem problem; + problem.batch_size = 1; + problem.compress_ratio = 4; + problem.cumulative_sequence_lengths = {0, token_count}; + problem.past_sequence_lengths = {visible_blocks * problem.compress_ratio}; + problem.head_size = 128; + problem.num_heads = 4; + problem.rotary_width = 64; + problem.token_budget = block_topk * problem.compress_ratio; + problem.state_capacity = 2051; + problem = MakeQsaPackedProblem(std::move(problem)); + return problem; +} + +void SetRadixTopKScoreGroups(QsaPackedProblem& problem, int high_score_blocks, int tied_score_blocks, + int high_score_index = -1) { + std::fill(problem.query.begin(), problem.query.end(), 1.0f); + std::fill(problem.query_norm_weight.begin(), problem.query_norm_weight.end(), 1.0f); + std::fill(problem.cos_cache.begin(), problem.cos_cache.end(), 1.0f); + std::fill(problem.sin_cache.begin(), problem.sin_cache.end(), 0.0f); + for (int batch = 0; batch < problem.batch_size; ++batch) { + const int visible_blocks = problem.past_sequence_lengths[batch] / problem.compress_ratio; + for (int block = 0; block < visible_blocks; ++block) { + const float value = + block == high_score_index || block < high_score_blocks ? 3.0f : block < tied_score_blocks ? 2.0f + : 1.0f; + const size_t offset = (static_cast(batch) * problem.state_capacity + block) * problem.head_size; + std::fill_n(problem.past_key_state.begin() + offset, problem.head_size, value); + } + } +} + +void RunRadixTopKCase(const char* name, QsaPackedProblem problem) { + SCOPED_TRACE(name); + ASSERT_LE(problem.TotalTokens(), 64); + ASSERT_GE(problem.state_capacity, 2048); + ASSERT_EQ(problem.head_size, 128); + ASSERT_EQ(problem.compress_ratio, 4); + ASSERT_EQ(problem.num_heads, 4); + ASSERT_GT(problem.token_budget / problem.compress_ratio, 32); + ASSERT_LE(problem.token_budget / problem.compress_ratio, 512); + RunQsaPackedTest(1.0e-5f, std::move(problem), ProviderKind::Cuda); +} + +void RunRadixTopKRegressionCases() { + const std::string flag = Env::Default().GetEnvironmentVar("ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK"); + ASSERT_TRUE(flag == "0" || flag == "1"); + + auto random_scores = MakeRadixTopKProblem(2051, 35); + std::mt19937 generator(33393); + std::uniform_real_distribution distribution(-0.5f, 0.5f); + for (float& value : random_scores.query) value = distribution(generator); + for (float& value : random_scores.past_key_state) value = distribution(generator); + RunRadixTopKCase("random_scores", std::move(random_scores)); + + auto threshold_ties = MakeRadixTopKProblem(2051, 35); + SetRadixTopKScoreGroups(threshold_ties, 33, 71); + RunRadixTopKCase("threshold_ties", std::move(threshold_ties)); + + auto empty_row = MakeRadixTopKProblem(2051, 35, 2); + empty_row.batch_size = 2; + empty_row.cumulative_sequence_lengths = {0, 1, 2}; + empty_row.past_sequence_lengths = {0, 8204}; + empty_row.past_state_lengths.clear(); + empty_row = MakeQsaPackedProblem(std::move(empty_row)); + SetRadixTopKScoreGroups(empty_row, 33, 71); + RunRadixTopKCase("empty_row", std::move(empty_row)); + + auto partial_tile = MakeRadixTopKProblem(2051, 35); + SetRadixTopKScoreGroups(partial_tile, 34, 34, 2050); + RunRadixTopKCase("partial_tile", std::move(partial_tile)); + + auto fully_selected_bucket = MakeRadixTopKProblem(64, 64); + SetRadixTopKScoreGroups(fully_selected_bucket, 64, 64); + RunRadixTopKCase("fully_selected_bucket", std::move(fully_selected_bucket)); +} +#endif + } // namespace TEST(PackedSparseAttentionIndexerTest, QsaFloat) { RunQsaPackedTest(1.0e-5f); } @@ -1255,6 +1344,57 @@ TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalTopK) { RunQsaPackedTest(1.0e-5f, std::move(problem)); } +TEST(PackedSparseAttentionIndexerCudaTest, QsaRadixTopKReferenceParity) { +#if GTEST_HAS_DEATH_TEST + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA execution provider is not available"; + } + +#if defined(GTEST_HAS_ABSL) && !defined(GTEST_NO_ABSL_FLAGS) + auto* death_test_style_flag = absl::FindCommandLineFlag("gtest_death_test_style"); + ASSERT_NE(death_test_style_flag, nullptr); + const std::string previous_style = death_test_style_flag->CurrentValue(); + std::string flag_error; + ASSERT_TRUE(death_test_style_flag->ParseFrom("threadsafe", &flag_error)) << flag_error; +#else + const auto previous_style = GTEST_FLAG_GET(death_test_style); + GTEST_FLAG_SET(death_test_style, "threadsafe"); +#endif + for (const char* radix_flag : {"0", "1"}) { + const char* dispatch_marker = + radix_flag[0] == '1' ? "QsaPartialTopKKernel" : "QsaMergeTileTopKKernel"; + // Plugin loggers cache their severity during registration in the re-executed test main. + ScopedEnvironmentVariables env_vars{EnvVarMap{{"ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK", radix_flag}, + {"ORT_UNIT_TEST_MAIN_LOG_LEVEL", "0"}}}; + ASSERT_EXIT( + { + logging::LoggingManager::SetDefaultLoggerSeverity(logging::Severity::kVERBOSE); + logging::LoggingManager::SetDefaultLoggerVerbosity(1); + RunRadixTopKRegressionCases(); + const auto* result = ::testing::UnitTest::GetInstance()->current_test_info()->result(); + for (int index = 0; index < result->total_part_count(); ++index) { + const auto& part = result->GetTestPartResult(index); + if (part.failed()) { + std::fprintf(stderr, "%s:%d: %s\n", part.file_name() ? part.file_name() : "unknown", + part.line_number(), part.message()); + } + } + std::fflush(stderr); + std::_Exit(::testing::Test::HasFailure() ? EXIT_FAILURE : EXIT_SUCCESS); + }, + ::testing::ExitedWithCode(EXIT_SUCCESS), dispatch_marker); + } +#if defined(GTEST_HAS_ABSL) && !defined(GTEST_NO_ABSL_FLAGS) + flag_error.clear(); + EXPECT_TRUE(death_test_style_flag->ParseFrom(previous_style, &flag_error)) << flag_error; +#else + GTEST_FLAG_SET(death_test_style, previous_style); +#endif +#else + GTEST_SKIP() << "GoogleTest death tests are not available"; +#endif +} + TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalPackedVerification) { QsaPackedProblem problem; problem.batch_size = 3; @@ -1341,6 +1481,52 @@ TEST(PackedSparseAttentionIndexerTest, QsaMalformedStateLengthsAreClamped) { RunQsaPackedTest(1.0e-5f, MakeQsaPackedProblem(std::move(problem))); } +TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalTileBoundariesPreserveStableTies) { + for (int live_blocks : {0, 31, 32, 33, 511, 512, 513, 1023, 1024, 1025, 2047, 2048, 2049, 2051}) { + SCOPED_TRACE(live_blocks); + QsaPackedProblem problem; + problem.batch_size = 2; + problem.cumulative_sequence_lengths = {0, 1, 2}; + problem.past_sequence_lengths = {live_blocks * 4, std::max(0, live_blocks - 1) * 4}; + problem.head_size = 128; + problem.num_heads = 4; + problem.rotary_width = 64; + problem.compress_ratio = 4; + problem.token_budget = 2048; + problem.state_capacity = 2051; + problem.max_position = std::max(32, live_blocks * 4 + 2); + problem = MakeQsaPackedProblem(std::move(problem)); + std::fill(problem.query.begin(), problem.query.end(), 0.0f); + std::fill(problem.cos_cache.begin(), problem.cos_cache.end(), 1.0f); + std::fill(problem.sin_cache.begin(), problem.sin_cache.end(), 0.0f); + RunQsaPackedTest(2.0e-3f, problem); + RunQsaPackedTest(2.0e-2f, std::move(problem)); + } +} + +TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalPackedRowsWithLargeCapacity) { + for (int tokens : {1, 2, 3, 8, 64}) { + SCOPED_TRACE(tokens); + QsaPackedProblem problem; + problem.batch_size = 1; + problem.cumulative_sequence_lengths = {0, tokens}; + problem.past_sequence_lengths = {4097 * 4 + 3}; + problem.head_size = 128; + problem.num_heads = 4; + problem.rotary_width = 64; + problem.compress_ratio = 4; + problem.token_budget = 2048; + problem.state_capacity = 65536; + problem.max_position = problem.past_sequence_lengths[0] + tokens + 1; + problem = MakeQsaPackedProblem(std::move(problem)); + std::fill(problem.query.begin(), problem.query.end(), 0.0f); + std::fill(problem.cos_cache.begin(), problem.cos_cache.end(), 1.0f); + std::fill(problem.sin_cache.begin(), problem.sin_cache.end(), 0.0f); + RunQsaPackedTest(2.0e-3f, problem, ProviderKind::Cuda); + RunQsaPackedTest(2.0e-2f, std::move(problem), ProviderKind::Cuda); + } +} + #ifdef USE_WEBGPU TEST(PackedSparseAttentionIndexerWebGpuTest, QsaFloat) { RunQsaPackedTest(1.0e-5f, MakeQsaPackedProblem(), ProviderKind::WebGpu); diff --git a/onnxruntime/test/contrib_ops/sparse_paged_attention_op_test.cc b/onnxruntime/test/contrib_ops/sparse_paged_attention_op_test.cc index 683250ef4bc25..00a5aa975acbc 100644 --- a/onnxruntime/test/contrib_ops/sparse_paged_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/sparse_paged_attention_op_test.cc @@ -268,11 +268,36 @@ TEST(SparsePagedAttention, Cuda_SuppressedCurrentWriteReadsMappedCache) { RunCuda(tester); } -TEST(SparsePagedAttention, Cuda_Int8CurrentTokenUsesQuantizedCacheValue) { +TEST(SparsePagedAttention, Cuda_TiledSuppressedCurrentWriteReadsMappedCache) { if (DefaultCudaExecutionProvider() == nullptr) { GTEST_SKIP() << "CUDA EP not available."; } + constexpr int heads = 4; + constexpr int channels = 256; + OpTester tester("SparsePagedAttention", 1, kMSDomain); + tester.AddAttribute("num_heads", heads); + tester.AddAttribute("kv_num_heads", 1); + tester.AddInput("query", {1, heads * channels}, HalfVector(0.0f, heads * channels)); + tester.AddInput("key", {1, channels}, HalfVector(0.0f, channels)); + tester.AddInput("value", {1, channels}, HalfVector(1.0f, channels)); + tester.AddInput("key_cache", {1, kBlockSize, 1, channels}, + HalfVector(0.0f, kBlockSize * channels)); + tester.AddInput("value_cache", {1, kBlockSize, 1, channels}, + HalfVector(7.0f, kBlockSize * channels)); + tester.AddInput("cumulative_sequence_length", {2}, {0, 1}); + tester.AddInput("past_seqlens", {1}, {0}); + tester.AddInput("block_table", {1, 1}, {0}); + tester.AddInput("slot_mapping", {1}, {-1}); + tester.AddInput("selected_indices", {1, 1}, {0}); + tester.AddInput("selected_counts", {1}, {1}); + tester.AddOutput("output", {1, heads * channels}, HalfVector(7.0f, heads * channels)); + RunCuda(tester); +} +TEST(SparsePagedAttention, Cuda_Int8CurrentTokenUsesQuantizedCacheValue) { + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } auto query = HalfVector(0.0f); query[0] = MLFloat16(1.0f); auto key = HalfVector(0.0f); @@ -660,6 +685,214 @@ TEST(SparsePagedAttention, Cuda_MergesMultipleSelectedMainSplits) { RunCuda(tester); } +void RunSparsePagedAttentionMultiTileReference(int token_count, int num_heads = 8, int past_length = 257, + int reference_head_size = 0) { + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + const int selected_count = past_length + 2; + const int num_blocks = (past_length + token_count + kBlockSize - 1) / kBlockSize; + for (int head_size : {8, 128, 192, 256, 384}) { + if (reference_head_size != 0 && head_size != reference_head_size) continue; + SCOPED_TRACE(head_size); + const float scale = 1.0f / std::sqrt(static_cast(head_size)); + auto query = HalfVector(0.0f, token_count * num_heads * head_size); + auto key_cache = HalfVector(0.0f, num_blocks * kBlockSize * head_size); + auto value_cache = key_cache; + for (size_t i = 0; i < query.size(); ++i) { + query[i] = MLFloat16(static_cast(static_cast(i % 23) - 11) / 16.0f); + } + for (size_t i = 0; i < key_cache.size(); ++i) { + key_cache[i] = MLFloat16(static_cast(static_cast(i % 31) - 15) / 16.0f); + value_cache[i] = MLFloat16(static_cast(static_cast(i % 19) - 9) / 8.0f); + } + std::vector selected(token_count * selected_count); + for (int token = 0; token < token_count; ++token) { + std::iota(selected.begin() + token * selected_count, + selected.begin() + token * selected_count + past_length, 0); + selected[token * selected_count + past_length] = -1; + selected[token * selected_count + past_length + 1] = past_length + token_count; + } + std::vector blocks(num_blocks); + std::iota(blocks.begin(), blocks.end(), 0); + std::vector slots(token_count); + std::iota(slots.begin(), slots.end(), past_length); + auto expected = HalfVector(0.0f, static_cast(query.size())); + for (int token_head = 0; token_head < token_count * num_heads; ++token_head) { + std::vector logits(past_length); + float maximum = -std::numeric_limits::infinity(); + for (int position = 0; position < past_length; ++position) { + float dot = 0.0f; + for (int c = 0; c < head_size; ++c) { + dot += query[token_head * head_size + c].ToFloat() * key_cache[position * head_size + c].ToFloat(); + } + logits[position] = dot * scale; + maximum = std::max(maximum, logits[position]); + } + float sum = 0.0f; + for (float& logit : logits) { + logit = std::exp(logit - maximum); + sum += logit; + } + for (int c = 0; c < head_size; ++c) { + float value = 0.0f; + for (int position = 0; position < past_length; ++position) { + value += logits[position] * value_cache[position * head_size + c].ToFloat(); + } + expected[token_head * head_size + c] = MLFloat16(value / sum); + } + } + OpTester tester("SparsePagedAttention", 1, kMSDomain); + tester.AddAttribute("num_heads", num_heads); + tester.AddAttribute("kv_num_heads", 1); + tester.AddInput("query", {token_count, num_heads * head_size}, query); + tester.AddInput("key", {token_count, head_size}, HalfVector(0.0f, token_count * head_size)); + tester.AddInput("value", {token_count, head_size}, HalfVector(0.0f, token_count * head_size)); + tester.AddInput("key_cache", {num_blocks, kBlockSize, 1, head_size}, key_cache); + tester.AddInput("value_cache", {num_blocks, kBlockSize, 1, head_size}, value_cache); + tester.AddInput("cumulative_sequence_length", {2}, {0, token_count}); + tester.AddInput("past_seqlens", {1}, {past_length}); + tester.AddInput("block_table", {1, num_blocks}, blocks); + tester.AddInput("slot_mapping", {token_count}, slots); + tester.AddInput("selected_indices", {token_count, selected_count}, selected); + tester.AddInput("selected_counts", {token_count}, std::vector(token_count, selected_count)); + tester.AddOutput("output", {token_count, num_heads * head_size}, expected, false, 0.002f, 0.002f); + RunCuda(tester); + } +} + +TEST(SparsePagedAttention, Cuda_MultiTileSoftmaxWithLargeHeadsMatchesReference) { + RunSparsePagedAttentionMultiTileReference(40); +} + +TEST(SparsePagedAttention, Cuda_MultiSplitDecodeWithLargeHeadsMatchesReference) { + RunSparsePagedAttentionMultiTileReference(1); +} + +TEST(SparsePagedAttention, Cuda_ShortVerificationQueriesWithLargeSelectionsMatchReference) { + for (int token_count : {2, 4, 8}) { + SCOPED_TRACE(token_count); + RunSparsePagedAttentionMultiTileReference(token_count, 24, 2048); + } +} + +TEST(SparsePagedAttention, Cuda_TwelveHeadGroupedPrefillMatchesReference) { + for (int past_length : {254, 255, 256, 257}) { + SCOPED_TRACE(past_length); + RunSparsePagedAttentionMultiTileReference(129, 24, past_length, 256); + } +} + +TEST(SparsePagedAttention, Cuda_GqaTiledPrefillHandlesMixedSourcesAndMasks) { + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + constexpr int token_count = 40; + constexpr int num_heads = 8; + constexpr int kv_heads = 2; + constexpr int head_size = 256; + constexpr int past_length = 257; + constexpr int main_length = past_length + token_count; + constexpr int selected_width = main_length + 2; + constexpr int num_blocks = (main_length + kBlockSize - 1) / kBlockSize; + constexpr float scale = 0.0625f; + constexpr float softcap = 1.7f; + auto query = HalfVector(0.0f, token_count * num_heads * head_size); + auto key = HalfVector(0.0f, token_count * kv_heads * head_size); + auto value = key; + auto key_cache = HalfVector(0.0f, num_blocks * kBlockSize * kv_heads * head_size); + auto value_cache = key_cache; + auto sink = HalfVector(0.3f, num_heads); + std::vector blocks(num_blocks); + std::iota(blocks.begin(), blocks.end(), 0); + blocks[2] = -1; + std::vector slots(token_count); + std::iota(slots.begin(), slots.end(), past_length); + slots[5] = -1; + std::vector selected(token_count * selected_width); + std::vector counts(token_count); + for (int token = 0; token < token_count; ++token) { + std::iota(selected.begin() + token * selected_width, + selected.begin() + (token + 1) * selected_width, 0); + selected[token * selected_width + main_length] = -1; + counts[token] = token == 0 ? 0 : (token % 3 == 0 ? 64 : (token % 3 == 1 ? 65 : selected_width)); + if (token == 1) { + std::fill_n(selected.begin() + token * selected_width, selected_width, -1); + } + for (int head = 0; head < num_heads; ++head) { + std::fill_n(query.begin() + (token * num_heads + head) * head_size, + head_size, MLFloat16(static_cast(head + 1) / 32.0f)); + } + } + for (int position = 0; position < num_blocks * kBlockSize; ++position) { + for (int head = 0; head < kv_heads; ++head) { + const auto key_element = MLFloat16(static_cast((position + head) % 17 - 8) / 32.0f); + const auto value_element = MLFloat16(static_cast((position + 3 * head) % 23 - 11) / 16.0f); + std::fill_n(key_cache.begin() + (position * kv_heads + head) * head_size, head_size, key_element); + std::fill_n(value_cache.begin() + (position * kv_heads + head) * head_size, head_size, value_element); + if (position >= past_length && position < main_length) { + const int offset = ((position - past_length) * kv_heads + head) * head_size; + std::fill_n(key.begin() + offset, head_size, key_element); + std::fill_n(value.begin() + offset, head_size, value_element); + } + } + } + for (int head = 0; head < kv_heads; ++head) { + const int cached_offset = ((past_length + 5) * kv_heads + head) * head_size; + const int current_offset = (5 * kv_heads + head) * head_size; + std::fill_n(key_cache.begin() + cached_offset, head_size, MLFloat16(0.75f)); + std::fill_n(value_cache.begin() + cached_offset, head_size, MLFloat16(4.0f)); + std::fill_n(key.begin() + current_offset, head_size, MLFloat16(-0.75f)); + std::fill_n(value.begin() + current_offset, head_size, MLFloat16(-4.0f)); + } + auto expected = HalfVector(0.0f, static_cast(query.size())); + for (int token = 0; token < token_count; ++token) { + for (int head = 0; head < num_heads; ++head) { + const int kv_head = head / (num_heads / kv_heads); + float denominator = std::exp(sink[head].ToFloat()); + float numerator = 0.0f; + for (int index = 0; index < counts[token]; ++index) { + const int position = selected[token * selected_width + index]; + if (position < 0 || position > past_length + token || blocks[position / kBlockSize] < 0) { + continue; + } + const int offset = (position * kv_heads + kv_head) * head_size; + const float dot = head_size * query[(token * num_heads + head) * head_size].ToFloat() * + key_cache[offset].ToFloat() * scale; + const float weight = std::exp(std::tanh(dot / softcap) * softcap); + denominator += weight; + numerator += weight * value_cache[offset].ToFloat(); + } + std::fill_n(expected.begin() + (token * num_heads + head) * head_size, + head_size, MLFloat16(numerator / denominator)); + } + } + OpTester tester("SparsePagedAttention", 1, kMSDomain); + tester.AddAttribute("num_heads", num_heads); + tester.AddAttribute("kv_num_heads", kv_heads); + tester.AddAttribute("scale", scale); + tester.AddAttribute("softcap", softcap); + tester.AddInput("query", {token_count, num_heads * head_size}, query); + tester.AddInput("key", {token_count, kv_heads * head_size}, key); + tester.AddInput("value", {token_count, kv_heads * head_size}, value); + tester.AddInput("key_cache", {num_blocks, kBlockSize, kv_heads, head_size}, key_cache); + tester.AddInput("value_cache", {num_blocks, kBlockSize, kv_heads, head_size}, value_cache); + tester.AddInput("cumulative_sequence_length", {2}, {0, token_count}); + tester.AddInput("past_seqlens", {1}, {past_length}); + tester.AddInput("block_table", {1, num_blocks}, blocks); + tester.AddInput("slot_mapping", {token_count}, slots); + tester.AddInput("selected_indices", {token_count, selected_width}, selected); + tester.AddInput("selected_counts", {token_count}, counts); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddInput("head_sink", {num_heads}, sink); + tester.AddOutput("output", {token_count, num_heads * head_size}, expected, false, 0.002f, 0.002f); + RunCuda(tester); +} + class SparsePagedAttentionInt8Test : public ::testing::TestWithParam> {}; TEST_P(SparsePagedAttentionInt8Test, Cuda_QuantizedMainCache) { diff --git a/onnxruntime/test/python/transformers/test_packed_sparse_indexer_warp_heads.py b/onnxruntime/test/python/transformers/test_packed_sparse_indexer_warp_heads.py new file mode 100644 index 0000000000000..fe04899988559 --- /dev/null +++ b/onnxruntime/test/python/transformers/test_packed_sparse_indexer_warp_heads.py @@ -0,0 +1,133 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import argparse +import os +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest +from onnx import TensorProto, helper + +import onnxruntime as ort + + +def run(reference, artifacts, dtype): + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + cases = 0 + for capacity, live, ties in [ + (2051, 2051, False), + (2051, 513, True), + (65536, 2048, False), + (65536, 0, True), + (2051, 31, False), + (2051, 32, False), + (2051, 33, False), + (2051, 2048, True), + (65536, 32768, False), + (65536, 32767, True), + ]: + rng = np.random.default_rng(481) + batch, tokens, channels = 2, 130, 128 + past = np.array([live * 4, max(0, live - 1) * 4], dtype=np.int32) + feed = { + "query": (np.zeros((tokens, 512)) if ties else rng.normal(size=(tokens, 512))).astype(dtype), + "key": rng.normal(size=(tokens, channels)).astype(dtype), + "qn": np.ones(channels, dtype=dtype), + "kn": np.ones(channels, dtype=dtype), + "cos": np.ones((32, 64), dtype=dtype), + "sin": np.zeros((32, 64), dtype=dtype), + "cumulative": np.array([0, 65, 130], dtype=np.int32), + "past": past, + "positions": np.array([0, 0], dtype=np.int64), + "state": rng.normal(size=(batch, capacity, channels)).astype(dtype), + "buffer": np.zeros((batch, 7, channels), dtype=dtype), + "lengths": np.array([[live, 0], [max(0, live - 1), 0]], dtype=np.int32), + } + # Position IDs bound RoPE accesses but also set visibility; use identity caches + # large enough to represent the actual causal positions for both requests. + feed.pop("positions") + feed["cos"] = np.ones((max(int(past.max()) + 66, 32), 64), dtype=dtype) + feed["sin"] = np.zeros_like(feed["cos"]) + names = [ + "query", + "key", + "qn", + "kn", + "cos", + "sin", + "cumulative", + "past", + "", + "", + "", + "", + "state", + "buffer", + "", + "lengths", + ] + tensor_type = TensorProto.FLOAT16 if dtype == np.float16 else TensorProto.FLOAT + inputs = [ + helper.make_tensor_value_info(n, tensor_type if a.dtype == dtype else TensorProto.INT32, a.shape) + for n, a in feed.items() + ] + output_names = ["indices", "counts", "next_state", "next_buffer", "", "next_lengths"] + outputs = [ + helper.make_tensor_value_info("indices", TensorProto.INT32, [tokens, 2051]), + helper.make_tensor_value_info("counts", TensorProto.INT32, [tokens]), + helper.make_tensor_value_info("next_state", tensor_type, [batch, capacity, channels]), + helper.make_tensor_value_info("next_buffer", tensor_type, [batch, 7, channels]), + helper.make_tensor_value_info("next_lengths", TensorProto.INT32, [batch, 2]), + ] + node = helper.make_node( + "PackedSparseAttentionIndexer", + names, + output_names, + domain="com.microsoft", + policy_mode="qsa", + compress_ratio=4, + state_capacity=capacity, + token_budget=2048, + ) + model = helper.make_model( + helper.make_graph([node], "indexer_topk", inputs, outputs), + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)], + ) + model.ir_version = 10 + session = ort.InferenceSession(model.SerializeToString(), options, providers=["CUDAExecutionProvider"]) + assert session.get_providers()[0] == "CUDAExecutionProvider" + actual = session.run(None, feed) + path = artifacts / f"indexer_topk_reference_{cases}.npz" + if reference: + np.savez(path, *actual) + else: + with np.load(path) as expected: + for i, output in enumerate(actual): + np.testing.assert_array_equal(output, expected[f"arr_{i}"]) + print(f"case={cases} capacity={capacity} live={live} ties={ties} counts={actual[1].tolist()}", flush=True) + cases += 1 + print(f"{cases} cases: {'saved baseline' if reference else 'all outputs bitwise equal to baseline'}") + + +@pytest.mark.skipif("CUDAExecutionProvider" not in ort.get_available_providers(), reason="CUDA required") +@pytest.mark.parametrize("dtype", ["float16", "float32"]) +def test_warp_heads_matches_existing_prefill(tmp_path, dtype): + for flag in ("0", "1"): + environment = dict(os.environ, ORT_PACKED_SPARSE_INDEXER_WARP_HEADS=flag) + command = [sys.executable, str(Path(__file__).resolve()), "--artifacts", str(tmp_path), "--dtype", dtype] + if flag == "0": + command.append("--reference") + subprocess.run(command, env=environment, check=True) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--reference", action="store_true") + parser.add_argument("--artifacts", type=Path, required=True) + parser.add_argument("--dtype", choices=["float16", "float32"], required=True) + args = parser.parse_args() + run(args.reference, args.artifacts, np.dtype(args.dtype).type) diff --git a/onnxruntime/test/python/transformers/test_qmoe_int2_fake_int4_activations_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_int2_fake_int4_activations_cuda.py new file mode 100644 index 0000000000000..14ba4714f4323 --- /dev/null +++ b/onnxruntime/test/python/transformers/test_qmoe_int2_fake_int4_activations_cuda.py @@ -0,0 +1,158 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import numpy as np +import pytest +import torch +from onnx import TensorProto, helper, numpy_helper + +import onnxruntime as ort + +pytestmark = pytest.mark.skipif( + "CUDAExecutionProvider" not in ort.get_available_providers() + or not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] < 8, + reason="CUDA SM80+ required", +) + + +def _fake_int4(values): + groups = values.float().reshape(*values.shape[:-1], -1, 64) + scales = groups.abs().amax(-1, keepdim=True) / 7 + safe_scales = torch.where(scales == 0, torch.ones_like(scales), scales) + return (torch.round(groups / safe_scales).clamp(-7, 7) * scales).reshape_as(values).to(values.dtype) + + +def _initializer(name, values): + if values.dtype == torch.uint8: + return numpy_helper.from_array(values.numpy(), name) + tensor = numpy_helper.from_array(values.view(torch.uint16).numpy(), name) + tensor.data_type = TensorProto.BFLOAT16 if values.dtype == torch.bfloat16 else TensorProto.FLOAT16 + return tensor + + +def _case(rows, dtype, fake_activations=True): + experts, hidden, inter = 512, 128, 128 + generator = torch.Generator().manual_seed(42) + codes = [ + torch.randint(0, 4, (experts, inter * 2, hidden), generator=generator, dtype=torch.uint8), + torch.randint(0, 4, (experts, hidden, inter), generator=generator, dtype=torch.uint8), + ] + codes[0][0, :128] = 2 + arrays = {} + matrices = [] + for index, code in enumerate(codes, 1): + packed = code[..., 0::4] | (code[..., 1::4] << 2) | (code[..., 2::4] << 4) | (code[..., 3::4] << 6) + scales = (torch.rand((*code.shape[:-1], code.shape[-1] // 64), generator=generator) * 0.1 - 0.05).to(dtype) + arrays[f"w{index}"] = packed + arrays[f"s{index}"] = scales + arrays[f"b{index}"] = (torch.randn(code.shape[:-1], generator=generator) * 0.01).to(dtype) + matrices.append( + ((code[:10].float() - 2).reshape(10, code.shape[1], -1, 64) * scales[:10].float().unsqueeze(-1)).reshape( + 10, code.shape[1], code.shape[2] + ) + ) + arrays["b1"][0, :128] = 0 + element_type = TensorProto.BFLOAT16 if dtype == torch.bfloat16 else TensorProto.FLOAT16 + node = helper.make_node( + "QMoE", + ["x", "router", "w1", "s1", "b1", "w2", "s2", "b2"], + ["y"], + domain="com.microsoft", + expert_weight_bits=2, + block_size=64, + weights_prepacked=0, + quant_type="int", + k=10, + activation_type="swiglu", + swiglu_fusion=1, + normalize_routing_weights=1, + ) + graph = helper.make_graph( + [node], + "int2_fake_int4", + [ + helper.make_tensor_value_info("x", element_type, [rows, hidden]), + helper.make_tensor_value_info("router", element_type, [rows, experts]), + ], + [helper.make_tensor_value_info("y", element_type, [rows, hidden])], + [_initializer(name, values) for name, values in arrays.items()], + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)] + ) + model.ir_version = 10 + x = (torch.randn((rows, hidden), generator=generator) * 0.2).to(dtype) + x[0, :64] = 0 + x[0, 64:68] = torch.tensor([7, 0.5, 1.5, -0.5], dtype=dtype) + router = torch.full((rows, experts), -100, dtype=dtype) + router[:, :10] = torch.linspace(-1, 1, 10).to(dtype) + routing = torch.softmax(router[:, :10].float(), -1) + expected = torch.zeros((rows, hidden)) + quantized_x = _fake_int4(x) if fake_activations else x + for expert in range(10): + gate_up = (quantized_x.float() @ matrices[0][expert].T).to(dtype).float() + gate_up += arrays["b1"][expert].float() + activated = (torch.nn.functional.silu(gate_up[:, 0::2]) * gate_up[:, 1::2]).to(dtype) + down_input = _fake_int4(activated) if fake_activations else activated + down = (down_input.float() @ matrices[1][expert].T).to(dtype).float() + expected += (down + arrays["b2"][expert].float()) * routing[:, expert : expert + 1] + return model, {"x": x, "router": router}, expected.to(dtype).float().numpy(), element_type + + +@pytest.mark.skipif("CUDAExecutionProvider" not in ort.get_available_providers(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("rows", [1, 3, 32]) +def test_int2_fake_int4_reference_and_graph(monkeypatch, dtype, rows): + model, feeds, reference, element_type = _case(rows, dtype) + monkeypatch.setenv("ORT_QMOE_INT2_LARGE_EXPERTS", "1") + monkeypatch.setenv("ORT_ENABLE_QMOE_INT2_PREFILL", "1") + monkeypatch.setenv("ORT_QMOE_INT2_FAKE_INT4_ACTIVATIONS", "1") + options = ort.SessionOptions() + options.add_session_config_entry("ep.cuda.qmoe_int_dequant_max_scratch_bytes", "1") + session = ort.InferenceSession( + model.SerializeToString(), + sess_options=options, + providers=[("CUDAExecutionProvider", {"enable_cuda_graph": "1"})], + ) + device_feeds = {name: value.cuda() for name, value in feeds.items()} + output = torch.empty((rows, 128), dtype=dtype, device="cuda") + binding = session.io_binding() + for name, value in device_feeds.items(): + binding.bind_input(name, "cuda", 0, element_type, tuple(value.shape), value.data_ptr()) + binding.bind_output("y", "cuda", 0, element_type, tuple(output.shape), output.data_ptr()) + torch.cuda.synchronize() + session.run_with_iobinding(binding) + torch.cuda.synchronize() + captured = output.float().cpu().numpy().copy() + session.run_with_iobinding(binding) + torch.cuda.synchronize() + np.testing.assert_array_equal(captured, output.float().cpu().numpy()) + assert np.isfinite(captured).all() + np.testing.assert_allclose(captured, reference, atol=0.0005 if dtype == torch.float16 else 0.003, rtol=0.03) + unquantized_reference = _case(rows, dtype, fake_activations=False)[2] + assert np.max(np.abs(captured - unquantized_reference)) > 0.004 + np.testing.assert_array_equal(device_feeds["x"].cpu().float().numpy(), feeds["x"].float().numpy()) + + +@pytest.mark.skipif("CUDAExecutionProvider" not in ort.get_available_providers(), reason="CUDA required") +def test_int2_fake_int4_rejects_unquantized_fallback(monkeypatch): + model, feeds, _, _ = _case(1, torch.float16) + monkeypatch.setenv("ORT_QMOE_INT2_FAKE_INT4_ACTIVATIONS", "1") + monkeypatch.setenv("ORT_ENABLE_QMOE_INT2_PREFILL", "1") + monkeypatch.setenv("ORT_QMOE_INT2_LARGE_EXPERTS", "0") + session = ort.InferenceSession(model.SerializeToString(), providers=["CUDAExecutionProvider"]) + with pytest.raises(ort.capi.onnxruntime_pybind11_state.Fail, match="refusing an unquantized fallback"): + session.run(None, {name: value.numpy() for name, value in feeds.items()}) + + +@pytest.mark.skipif("CUDAExecutionProvider" not in ort.get_available_providers(), reason="CUDA required") +def test_int2_fake_int4_rejects_disabled_grouped_path(monkeypatch): + model, _, _, _ = _case(1, torch.float16) + monkeypatch.setenv("ORT_QMOE_INT2_FAKE_INT4_ACTIVATIONS", "1") + monkeypatch.setenv("ORT_ENABLE_QMOE_INT2_PREFILL", "0") + with pytest.raises( + (ort.capi.onnxruntime_pybind11_state.RuntimeException, ort.capi.onnxruntime_pybind11_state.Fail), + match="INT4 activation simulation requires ORT_ENABLE_QMOE_INT2_PREFILL", + ): + ort.InferenceSession(model.SerializeToString(), providers=["CUDAExecutionProvider"]) diff --git a/onnxruntime/test/python/transformers/test_qmoe_int2_large_experts_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_int2_large_experts_cuda.py new file mode 100644 index 0000000000000..34f24849a4417 --- /dev/null +++ b/onnxruntime/test/python/transformers/test_qmoe_int2_large_experts_cuda.py @@ -0,0 +1,91 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import numpy as np +import pytest +import torch +from onnx import TensorProto, helper, numpy_helper + +import onnxruntime as ort + +pytestmark = pytest.mark.skipif( + "CUDAExecutionProvider" not in ort.get_available_providers() + or not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] < 8, + reason="CUDA SM80+ required", +) + + +@pytest.mark.skipif("CUDAExecutionProvider" not in ort.get_available_providers(), reason="CUDA required") +@pytest.mark.parametrize("block_size", [64, 128]) +@pytest.mark.parametrize("rows", [1, 3, 32]) +def test_int2_large_experts_packed_graph_parity(monkeypatch, block_size, rows): + monkeypatch.setenv("ORT_QMOE_INT2_FAKE_INT4_ACTIVATIONS", "0") + experts, hidden, inter = 512, 128, 128 + rng = np.random.default_rng(42) + arrays = { + "w1": rng.integers(0, 256, (experts, inter * 2, hidden // 4), dtype=np.uint8), + "w2": rng.integers(0, 256, (experts, hidden, inter // 4), dtype=np.uint8), + "s1": rng.uniform(-0.02, 0.02, (experts, inter * 2, hidden // block_size)).astype(np.float16), + "s2": rng.uniform(-0.02, 0.02, (experts, hidden, inter // block_size)).astype(np.float16), + "b1": rng.normal(0, 0.01, (experts, inter * 2)).astype(np.float16), + "b2": rng.normal(0, 0.01, (experts, hidden)).astype(np.float16), + } + node = helper.make_node( + "QMoE", + ["x", "router", "w1", "s1", "b1", "w2", "s2", "b2"], + ["y"], + domain="com.microsoft", + expert_weight_bits=2, + block_size=block_size, + weights_prepacked=0, + quant_type="int", + k=10, + activation_type="swiglu", + swiglu_fusion=1, + normalize_routing_weights=1, + ) + graph = helper.make_graph( + [node], + "int2_large_experts", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT16, [rows, hidden]), + helper.make_tensor_value_info("router", TensorProto.FLOAT16, [rows, experts]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT16, [rows, hidden])], + [numpy_helper.from_array(value, name) for name, value in arrays.items()], + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)] + ) + model.ir_version = 10 + feeds = { + "x": rng.normal(size=(rows, hidden)).astype(np.float16), + "router": rng.normal(size=(rows, experts)).astype(np.float16), + } + monkeypatch.setenv("ORT_ENABLE_QMOE_INT2_PREFILL", "1") + monkeypatch.setenv("ORT_QMOE_INT2_LARGE_EXPERTS", "0") + reference_session = ort.InferenceSession(model.SerializeToString(), providers=["CUDAExecutionProvider"]) + reference = reference_session.run(None, feeds)[0] + del reference_session + monkeypatch.setenv("ORT_QMOE_INT2_LARGE_EXPERTS", "1") + options = ort.SessionOptions() + options.add_session_config_entry("ep.cuda.qmoe_int_dequant_max_scratch_bytes", "1") + session = ort.InferenceSession( + model.SerializeToString(), + sess_options=options, + providers=[("CUDAExecutionProvider", {"enable_cuda_graph": "1"})], + ) + values = {name: ort.OrtValue.ortvalue_from_numpy(data, "cuda", 0) for name, data in feeds.items()} + binding = session.io_binding() + for name, value in values.items(): + binding.bind_ortvalue_input(name, value) + binding.bind_output("y", "cuda", 0) + session.run_with_iobinding(binding) + captured = binding.copy_outputs_to_cpu()[0] + session.run_with_iobinding(binding) + replayed = binding.copy_outputs_to_cpu()[0] + assert captured.shape == (rows, hidden) + assert np.isfinite(captured).all() + np.testing.assert_array_equal(captured, replayed) + np.testing.assert_allclose(captured, reference, atol=0.005, rtol=0.02) diff --git a/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py index 2b95ffa455fe4..c0f28e85c02ce 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py @@ -153,7 +153,7 @@ def quantize_weight_to_nvfp4(weight, block_size=NVFP4_BLOCK_SIZE, global_scale=N # ============================================================================ -def swiglu_ref(x, alpha=1.702, limit=7.0): +def swiglu_ref(x, alpha=1.702, limit=7.0, beta=1.0): """SwiGLU activation matching the QMoE kernel implementation.""" dim = x.shape[-1] x = x.view(-1, dim // 2, 2) @@ -161,7 +161,7 @@ def swiglu_ref(x, alpha=1.702, limit=7.0): if limit is not None: g = g.clamp(max=limit) l_val = l_val.clamp(min=-limit, max=limit) - return g * torch.sigmoid(alpha * g) * (l_val + 1) + return g * torch.sigmoid(alpha * g) * (l_val + beta) # ============================================================================ @@ -187,6 +187,7 @@ def create_nvfp4_moe_onnx_graph( fc1_bias=None, fc2_bias=None, weights_prepacked=None, + swiglu_params=(1.702, 1.0, 7.0), ): """Build ONNX model with QMoE operator in NVFP4 mode.""" inputs = [ @@ -224,9 +225,9 @@ def create_nvfp4_moe_onnx_graph( quant_type="nvfp4", block_size=block_size, swiglu_fusion=1 if use_swiglu else 0, - swiglu_limit=7.0, - activation_alpha=1.702, - activation_beta=1.0, + swiglu_limit=swiglu_params[2], + activation_alpha=swiglu_params[0], + activation_beta=swiglu_params[1], domain="com.microsoft", ), ] @@ -342,6 +343,7 @@ def _run_nvfp4_moe_test( global_scale=None, fc1_global_scale=None, fc2_global_scale=None, + swiglu_params=(1.702, 1.0, 7.0), ): self._skip_if_no_fp4() @@ -413,6 +415,7 @@ def _run_nvfp4_moe_test( fc1_bias=fc1_bias, fc2_bias=fc2_bias, weights_prepacked=1 if offline_prepacked else weights_prepacked, + swiglu_params=swiglu_params, ) opts = onnxruntime.SessionOptions() @@ -488,6 +491,7 @@ def _run_nvfp4_moe_test( torch_dtype, fc1_bias, fc2_bias, + swiglu_params, ) max_diff = (ort_output.float() - ref_output.float()).abs().max().item() @@ -618,6 +622,7 @@ def _compute_reference( torch_dtype, fc1_bias=None, fc2_bias=None, + swiglu_params=(1.702, 1.0, 7.0), ): """Reference MoE forward pass using dequantized weights.""" num_tokens = input_tensor.shape[0] @@ -644,7 +649,11 @@ def _compute_reference( h = tokens @ w1.T if fc1_bias is not None: h = h + fc1_bias[e].float() - h = swiglu_ref(h) if use_swiglu else F.silu(h) + h = ( + swiglu_ref(h, alpha=swiglu_params[0], beta=swiglu_params[1], limit=swiglu_params[2]) + if use_swiglu + else F.silu(h) + ) h = h @ w2.T if fc2_bias is not None: h = h + fc2_bias[e].float() @@ -1696,5 +1705,321 @@ def test_nvfp4_kpacked_mtp_qwen38_official_shape(onnx_dtype, num_tokens, capfd, assert "DequantWeightBytes=" not in diagnostics +def _enable_nvfp4_packed_prefill(monkeypatch): + monkeypatch.setenv("ORT_QMOE_NVFP4_PACKED_PREFILL", "1") + monkeypatch.setenv("ORT_ENABLE_NVFP4_CUTLASS_GEMM", "0") + + +@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") +@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT16, TensorProto.BFLOAT16]) +def test_nvfp4_packed_prefill_default_matches_explicit(onnx_dtype, capfd, monkeypatch): + monkeypatch.setenv("ORT_ENABLE_NVFP4_CUTLASS_GEMM", "0") + monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + monkeypatch.delenv("ORT_QMOE_NVFP4_PACKED_PREFILL", raising=False) + shape = dict( + hidden_size=128, + inter_size=64, + num_experts=16, + top_k=10, + num_tokens=17, + onnx_dtype=onnx_dtype, + use_swiglu=True, + gemv_mode="0", + offline_prepacked=True, + input_scale=0.1, + ) + automatic = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) + assert "Route=packed_nvfp4_prefill" in capfd.readouterr().out + monkeypatch.setenv("ORT_QMOE_NVFP4_PACKED_PREFILL", "1") + explicit = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) + assert "Route=packed_nvfp4_prefill" in capfd.readouterr().out + torch.testing.assert_close(automatic, explicit, atol=0, rtol=0) + monkeypatch.setenv("ORT_QMOE_NVFP4_PACKED_PREFILL", "0") + TestQMoENVFP4()._run_nvfp4_moe_test(**shape) + assert "Route=packed_nvfp4_prefill" not in capfd.readouterr().out + + +@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") +@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT16, TensorProto.BFLOAT16]) +@pytest.mark.parametrize("num_tokens", [17, 65, 257]) +@pytest.mark.parametrize("hidden_size,inter_size", [(128, 64), (64, 128)]) +def test_nvfp4_packed_prefill_reference(onnx_dtype, num_tokens, hidden_size, inter_size, capfd, monkeypatch): + _enable_nvfp4_packed_prefill(monkeypatch) + monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + shape = dict( + hidden_size=hidden_size, + inter_size=inter_size, + num_experts=16, + top_k=10, + num_tokens=num_tokens, + onnx_dtype=onnx_dtype, + use_swiglu=True, + use_bias=True, + gemv_mode="0", + offline_prepacked=True, + input_scale=0.1, + enable_cuda_graph=True, + ) + packed = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) + assert "Route=packed_nvfp4_prefill" in capfd.readouterr().out + monkeypatch.setenv("ORT_QMOE_NVFP4_PACKED_PREFILL", "0") + dense = TestQMoENVFP4()._run_nvfp4_moe_test(**{**shape, "enable_cuda_graph": False}) + tolerance = 0.03 if onnx_dtype == TensorProto.BFLOAT16 else 0.005 + torch.testing.assert_close(packed, dense, atol=tolerance, rtol=tolerance) + + +@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") +@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT16, TensorProto.BFLOAT16]) +def test_nvfp4_packed_prefill_qwen_dimensions(onnx_dtype, capfd, monkeypatch): + _enable_nvfp4_packed_prefill(monkeypatch) + monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + shape = dict( + hidden_size=2560, + inter_size=640, + num_experts=16, + top_k=10, + num_tokens=65, + onnx_dtype=onnx_dtype, + use_swiglu=True, + use_bias=True, + gemv_mode="0", + offline_prepacked=True, + input_scale=0.05, + ) + packed = TestQMoENVFP4()._run_nvfp4_moe_test(**shape, enable_cuda_graph=True) + assert "Route=packed_nvfp4_prefill" in capfd.readouterr().out + monkeypatch.setenv("ORT_QMOE_NVFP4_PACKED_PREFILL", "0") + dense = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) + tolerance = 0.03 if onnx_dtype == TensorProto.BFLOAT16 else 0.005 + torch.testing.assert_close(packed, dense, atol=tolerance, rtol=tolerance) + + +@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") +@pytest.mark.parametrize("fused", ["0", "1"]) +@pytest.mark.parametrize("reuse_scratch", ["0", "1"]) +@pytest.mark.parametrize("tiled_down", ["0", "1"]) +def test_nvfp4_packed_prefill_empty_experts(fused, reuse_scratch, tiled_down, capfd, monkeypatch): + _enable_nvfp4_packed_prefill(monkeypatch) + monkeypatch.setenv("ORT_QMOE_NVFP4_FUSED_SWIGLU", fused) + monkeypatch.setenv("ORT_QMOE_NVFP4_REUSE_ROUTING_SCRATCH", reuse_scratch) + monkeypatch.setenv("ORT_QMOE_NVFP4_TILED_DOWN_OUTPUT", tiled_down) + monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + logits = numpy.tile(numpy.arange(16, dtype=numpy.float32), (65, 1)) + shape = dict( + hidden_size=64, + inter_size=64, + num_experts=16, + top_k=2, + num_tokens=65, + onnx_dtype=TensorProto.FLOAT16, + use_swiglu=True, + use_bias=True, + gemv_mode="0", + offline_prepacked=True, + input_scale=0.1, + router_logits_override=logits, + row_tile_size=17, + enable_cuda_graph=True, + ) + packed = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) + assert "Route=packed_nvfp4_prefill" in capfd.readouterr().out + monkeypatch.setenv("ORT_QMOE_NVFP4_PACKED_PREFILL", "0") + dense = TestQMoENVFP4()._run_nvfp4_moe_test(**{**shape, "enable_cuda_graph": False}) + torch.testing.assert_close(packed, dense, atol=0.005, rtol=0.005) + + +@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") +@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT16, TensorProto.BFLOAT16]) +@pytest.mark.parametrize("num_tokens", [17, 65]) +@pytest.mark.parametrize("hidden_size,inter_size", [(128, 64), (2560, 640)]) +@pytest.mark.parametrize("use_bias", [False, True]) +@pytest.mark.parametrize( + "swiglu_params", + [ + (1.702, 1.0, 7.0), + (1.0, 0.0, float("inf")), + (0.5, 0.2, 0.1), + ], +) +def test_nvfp4_packed_prefill_fused_swiglu( + onnx_dtype, num_tokens, hidden_size, inter_size, use_bias, swiglu_params, capfd, monkeypatch +): + _enable_nvfp4_packed_prefill(monkeypatch) + monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + shape = dict( + hidden_size=hidden_size, + inter_size=inter_size, + num_experts=16, + top_k=10, + num_tokens=num_tokens, + onnx_dtype=onnx_dtype, + use_swiglu=True, + use_bias=use_bias, + gemv_mode="0", + offline_prepacked=True, + input_scale=0.05, + enable_cuda_graph=True, + swiglu_params=swiglu_params, + ) + monkeypatch.setenv("ORT_QMOE_NVFP4_FUSED_SWIGLU", "0") + unfused = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) + unfused_diagnostics = capfd.readouterr().out + assert "Route=packed_nvfp4_prefill" in unfused_diagnostics + monkeypatch.setenv("ORT_QMOE_NVFP4_FUSED_SWIGLU", "1") + fused = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) + fused_diagnostics = capfd.readouterr().out + assert "Route=packed_nvfp4_prefill" in fused_diagnostics + unfused_workspace = int(re.search(r"RunnerWorkspaceBytes=(\d+)", unfused_diagnostics).group(1)) + fused_workspace = int(re.search(r"RunnerWorkspaceBytes=(\d+)", fused_diagnostics).group(1)) + gate_up_bytes = num_tokens * 10 * inter_size * 2 * 2 + assert unfused_workspace - fused_workspace == (gate_up_bytes + 255) // 256 * 256 + torch.testing.assert_close(fused, unfused, atol=0, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") +@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT16, TensorProto.BFLOAT16]) +@pytest.mark.parametrize("fused", ["0", "1"]) +@pytest.mark.parametrize("row_tile_size", [0, 17]) +@pytest.mark.parametrize( + "hidden_size,inter_size,num_experts,top_k", + [ + (128, 64, 16, 10), + (2560, 640, 16, 10), + (64, 64, 512, 1), + (2560, 640, 512, 10), + ], +) +def test_nvfp4_packed_prefill_reuse_routing_scratch( + onnx_dtype, fused, row_tile_size, hidden_size, inter_size, num_experts, top_k, capfd, monkeypatch +): + _enable_nvfp4_packed_prefill(monkeypatch) + monkeypatch.setenv("ORT_QMOE_NVFP4_FUSED_SWIGLU", fused) + monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + shape = dict( + hidden_size=hidden_size, + inter_size=inter_size, + num_experts=num_experts, + top_k=top_k, + num_tokens=65, + onnx_dtype=onnx_dtype, + use_swiglu=True, + use_bias=True, + gemv_mode="0", + offline_prepacked=True, + input_scale=0.05, + enable_cuda_graph=True, + row_tile_size=row_tile_size, + ) + monkeypatch.setenv("ORT_QMOE_NVFP4_REUSE_ROUTING_SCRATCH", "0") + separate = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) + separate_diagnostics = capfd.readouterr().out + monkeypatch.setenv("ORT_QMOE_NVFP4_REUSE_ROUTING_SCRATCH", "1") + reused = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) + reused_diagnostics = capfd.readouterr().out + assert "Route=packed_nvfp4_prefill" in separate_diagnostics + assert "Route=packed_nvfp4_prefill" in reused_diagnostics + separate_bytes = int(re.search(r"RunnerWorkspaceBytes=(\d+)", separate_diagnostics).group(1)) + reused_bytes = int(re.search(r"RunnerWorkspaceBytes=(\d+)", reused_diagnostics).group(1)) + rows = row_tile_size or 65 + block_rows = 32 + while block_rows < 1024 and (rows + block_rows - 1) // block_rows * num_experts > block_rows: + block_rows *= 2 + blocks = (rows + block_rows - 1) // block_rows + + def align(value): + return (value + 255) // 256 * 256 + + routing_bytes = 2 * align(num_experts * blocks * 4) + align(rows * num_experts * 4) + input_bytes = align(rows * top_k * hidden_size * 2) + assert separate_bytes - reused_bytes == min(routing_bytes, input_bytes) + torch.testing.assert_close(reused, separate, atol=0, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") +@pytest.mark.parametrize("onnx_dtype", [TensorProto.FLOAT16, TensorProto.BFLOAT16]) +@pytest.mark.parametrize( + "hidden_size,inter_size,num_experts,top_k", + [ + (64, 64, 16, 2), + (1152, 128, 16, 10), + (2560, 640, 16, 10), + (1152, 64, 512, 10), + ], +) +@pytest.mark.parametrize("row_tile_size", [0, 17]) +@pytest.mark.parametrize("reuse_scratch", ["0", "1"]) +def test_nvfp4_packed_prefill_tiled_down( + onnx_dtype, hidden_size, inter_size, num_experts, top_k, row_tile_size, reuse_scratch, capfd, monkeypatch +): + _enable_nvfp4_packed_prefill(monkeypatch) + monkeypatch.setenv("ORT_QMOE_NVFP4_FUSED_SWIGLU", "1") + monkeypatch.setenv("ORT_QMOE_NVFP4_REUSE_ROUTING_SCRATCH", reuse_scratch) + monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + shape = dict( + hidden_size=hidden_size, + inter_size=inter_size, + num_experts=num_experts, + top_k=top_k, + num_tokens=65, + onnx_dtype=onnx_dtype, + use_swiglu=True, + use_bias=row_tile_size == 0, + gemv_mode="0", + offline_prepacked=True, + input_scale=0.05, + enable_cuda_graph=True, + row_tile_size=row_tile_size, + ) + monkeypatch.setenv("ORT_QMOE_NVFP4_TILED_DOWN_OUTPUT", "0") + baseline = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) + baseline_bytes = int(re.search(r"RunnerWorkspaceBytes=(\d+)", capfd.readouterr().out).group(1)) + monkeypatch.setenv("ORT_QMOE_NVFP4_TILED_DOWN_OUTPUT", "1") + tiled = TestQMoENVFP4()._run_nvfp4_moe_test(**shape) + tiled_bytes = int(re.search(r"RunnerWorkspaceBytes=(\d+)", capfd.readouterr().out).group(1)) + rows = row_tile_size or 65 + + def aligned(value): + return (value + 255) // 256 * 256 + + input_bytes = aligned(rows * top_k * hidden_size * 2) + tiled_bytes_reservation = aligned(rows * top_k * min(hidden_size, 1024) * 2) + if reuse_scratch == "1": + block_rows = 32 + while block_rows < 1024 and (rows + block_rows - 1) // block_rows * num_experts > block_rows: + block_rows *= 2 + blocks = (rows + block_rows - 1) // block_rows + sorting_bytes = 2 * aligned(num_experts * blocks * 4) + aligned(rows * num_experts * 4) + input_bytes = max(input_bytes, sorting_bytes) + tiled_bytes_reservation = max(tiled_bytes_reservation, sorting_bytes) + assert baseline_bytes - tiled_bytes == input_bytes - tiled_bytes_reservation + torch.testing.assert_close(tiled, baseline, atol=0, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp4_qmoe, reason="Requires CUDA NVFP4 QMoE") +@pytest.mark.parametrize("mode", ["raw_layout", "unsupported_shape", "decode"]) +def test_nvfp4_packed_prefill_fallback(mode, capfd, monkeypatch): + _enable_nvfp4_packed_prefill(monkeypatch) + monkeypatch.setenv("ORT_QMOE_NVFP4_FUSED_SWIGLU", "1") + monkeypatch.setenv("ORT_QMOE_NVFP4_REUSE_ROUTING_SCRATCH", "1") + monkeypatch.setenv("ORT_QMOE_NVFP4_TILED_DOWN_OUTPUT", "1") + monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + decode = mode == "decode" + TestQMoENVFP4()._run_nvfp4_moe_test( + hidden_size=512 if decode else (80 if mode == "unsupported_shape" else 64), + inter_size=512 if decode else 64, + num_experts=16, + top_k=10, + num_tokens=1 if decode else 17, + onnx_dtype=TensorProto.FLOAT16, + use_swiglu=True, + gemv_mode="1" if decode else "0", + offline_prepacked=mode != "raw_layout", + input_scale=0.1, + ) + diagnostics = capfd.readouterr().out + assert "Route=packed_nvfp4_prefill" not in diagnostics + assert f"Route={'fp4_gemv_raw' if decode else 'grouped_moe'}" in diagnostics + + if __name__ == "__main__": unittest.main() diff --git a/onnxruntime/test/python/transformers/test_sparse_paged_attention_prefill_cuda.py b/onnxruntime/test/python/transformers/test_sparse_paged_attention_prefill_cuda.py new file mode 100644 index 0000000000000..a5f5be07e2985 --- /dev/null +++ b/onnxruntime/test/python/transformers/test_sparse_paged_attention_prefill_cuda.py @@ -0,0 +1,197 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import argparse +import os +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest +from onnx import TensorProto, helper + +import onnxruntime as ort + +torch = pytest.importorskip("torch") + + +def _run_case(dtype_name, mode): + grouped = mode.startswith("grouped") + dtype = torch.bfloat16 if dtype_name == "bfloat16" else torch.float16 + element_type = TensorProto.BFLOAT16 if dtype == torch.bfloat16 else TensorProto.FLOAT16 + rng = np.random.default_rng(33393) + tokens, heads, kv_heads, width, past, block_size = 130, 24, 2, 256, 257, 16 + length = past + tokens + blocks = (length + block_size - 1) // block_size + selected_width = length + 3 + + def floating(shape): + return torch.tensor(rng.normal(0, 0.2, shape), dtype=dtype, device="cuda") + + query = floating((tokens, heads, width)) + key = floating((tokens, kv_heads, width)) + value = floating((tokens, kv_heads, width)) + key_cache = floating((blocks, block_size, kv_heads, width)) + value_cache = floating((blocks, block_size, kv_heads, width)) + # A suppressed write must read these deliberately different prefix-cache values. + key_cache.reshape(-1, kv_heads, width)[past + 5] = 0.5 + value_cache.reshape(-1, kv_heads, width)[past + 5] = 5 + key[5] = -0.5 + value[5] = -5 + slots = np.arange(past, length, dtype=np.int32) + slots[5] = -1 + block_table = np.arange(blocks, dtype=np.int32) + block_table[2] = -1 + selected = np.tile(np.arange(selected_width, dtype=np.int32), (tokens, 1)) + selected[:, -3:] = [-1, length, length + 5] + counts = np.full(tokens, selected_width, dtype=np.int32) + counts[0] = 0 + counts[1] = 65 + counts[2] = 129 + selected[3] = -1 + sink = floating((heads,)) + sink[0] = 4.5 + + q = query.float().cpu().numpy().astype(np.float64) + k = key.float().cpu().numpy().astype(np.float64) + v = value.float().cpu().numpy().astype(np.float64) + kc = key_cache.float().cpu().numpy().reshape(-1, kv_heads, width).astype(np.float64) + vc = value_cache.float().cpu().numpy().reshape(-1, kv_heads, width).astype(np.float64) + sink_ref = sink.float().cpu().numpy().astype(np.float64) + scale, softcap = 0.0625, 1.7 + expected = np.zeros((tokens, heads, width), dtype=np.float64) + for token in range(tokens): + positions = selected[token, : counts[token]] + positions = positions[(positions >= 0) & (positions <= past + token) & (positions < length)] + positions = positions[block_table[positions // block_size] >= 0] + if positions.size == 0: + continue + keys, values = kc[positions].copy(), vc[positions].copy() + for index, position in enumerate(positions): + if position >= past and slots[position - past] >= 0: + keys[index], values[index] = k[position - past], v[position - past] + for head in range(heads): + kv_head = head // (heads // kv_heads) + logits = keys[:, kv_head] @ q[token, head] * scale + logits = softcap * np.tanh(logits / softcap) + maximum = max(logits.max(), sink_ref[head]) + weights = np.exp(logits - maximum) + expected[token, head] = weights @ values[:, kv_head] / (weights.sum() + np.exp(sink_ref[head] - maximum)) + + feed = { + "query": query.reshape(tokens, heads * width), + "key": key.reshape(tokens, kv_heads * width), + "value": value.reshape(tokens, kv_heads * width), + "key_cache": key_cache, + "value_cache": value_cache, + "cumulative_sequence_length": torch.tensor([0, tokens], dtype=torch.int32, device="cuda"), + "past_seqlens": torch.tensor([past], dtype=torch.int32, device="cuda"), + "block_table": torch.tensor(block_table[None], device="cuda"), + "slot_mapping": torch.tensor(slots, device="cuda"), + "selected_indices": torch.tensor(selected, device="cuda"), + "selected_counts": torch.tensor(counts, device="cuda"), + "head_sink": sink, + } + names = [*list(feed)[:-1], "", "", "", "", "", "head_sink"] + inputs = [ + helper.make_tensor_value_info( + name, TensorProto.INT32 if tensor.dtype == torch.int32 else element_type, list(tensor.shape) + ) + for name, tensor in feed.items() + ] + node = helper.make_node( + "SparsePagedAttention", + names, + ["output"], + domain="com.microsoft", + num_heads=heads, + kv_num_heads=kv_heads, + scale=scale, + softcap=softcap, + ) + model = helper.make_model( + helper.make_graph( + [node], + "sparse_prefill_reference", + inputs, + [helper.make_tensor_value_info("output", element_type, [tokens, heads * width])], + ), + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)], + ) + model.ir_version = 10 + ort.set_default_logger_severity(0) + ort.set_default_logger_verbosity(1) + options = ort.SessionOptions() + options.log_severity_level = 0 + options.log_verbosity_level = 1 + options.add_session_config_entry("ep.cuda.sparse_paged_attention_grouped", str(int(grouped))) + options.add_session_config_entry( + "ep.cuda.sparse_paged_attention_grouped_vectorized", str(int(mode == "grouped_vectorized")) + ) + session = ort.InferenceSession( + model.SerializeToString(), + options, + providers=[("CUDAExecutionProvider", {"device_id": torch.cuda.current_device()})], + ) + assert session.get_providers()[0] == "CUDAExecutionProvider" + binding = session.io_binding() + device = torch.cuda.current_device() + for name, tensor in feed.items(): + tensor_type = TensorProto.INT32 if tensor.dtype == torch.int32 else element_type + binding.bind_input(name, "cuda", device, tensor_type, tuple(tensor.shape), tensor.data_ptr()) + output = torch.empty((tokens, heads * width), dtype=dtype, device="cuda") + binding.bind_output("output", "cuda", device, element_type, tuple(output.shape), output.data_ptr()) + torch.cuda.synchronize() + session.run_with_iobinding(binding) + binding.synchronize_outputs() + actual = output.float().cpu().numpy().reshape(tokens, heads, width) + tolerance = 0.004 if dtype == torch.float16 else 0.015 + np.testing.assert_allclose(actual, expected, atol=tolerance, rtol=tolerance) + print("sparse prefill reference passed", flush=True) + + +@pytest.mark.skipif("CUDAExecutionProvider" not in ort.get_available_providers(), reason="CUDA EP required") +@pytest.mark.parametrize( + "dtype,mode", + [ + ("float16", "scalar"), + ("float16", "qk"), + ("float16", "qk_pv"), + ("bfloat16", "qk"), + ("float16", "grouped"), + ("float16", "grouped_vectorized"), + ], +) +def test_sparse_prefill_specializations_match_reference(dtype, mode): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8: + pytest.skip("SM80+ CUDA required") + environment = dict( + os.environ, + ORT_SPARSE_PREFILL_TENSOR_CORE_QK=str(int(mode in ("qk", "qk_pv"))), + ORT_SPARSE_PREFILL_TENSOR_CORE_PV=str(int(mode == "qk_pv")), + ) + result = subprocess.run( + [sys.executable, str(Path(__file__).resolve()), "--dtype", dtype, "--mode", mode], + env=environment, + capture_output=True, + text=True, + check=False, + timeout=180, + ) + assert result.returncode == 0, result.stdout + result.stderr + expected_dispatch = ( + f"SparsePagedAttention dispatch: grouped, vectorized={int(mode == 'grouped_vectorized')}" + if mode.startswith("grouped") + else f"SparsePagedAttention dispatch: tiled, heads=12, qk={int(mode != 'scalar')}, pv={int(mode == 'qk_pv')}" + ) + assert expected_dispatch in result.stderr, result.stdout + result.stderr + assert "sparse prefill reference passed" in result.stdout + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--dtype", choices=["float16", "bfloat16"], required=True) + parser.add_argument("--mode", choices=["scalar", "qk", "qk_pv", "grouped", "grouped_vectorized"], required=True) + args = parser.parse_args() + _run_case(args.dtype, args.mode)