Repository navigation
[CUDA] QMoE: select the 64-row CTA tile for repacked INT3 prefill when experts average >= 24 rows - #33714
Merged
Tianlei Wu (tianleiwu) merged 1 commit intoOct 29, 2026
Conversation
…n experts average >= 24 rows RunPackedGroupedGemm hard-coded CtaShape32x128x64 for the INT4 grouped GEMM that INT3 prefill repacks into. At Qwen3.8 prefill sizes (>= 2048 tokens, ~40 rows per expert) that tile is 7-9% slower per launch than CtaShape64x128x64, which costs 2-3% TTFT on INT3 variants compared with the previous selection (64x128x64 once rows_per_expert >= 24, 32x128x64 below it). Restore that selection, using the same crossover as PickPrefillTactic. Only the CTA/warp tiling changes; the K accumulation order of every output element is unchanged. Extend QMoETest_Int3AndMixedWidthCudaPrefillParity with 64-row cases (32 rows per expert on average) so the 64-row tile is covered. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
🟢 Approval recommended
No active review issues remain for the INT2/INT3 prefill tile-selection change.
No open issues 🎉
What changed in this PR
This PR tunes CUDA QMoE prefill for repacked INT2/INT3 weights by selecting a larger grouped-GEMM tile when experts average at least 24 rows, matching the existing INT4 prefill crossover.
Changes:
- Selects the 64-row CTA tile at the threshold and retains the 32-row tile below it.
- Adds three 64-row INT3 and mixed-width parity cases to exercise the larger tile with FP16 and BF16.
| File | Description |
|---|---|
onnxruntime/test/contrib_ops/moe_test.cc |
Adds parity cases that exercise the 64-row tile. |
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_int2_prefill.cu |
Selects the grouped-GEMM tile using average rows per expert. |
Review effort: Balanced
Akshay Sonawane (apsonawane)
approved these changes
Oct 29, 2026
Tianlei Wu (tianleiwu)
deleted the
tlwu/20261029/qmoe-int3-prefill-64row-tile
branch
October 29, 2026 20:36
Tianlei Wu (tianleiwu)
added a commit
that referenced
this pull request
Oct 31, 2026
…arge-M dequantize + cuBLAS GEMM (#33689) ### Description This PR adds tensor-core CUDA kernels for the `QMoE` and `MatMulNBits` hot spots of Qwen3.8-Flash-Next (B = 1, `S = 8` DFlash2 speculative decoding): `QMoE` INT4 and INT3 grouped GEMV for BF16 activations with packed bf16x2 INT3 dequantization, a warp-coalesced weight/scale layout for constant INT3 weights that the INT3 decode kernel reads with one 16-byte load per lane, a compact 80-route variant of the INT4 FC2 kernel, a dense symmetric-INT8 `MatMulNBits` GEMV for `M <= 8` (SM 9.0), and a large-`M` route for the same dense INT8 layers (`M >= 512`, default on SM 9.0) that dequantizes the weight once and runs one cuBLAS GEMM instead of the row-chunked CUTLASS mixed-input GEMM. It also triggers the programmatic-dependent-launch (PDL) successor early in the two fpA_intB GEMV kernels. The PR was split by operator. The `DynamicSparseAttention` kernels, the gate-mul early PDL trigger and the hyper-connection sigmoid, which this PR carried before, are now in #33707 (see "Related PRs and review guide"). #33740, which was stacked on this PR, is merged into this branch, so its scope (the coalesced INT3 prepack, the prefill repack that consumes it and the compact INT4 FC2 kernel) is part of this PR and of this description. The last commit also addresses the open review feedback of both PRs: four open threads of this PR and two threads of #33740 whose concerns were still present when it merged (see "7. Review follow-up"). On H200 (`B = 1`, DFlash2 `S = 8`, 2 rotated GPU rounds x median of 5 runs), comparing the **final head of this PR** (`d6813ee3ab`) against a baseline built from the exact same companion PRs (`main` `0a12bc0480` + #33714 `ebde3f13e5` + #33707 `204d6593c0` + #33650 / #32126 `78d59a59c7`; GenAI #2808 `5d561a2f5`; package options `prefill.chunk = 8192` and `qmoe_row_tile_size = 8192` held constant): - **Decode (`ms/cycle` at 2K / 8K / 32K context)** drops from `13.82 / 13.96 / 14.02` ms to `10.85 / 10.99 / 11.24` ms (**-21.4% / -21.3% / -19.9%**) for `all_int4`, `14.29 / 14.66 / 14.59` ms to `10.70 / 10.91 / 11.16` ms (**-25.1% / -25.6% / -23.5%**) for `top16`, and `14.42 / 14.51 / 14.92` ms to `10.71 / 10.87 / 11.23` ms (**-25.7% / -25.1% / -24.7%**) for `top08`. - **Prefill TTFT (`8192 / 32768` prompt tokens)** drops from `648.4 / 2760.9` ms to `579.2 / 2444.2` ms (**-10.7% / -11.5%**) for `all_int4`, `668.2 / 2838.6` ms to `593.1 / 2522.0` ms (**-11.2% / -11.2%**) for `top16`, and `672.3 / 2850.2` ms to `599.2 / 2544.0` ms (**-10.9% / -10.7%**) for `top08` (and **-11.0% / -11.8% / -11.2%** at 8192, **-11.2% / -11.2% / -10.8%** at 32768 when toggling `ORT_FPA_INTB_DEQUANT_GEMM_MIN_M` from `0` to `512` on the final head). Every new tensor-core kernel and the new route have a kill switch, available as an environment variable and as a CUDA session option (`ep.cuda.*`); the bit-identical INT3 bf16x2 dequantization and the early PDL trigger have none. The tensor-core kernels fall back to the existing kernels outside their gates or when the loaded PTX target is older than SM80. #### Summary of Changes ##### 1. QMoE decode (BF16): INT3 bf16x2 dequantization, INT4 and INT3 tensor-core GEMV | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu` | `DequantizeInt3x8Bf16`: two INT3 codes (`j` and `j + 4`) per 32-bit register from one `lop3.b32`, one `HSUB2` and one `HMUL2`. It is bit-identical to the scalar chain (about 6 instructions per weight); the static instruction count of `FusedFc1FastVecKernel<bf16, 3, 64>` drops from 4096 to 3448. `int3_tc::TryLaunchBf16` is tried first on the vectorized decode path. | | `onnxruntime/contrib_ops/cuda/moe/qmoe_int3_tc.cuh` (new) | Tensor-core grouped FC1 (with the interleaved SwiGLU) and FC2 GEMV for INT3 block-64 weights: `mma.sync.m16n8k16` BF16 with FP32 accumulation, up to 8 routes of one expert per CTA. Every warp owns whole-K m-blocks (FC1 `Int3TcKernel<true, 4, 2, 3>`: 4 warps x 2 m-blocks = 128 rows per CTA, 3 CTAs per SM; FC2 `<false, 4, 1, 5>`: 64 rows, 5 CTAs per SM), so there is no cross-warp K split, no shared-memory reduction and no barrier, and the dynamic shared memory holds only the staged activations (`16 * hidden` or `16 * intermediate` B, at most 48 KB). The kernel is one rolled loop over (m-block, 512-wide super-step) items that requests the weights and scales of the next item (two register sets) before it computes the current one. A compile-time unrolled window does not work, because `ptxas` hoists every independent load of a basic block to its top and spills; only a real loop back edge keeps the prefetch distance (FC1: 116 registers instead of 167). The five K partitions (super-step `s % 5`) keep separate FP32 sums that are added in partition order, which is bit-identical to a cross-warp K split. Dequantization (`bd53cf38e2`) uses the range-free packed `bf16x2` sequence (`lop3.b32` into normal BF16 `{8 + c_lo, 128 + c_hi}`, exact centering by `{12, 132}` via `HSUB2`, and `HMUL2` by the block scale) with unscaled shared-memory activations, so zero effective weights (`code = 4`) center to `+0.0` for every finite BF16 scale and all finite BF16 scales and activations are supported without an exponent-window restriction. | | `onnxruntime/contrib_ops/cuda/moe/qmoe_int3_tc.cuh` (#33740) | `PrePackInt3WeightsKernel` and `PrePackInt3ScalesKernel` rewrite constant INT3 weights and BF16 scales once, at session creation, into a warp-coalesced `mma.sync.m16n8k16` layout (512 K per super-step, 16 N per m-block, one `uint4` per lane per 128-K chunk). `Int3TcCoalescedKernel` replaces the scattered row-local 3-bit reads of `Int3TcKernel` with those coalesced loads. It is bit-identical to `Int3TcKernel` (same dequantization, accumulation order and finalize) and has the same numerics contract. | | `onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_int2_prefill.{h,cu}` (#33740) | `RepackPrepackedInt3ToSm80Int4ActiveExpertsKernel` and `TransposeScalesActiveExpertsVec8Kernel<T, true>` build the SM80 INT4 operands of the active experts for the INT3 prefill GEMM from the prepacked weights and scales, so the prefill path needs no raw-layout fallback. | | `onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_tc.cuh` (new) | Tensor-core INT4 (group 32) GEMV for FC1, interleaved-SwiGLU FC1 and FC2. The weights are the A operand of `mma.sync.m16n8k16`; up to 8 consecutive rows of one expert are its N side and share every dequantized weight tile. The K loop is software pipelined: the packed weights and scales of step `s + 1` are in registers while step `s` is computed, and lanes without a route read the activations of the last route so that every load is unconditional. The loop is deliberately not peeled: `__launch_bounds__` carries the occupancy target (5, 6 or 9 CTAs per SM for 4, 2 or 1 m-blocks per warp), and peeling or carrying the activation registers over the loop makes `ptxas` allocate more registers and lose a resident CTA. FC1 is `MoeGemvTcInt4Kernel<true, 4, 4>` (the four warps split K); FC2 is `<false, 1, 1>` (a warp owns 16 columns over the whole K, 56 registers). The summation order is unchanged by the pipelining. #33740 adds a compact FC2 variant for `expanded_num_rows == 80` (`<false, 1, 2, true>`, `kCompact80`): the leaders of the expert row groups are found with warp ballots, which handles experts with more than 8 routes in several chunks without reading `expert_first_token_offset`; `LoadWeightsAfterMb` lowers the register pressure; and the kernel issues `griddepcontrol.launch_dependents` before its global output store. | | `onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu` | Try `moe_tc::TryLaunchInt4Bf16` before the scalar kernels in `launch_moe_gemv_int_symmetric` and `launch_moe_gemv_int_symmetric_interleaved_swiglu`. | | `onnxruntime/contrib_ops/cuda/moe/moe_quantization.{h,cc}`, `onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.{h,cu}`, `onnxruntime/contrib_ops/cuda/moe/qmoe_fused_route.cuh`, `onnxruntime/contrib_ops/cuda/llm/moe_gemm/{moe_gemv.h,moe_kernels.h,moe_kernels.cu}` | `QMoE` resolves the two disable switches and the two row thresholds once at kernel construction (session option, else environment variable, else default) and passes them to the launchers as plain parameters (a disable switch sets its threshold to 0). From #33740: `QMoE::PrePack` (with a target-device guard) runs the two prepack kernels (`LaunchPrePackInt3Weights`, `LaunchPrePackInt3Scales`) for the constant INT3 weights and scales that `CanUseInt3TcPrePack` and the constructor's shape checks admit, and the decode path then runs the tensor-core kernel on that layout (`Int3TcCoalescedKernel`) from the first route; `moe_kernels.cu` passes `early_routing_loads = fc2_did_gemv` to `finalizeMoeRoutingKernelLauncher`, and `qmoe_fused_route.cuh` skips the route-group packing when `route_groups == nullptr`. | | `onnxruntime/test/contrib_ops/moe_test.cc` | `QMoETest_Bf16Int4TensorCoreGemvParity` (block 32) and `QMoETest_Bf16Int3TensorCoreGemvParity` (block 64): 12 tokens x top-4 over 8 experts, so that expert row groups are full, partial, single or empty, covering the FC1 / FC2 K loops and K tails, compared with a host reference (per-weight BF16 rounding, FP32 accumulation, abs 0.05 / rel 0.01). No existing QMoE test reached these kernels. `QMoETest_Bf16TensorCoreGemvExtremeScales` runs INT3 and INT4 with FC1 block scales near 2^33 and FC2 block scales near 2^-47, and the mirror image, against the host reference (the activations and the tolerance follow the scale factors), which covers the exact scale window of the INT3 kernel. `QMoETensorCoreSessionOptions` runs the parity case with `ep.cuda.disable_qmoe_int3_tc` / `ep.cuda.disable_moe_gemv_tc` = 0 and 1 while the environment variables hold an invalid value, which shows that explicit options override the environment. #33740 extends both parity tests with `num_rows = 20` (80 routes, which reaches the compact INT4 FC2 kernel and the INT3 prefill path). The three tests of the review follow-up are listed under 7. | ##### 2. Early PDL trigger in the fpA_intB GEMV kernels | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/dispatcher.h` | The scalar fpA_intB GEMV `kernel` issues `griddepcontrol.launch_dependents` first (`__CUDA_ARCH__ >= 900`). | | `onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/int8_tc_gemv.cuh` | `Int8TcGemvKernel` (section 3) does the same. | Both kernels feed the hyper-connection down kernel (`HcDownKernel`), which is launched with PDL and reads their output only after `griddepcontrol.wait`; its constant-weight prefetch now overlaps with them. There is no effect for plain launches, no change of any result and no switch. The same trigger in `ProjectedSigmoidGateMulKernel` is part of #33707; together they shorten an `S = 8` decode cycle by 1.4-1.7%. ##### 3. Dense INT8 tensor-core GEMV for `MatMulNBits` (SM 9.0, `M <= 8`) | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/int8_tc_gemv.cuh` (new) | The weights are the A operand of `mma.sync.m16n8k16` BF16 (16 output columns x the 1-8 activation rows, FP32 accumulation). The runtime's SM80 interleaved weight layout is consumed as it is: the K permutation is absorbed by choosing which activation words feed each `mma`. Dequantization uses one `fma.rn.bf16x2` (`bytes * 2^126 - 1 = (q - 128) / 128`, exact), one `mul.rn.bf16x2` by `128` to form the exact BF16 integer `q - 128`, and one `mul.rn.bf16x2` by the BF16 group scale, with no post-accumulation unscaling. Every weight therefore equals the scalar kernel's `bf16_rn((q - 128) * scale)` bit-for-bit across the entire finite BF16 scale range (including subnormal scales below `2^-119` down to `2^-133`), with no fixed exponent shift that could overflow or flush a valid tiny weight. Register-ring prefetch, a fixed-order cross-warp reduction (a row's result does not depend on `M`), and `LaunchInt8TcGemvAuto` picks the K split (4 or 8 warps), the prefetch depth and the CTAs per SM from `N / 16` relative to the SM count. The launcher returns false when the loaded PTX target is older than SM80, because the kernel body is empty there. | | `onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/dispatcher.h` | `CanUseInt8TensorCore` and `exec_kernel_tc`; the gate sits ahead of the scalar kernel in `dispatcher`, which runs the scalar kernel when `exec_kernel_tc` declines. The gate also requires `alpha == 1`: the tensor-core kernel does not apply `alpha`, so a launch with another value keeps the scalar kernel, which does. | | `onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/fpA_intB_gemv.h` | `Params::tensor_core`. | | `onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.{h,cc}`, `onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.{h,cc}`, `onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h`, `onnxruntime/contrib_ops/cuda/llm/gemm_tactic_cache.{h,cc}` | `MatMulNBits` resolves the switch at kernel construction and sets `Params::tensor_core` for both the run and the tactic profiler, so the profiler times the kernel that runs. The setting is part of the in-memory tactic-cache key and of the persisted cache (a new `tensor_core_gemv` column; `kTacticSelectionVersion` 4 -> 5), so sessions with different settings do not reuse each other's profiling results. | | `onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemv_int8_tensor_core_test.cc` (new) | `Int8TensorCoreGemvTest.*` (8 tests, SM 9.0 only): one-hot activations isolate the dequantization and must match the scalar kernel as values (a zero may have either sign), for ordinary and wide scale ranges; uniform codes at extreme scales (+-2^100 to +-2^-119) against the BF16 rounding of the exact product, which a fixed exponent shift would fail; random activations against a double-precision reference for the Qwen3.8 projection shapes, all group sizes and row counts; a row's result must not depend on the row count (bitwise); the launch configuration follows the grid size; launches outside the gate, including one with `alpha != 1`, use the scalar kernel, and `NonUnitAlphaScalesTheOutput` checks that an `alpha = 2` launch runs the scalar kernel and returns the scaled reference. | | `onnxruntime/test/contrib_ops/cuda_kernels/gemm_tactic_cache_test.cc`, `onnxruntime/test/contrib_ops/matmul_nbits_fpa_intb_layout_test.cc` | `GemmTacticCacheTest.TensorCoreSettingsHaveSeparateProfiles` (in-memory key and persisted round trip) and `MatMulNBitsFpAIntBLayout.TensorCoreSessionOptions` (`ep.cuda.disable_fpa_intb_gemv_tc` = 0 and 1 on SM 9.0 while the environment variable holds an invalid value). | ##### 4. Large-`M` dense INT8 for `MatMulNBits`: dequantize once + cuBLAS (SM 9.0 default) | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/int8_dequant.cuh` (new) | `DequantizeInt8InterleavedKernel` expands the SM80-interleaved INT8 weight (64-wide K tiles, two real columns per interleaved row) into a row-major `[N, K]` BF16 / FP16 matrix. `(q - 128) * scale` is exact in FP32 (an 8-bit integer times a 16-bit float) and rounded once to the activation type, so every weight has the value the GEMV kernels and the CUTLASS converter use. | | `onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/fpA_intB_gemv.{h,cu}` | `dequantize_int8_interleaved` launches it and returns false (the caller keeps CUTLASS) for `N % 16 != 0`, `K % 64 != 0`, an unsupported group size, unaligned pointers, or `N * K / 32 >= 2^31`. | | `onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.{h,cc}` | `ComputeFpAIntBDequantGemmMinM` resolves the threshold once per node (session option, else `ORT_FPA_INTB_DEQUANT_GEMM_MIN_M`, else 512 on SM 9.0 and 0 elsewhere) and decides eligibility in the constructor. `ComputeInternal` takes the scratch matrix from the EP allocator, dequantizes into it and runs one cuBLAS GEMM instead of the row-chunked CUTLASS call; the route profiles only in a session with CUDA graphs enabled, where it first tunes the CUTLASS `M` buckets that the later capture will use (see 7). When it declines, the CUTLASS path runs as before. | | `onnxruntime/contrib_ops/cuda/quantization/matmul_nbits_workspace_estimate.h`, `include/onnxruntime/core/framework/resource_accountant.h`, `onnxruntime/core/framework/resource_accountant.cc`, `onnxruntime/core/providers/cuda/cuda_execution_provider.cc` | The scratch is charged by both memory estimates. Level 1 becomes `max(CUTLASS workspace, N * K * sizeof(T))` for every admitted `M` at or above the threshold (`WorkspaceEstimatorConfig::cuda_fpa_intb_dequant_gemm_min_m` carries the option from the execution provider to the estimator; an invalid value leaves the node without an estimate, like an invalid chunk size), and Level 2 (`DeclareWorkspaceRequirements`) declares the same maximum. The Level-2 figure is exact when the route runs and an upper bound when CUDA graph capture or deterministic compute keep CUTLASS. | | `onnxruntime/test/contrib_ops/matmul_8bits_test.cc`, `onnxruntime/test/contrib_ops/matmul_nbits_fpa_intb_layout_test.cc` | `MatMulNBits.Fp16_Int8_LargeM_DequantGemm_Cuda` and `MatMulNBits.BFloat16_Int8_LargeM_DequantGemm_Cuda` run the route against the reference. `MatMulNBitsFpAIntBLayout.Int8LargeMDequantGemmRoundTripsExactly` compares the dequantized matrix with the original weights exactly, and `Int8LargeMDequantGemmSessionOption` checks the option, its precedence over an invalid environment variable and the rejection of invalid values (it skips in a build without the fpA_intB kernels, which do not parse the option). `Int8LargeMDequantGemmCudaGraphCapture` is described under 7. | | `onnxruntime/test/providers/cuda/test_cases/matmul_nbits_e2e_workspace_test.cc`, `onnxruntime/test/framework/resource_accountant_test.cc` | `MatMulNBitsWorkspace.Level1DequantGemmScratch`, `DequantGemmScratchAgreementAcrossLevels` (Level 1 and Level 2 charge the same scratch) and `GetCapabilityBudgetChargesDequantGemmScratch`; the factory test of the accountant also checks `cuda_fpa_intb_dequant_gemm_min_m`. | ##### 5. Session options (shared by 1, 3 and 4) | File | Change | |------|--------| | `include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h` | Six `ep.cuda.*` keys: `disable_fpa_intb_gemv_tc`, `disable_moe_gemv_tc`, `disable_qmoe_int3_tc`, `moe_gemv_tc_min_rows`, `qmoe_int3_tc_min_routes` and `fpa_intb_dequant_gemm_min_m`. The comment of `qmoe_int3_tc_min_routes` states that a value above 1 keeps constant INT3 weights in their raw layout (see 7). | | `onnxruntime/contrib_ops/cuda/cuda_session_options.h` (new) | `ResolveCudaSessionOption`: an explicit session option wins, otherwise the environment variable (an empty variable counts as unset), otherwise the default; the whole string is parsed and range-checked, and an invalid value fails. `ReadCudaSessionOption` applies it to the configuration of a kernel at construction, and the memory estimator calls `ResolveCudaSessionOption` directly so that it and the kernel agree. #33707 adds the same file (same blob) for its own switches. | ##### 6. Documentation | File | Change | |------|--------| | `docs/contrib_ops/cuda/moe_qmoe.md` | BF16 INT4 tensor-core GEMV (including the pipelined K loop), packed bf16x2 INT3 dequantization and the INT3 tensor-core GEMV (gates, switches with their `ep.cuda.*` session options, numerics, the exact scale window and activation range, the SM80 PTX requirement, measured gains), and the prepacked INT3 layout of #33740 (conditions, effect on the route threshold, staging of misaligned bindings, measured gain). Corrects the GEMV row gate from `(0, 8]` to `(0, 80]` (`kMaxProfiledExpandedRows`). | | `docs/contrib_ops/cuda/matmul_nbits.md` | Section 6.5 (SM 9.0 tensor-core INT8 GEMV: gate including `alpha == 1`, dequantization and its scale range, the SM80 PTX requirement, the tactic-cache key), the new section 6.6 (large-`M` INT8: gate, threshold, scratch accounting, CUDA-graph warmup, numerics), the `ORT_DISABLE_FPA_INTB_GEMV_TC` and `ORT_FPA_INTB_DEQUANT_GEMM_MIN_M` rows with their `ep.cuda.*` options in section 8 and the tests in section 9. | | `docs/contrib_ops/cuda/gemm_profiler_cache.md` | Tactic-selection version 5 and the `tensor_core_gemv` key column of the persisted cache. | ##### 7. Review follow-up (this PR and #33740) `97dcea1380` fixes the four open review threads of this PR and the two threads of #33740 whose concerns were still present when it merged. Each fix has a regression test that fails on the code before it (checked by reverting the source changes and rerunning the new tests). | Thread | Problem | Fix | Regression test | |--------|---------|-----|-----------------| | `dispatcher.h` | The dense INT8 tensor-core GEMV has no `alpha` multiplier (the scalar epilogue computes `alpha * acc + bias`), so a launch with `alpha != 1` returned the unscaled product. | `CanUseInt8TensorCore` requires `alpha == 1`; other launches use the scalar kernel. | `Int8TensorCoreGemvTest.NonUnitAlphaScalesTheOutput`, and an `alpha = 2` launch in `UnsupportedLaunchesUseScalarKernel`. | | `qmoe_int3_tc.cuh` | The INT3 tensor-core kernel rescaled its FP32 accumulators by 2^60 at the end, but the shifted accumulators underflow in FP32 first, so every result below about 2^-89 became 0. | The activations are multiplied by 2^60 when they are staged in shared memory (exact) and the accumulators are not rescaled. The kernel now also requires activations of magnitude below 2^68 (`ep.cuda.disable_qmoe_int3_tc=1` otherwise). | `QMoETest_Bf16Int3TensorCoreTinyActivations` (expert output 2^-96, constant and non-constant weights). | | `matmul_nbits.{h,cc}` | In a CUDA-graph session the large-`M` route served the warmup runs, so the CUTLASS tactic of the `M` bucket that the capture falls back to was never tuned, and the capture failed with "No valid fpA_intB MatMulNBits tactic". | A graph-enabled node tunes the CUTLASS buckets of the row chunk and of the remainder before it allocates the scratch copy (profiling cannot run during capture). | `MatMulNBitsFpAIntBLayout.Int8LargeMDequantGemmCudaGraphCapture` (`M = 4096`, capture on a user stream). | | `matmul_nbits_fpa_intb_layout_test.cc` | `Int8LargeMDequantGemmSessionOption` failed in a build without the fpA_intB kernels, which do not parse the option. | The test skips there, as `matmul_4bits_test.cc` does under `USE_FPA_INTB_GEMM`. | the test itself | | `moe_quantization.cc` (#33740) | An externally bound input or output that is not 16-byte aligned made the dispatcher choose a byte-safe kernel, which read the prepacked INT3 weights as raw ones and returned wrong values. | With prepacked weights, a misaligned input or output is staged through aligned scratch, so only the tensor-core kernel runs (an `ORT_RETURN_IF_NOT` checks it). | `QMoETest_Bf16Int3PrepackedMisalignedExternalBuffers` (12-70 rows, bindings at misaligned element offsets, compared with the aligned run and with a raw-layout session). | | `moe_quantization.cc` (#33740) | The prepack ran the tensor-core GEMV from the first route and so ignored an explicit `ep.cuda.qmoe_int3_tc_min_routes` or `ORT_QMOE_INT3_TC_MIN_ROUTES`. | An explicitly configured threshold above 1 disables the prepack, so the raw layout honors it for every launch; an explicit 1 or no setting keeps the prepack. | `QMoETest_Bf16Int3ExplicitTensorCoreMinRoutes` (threshold 256 with FC1 scales of 2^-70, which only the vectorized kernels accept). | | `qmoe_int3_tc.cuh` (`bd53cf38e2`) | Shifted subnormal FMA dequantization (`kDequantShift = 60`) overflowed BF16 scale multipliers at `scale = 2^60` (producing `NaN` on zero effective weights `code = 4`) and scaled activations by `2^60`. | Replaced with range-free packed `bf16x2` dequantization (`lop3.b32` + `HSUB2` + `HMUL2`) and unscaled activations in both `Int3TcCoalescedKernel` and `Int3TcKernel` (FC1 and FC2). | `QMoETest_Bf16Int3TensorCoreZeroWeightsExtremeScalesAndActivations` and `QMoETest_Bf16TensorCoreGemvExtremeScales` (`±2^60`). | | `int8_tc_gemv.cuh` (`d6813ee3ab`) | Dividing weights by `128` before rounding to BF16 flushed or distorted tiny weights below `2^-119` (e.g., `code = 129, scale = 2^-130` flushed to `0`; `code = 131, scale = 2^-127` rounded to `0.03125` instead of `0.0234375` with activation `2^120`). | Multiply `(q - 128) / 128` by `128` in BF16 before multiplying by `scale` (`Bf16x2Mul(Bf16x2Mul(Bf16x2Fma(bytes, 2^126, -1), 128), scale)`), removing `kUnscale` and matching `bf16_rn((q - 128) * scale)` bit-for-bit across all 65,024 finite BF16 scales (down to `2^-133`). | `Int8TensorCoreGemvTest.ExtremeScalesMatchReferenceWeights` and `WideScaleRangeMatchesScalarKernel` (scales down to `2^-133` with `1.0` and `2^120` activations across `M = 1, 8` and all 4 launch configurations). | `moe_qmoe.md`, `matmul_nbits.md` and the comment of the `qmoe_int3_tc_min_routes` key describe the new behavior. ##### Gates and switches Every switch is available as an environment variable and as a CUDA session option (`ep.cuda.*`, declared in `onnxruntime_session_options_config_keys.h`). It is resolved once when the kernel is created, so a change needs a new session, and an explicit session option takes precedence over the environment variable. The 0/1 switches accept only `0` and `1`, the row and route thresholds a positive decimal integer and `fpa_intb_dequant_gemm_min_m` an integer in `[0, 2147483647]`; any other value (including an empty session-option string) fails session initialization. | Kernel | Used when | Environment variable / session option | |--------|-----------|--------| | INT3 bf16x2 dequantization | BF16 activations, 3-bit weights, vectorized decode kernels | none (bit-identical) | | INT3 tensor-core GEMV | BF16, FC1 and FC2 3-bit, block 64, no zero points, no `router_weights`, `33 <= routes <= 256`, `hidden % 512 == 0`, `intermediate % 128 == 0`, `hidden <= 3072`, `intermediate <= 3072`, 16-byte-aligned weights, 4-byte-aligned scales, SM80+. Weights prepacked at session creation (constant weights and BF16 scales, see 1) run it from the first route. Exact for block scales in [2^-66, 2^55) and zero and for activations of magnitude below 2^68; a model outside these ranges (real models use scales of about 2^-13 to 2^0) needs the switch. | `ORT_DISABLE_QMOE_INT3_TC=1` / `ep.cuda.disable_qmoe_int3_tc`; `ORT_QMOE_INT3_TC_MIN_ROUTES` / `ep.cuda.qmoe_int3_tc_min_routes` (default 33; a value above 1 also keeps the weights in the raw layout, so the threshold holds for every launch) | | INT4 tensor-core GEMV | BF16, 4-bit, group 32, `33 <= expanded_num_rows <= 80`, `hidden % 512 == 0`, `inter % 128 == 0`, SM80+ | `ORT_DISABLE_MOE_GEMV_TC=1` / `ep.cuda.disable_moe_gemv_tc`; `ORT_MOE_GEMV_TC_MIN_ROWS` / `ep.cuda.moe_gemv_tc_min_rows` (default 33) | | Early PDL trigger | SM 9.0+, PDL successor | none (no effect for plain launches) | | Dense INT8 GEMV | SM 9.0, BF16 x symmetric INT8 (no zero point, bias or activation scale) in the SM80 interleaved layout, `alpha == 1`, `1 <= M <= 8`, `N % 16 == 0`, `K % 64 == 0`, `K % block_size == 0`, 16-byte-aligned activations and weights | `ORT_DISABLE_FPA_INTB_GEMV_TC=1` / `ep.cuda.disable_fpa_intb_gemv_tc` (part of the tactic-cache key) | | Large-`M` dequantize + cuBLAS | FP16 or BF16 activations, symmetric INT8 without zero points or bias in the SM80 layout (weights prepacked for SM 9.0 keep CUTLASS), `N % 16 == 0`, `K % 64 == 0`, a power-of-two `block_size >= 16` that divides `K`, `N * K * sizeof(T) <= 256 MiB` (the 1.3 GB `lm_head` keeps CUTLASS), `M` at or above the threshold, not under CUDA graph capture (a graph-enabled session first tunes the CUTLASS `M` buckets that the capture uses) and not with deterministic compute | `ORT_FPA_INTB_DEQUANT_GEMM_MIN_M` / `ep.cuda.fpa_intb_dequant_gemm_min_m` (default 512 on SM 9.0 and 0 elsewhere; 0 disables) | The tensor-core kernels above (not the large-`M` route, which needs no PTX) also require a loaded PTX target of at least SM80. If only older PTX is available (for example a `compute_75` image that the driver JIT-compiles for a newer GPU), their `__CUDA_ARCH__ >= 800` bodies are empty stubs that would silently leave the output unwritten, so each launcher checks `cudaFuncGetAttributes(...).ptxVersion >= 80` and otherwise uses the scalar or vectorized kernel. ##### Numerics - INT3 bf16x2 dequantization is bit-identical to the scalar chain (generated-token hashes identical in 5 / 5 runs). - The QMoE tensor-core kernels keep the per-weight BF16 rounding of the scalar kernels and differ only in the FP32 accumulation order (about one BF16 ulp). The INT4 pipelining and the INT3 rewrite are bit-identical to the tensor-core kernels they replace (byte-compared outputs in standalone harnesses; generated-token hashes identical in 18 + 12 decode cells x 5 runs). - The INT3 tensor-core kernel decodes each weight to the once-rounded `(code - 4) * scale * 2^-60` and multiplies the activations by 2^60 when it stages them in shared memory (exact), so the FP32 accumulators hold the true dot products. It is exact for block scales in [2^-66, 2^55) and zero and for activations of magnitude below 2^68. The first version of the kernel accepted every scale, so the rewrite keeps a window far wider than the scales of real models; the harness found no mismatch against the previous kernel for scales multiplied by 2^0, 2^+-40 and 2^+-56, and almost every output differs outside the window (2^60, 2^-70), which shows that the check is sensitive. Before the review follow-up (7) the kernel rescaled the accumulators at the end instead, which returned 0 for every result below about 2^-89 (`QMoETest_Bf16Int3TensorCoreTinyActivations`). - The coalesced INT3 prepack of #33740 changes only where the kernel reads weights and scales: `Int3TcCoalescedKernel` is bit-exact with `Int3TcKernel` (measured in #33740), so it adds no numerical difference. End to end, `engine_behaviours.py` passed 27 / 27 on each of the RC5 `all_int4`, `top16` and `top08` packages (#33740). - The dense INT8 GEMV dequantizes to the same weight value as the scalar kernel for every weight that is finite in BF16 when the scale has magnitude >= 2^-119. A standalone harness compared all 256 code bytes with 210 scales of magnitude 2^-119 to 2^110 (three code-to-column assignments, all four launch configurations) against the BF16 rounding of the exact product, with no mismatch, and `OneHotActivationsMatchScalarKernel` and `ExtremeScalesMatchReferenceWeights` check it in the tests. It accumulates in FP32 in a fixed order, where the scalar kernel accumulates in BF16 chains. Compared with a version that pre-scales the weights by a fixed power of two, the two packed FMAs cost +0.6-2.4% kernel time (M = 1-8, N = 2560-16480, K = 2560-6144, group 32-128) and gave bit-identical outputs on realistic scales. - The large-`M` route produces the same per-weight values as the GEMV kernels and the CUTLASS converter (one rounding of the exact `(q - 128) * scale`). Only the FP32 accumulation order of the GEMM differs from CUTLASS, so a near-tie in the first generated token can flip; of the prompts of the TTFT benchmark this happened only at 2048 tokens, deterministically. - The early PDL trigger changes no result (generated tokens identical in 45 / 45 decode runs, measured with both triggers). #### Testing In-tree CUDA 13.0 Release build, `sm_90a`, one H200 (SM 9.0), on this branch at `97dcea1380` (`onnxruntime_provider_test`; the build has `onnxruntime_USE_FPA_INTB_GEMM=ON` and the flash, memory-efficient and lean attention kernels and the INT4 and FP8 KV caches off): | Command (`--gtest_filter=`) | Result | |---|---| | `*Bf16Int4TensorCoreGemvParity*:*Bf16Int3TensorCoreGemvParity*:*QMoETensorCoreSessionOptions*:*Bf16TensorCoreGemvExtremeScales*` | 4 run, **4 passed** | | `*Bf16Int3TensorCoreTinyActivations*:*Bf16Int3ExplicitTensorCoreMinRoutes*:*Bf16Int3PrepackedMisalignedExternalBuffers*` (new, see 7) | 3 run, **3 passed** | | `MoETest.*:*QMoE*:*Moe*:*MoE*` | 134 run, **121 passed**, 13 skipped (CPU-EP, Mixtral and plugin-only cases), 0 failed; includes the seven tests above | | `MatMulNBits*:MatMulNBitsFpAIntBLayout*:*8bits*:*Int8LargeM*` | 112 run, **101 passed**, 11 skipped (CPU-only and SM 12.1-only cases), 0 failed; includes `MatMulNBitsFpAIntBLayout.TensorCoreSessionOptions`, the two `LargeM_DequantGemm` tests and the three `Int8LargeMDequantGemm*` layout tests (`Int8LargeMDequantGemmCudaGraphCapture` is new) | | `CUDA_EP_Unittest.All`, run as `GTEST_FILTER=-GroupQueryAttentionWorkspaceEstimateTest.ValidatesQuantizedCacheMetadata ./onnxruntime_provider_test --gtest_filter=CUDA_EP_Unittest.All` (see the note below) | 516 run, **505 passed**, 11 skipped (FP8 native GEMV, external-data loader and the flash, lean and packed-attention estimator tests that need kernels this build omits), 0 failed; includes the 8 `Int8TensorCoreGemvTest.*` tests (one new), `GemmTacticCacheTest.TensorCoreSettingsHaveSeparateProfiles` and the 45 `MatMulNBitsWorkspace.*` tests | | `RealAccountantTest.*:ResourceAccountantTest.*` (`onnxruntime_test_all`) | 28 run, **28 passed** on the tree before the follow-up; not rerun, because the follow-up changes only a comment of `onnxruntime_session_options_config_keys.h` that these tests include | - `GroupQueryAttentionWorkspaceEstimateTest.ValidatesQuantizedCacheMetadata` is excluded from the `CUDA_EP_Unittest.All` run because it fails in this build, whose flash attention is off: `group_query_attention_workspace_estimate.cc` offers a backend for an INT4 quantized KV cache only under `USE_FLASH_ATTENTION` (the XQA backend takes 8-bit caches only), and only the INT4 iteration of the test fails. This PR changes neither that file nor its test. The 22 other tests of the suite pass. - Each new test fails on the code before the follow-up. I reverted the source changes, rebuilt, reran the new tests and then restored the sources: `Int8TensorCoreGemvTest.NonUnitAlphaScalesTheOutput` and the `alpha = 2` check of `UnsupportedLaunchesUseScalarKernel` (the tensor-core kernel returned 1088 where 2192 is expected); `Int8LargeMDequantGemmCudaGraphCapture` ("No valid fpA_intB MatMulNBits tactic for M=4096, N=256, K=512. The M bucket was not profiled before CUDA graph capture"); `QMoETest_Bf16Int3TensorCoreTinyActivations` (0 where 2^-96 is expected); `QMoETest_Bf16Int3ExplicitTensorCoreMinRoutes` (the explicit threshold of 256 had no effect, so the tensor-core kernel ran outside its scale window and returned wrong values); `QMoETest_Bf16Int3PrepackedMisalignedExternalBuffers` (first output element -0.042 where 0.326 is expected). The skip of `Int8LargeMDequantGemmSessionOption` has no failing state in this build, because it only matters in a build without `USE_FPA_INTB_GEMM`. - Under `nsys` the session-option tests launch the intended kernels. QMoE (`QMoETensorCoreSessionOptions`): `MoeGemvTcInt4Kernel` and `Int3TcKernel` (FC1 and FC2) with the switches at 0, and the scalar `moe_gemv_kernel`, `moe_gemv_interleaved_swiglu_kernel`, `FusedFc1FastVecKernel` and `Fc2RouteFastVecKernel` at 1. Dense INT8 (`MatMulNBitsFpAIntBLayout.TensorCoreSessionOptions`): 49 `Int8TcGemvKernel` launches in the session with the switch at 0, and the scalar `fpA_intB_gemv::kernel` with no tensor-core launch at 1 (measured again on the final tree, after the `alpha == 1` gate of 7: 49 tensor-core and 49 scalar launches in the test process). Large `M` (`Int8LargeMDequantGemmSessionOption`): with the threshold at 512, set as a session option or as an environment variable, `DequantizeInt8InterleavedKernel` followed by a cuBLAS GEMM (`nvjet_*`); at 0 the CUTLASS `GemmFpAIntB`. - `lintrunner` is clean on all 43 changed files. The CUDA translation units built from the changed files (`moe_gemv.cu`, `moe_kernels.cu`, `qmoe_kernels.cu`, `fpA_intB_gemv.cu` and the six `dispatcher_{bf16,fp16}_{int2,int4,int8}.cu`) also compile for `sm_75`, `sm_86` and `sm_120` (30 / 30 compilations before the follow-up; `moe_kernels.cu` for `sm_120a`, which its existing FP4 conversions require). After the follow-up I compiled again the eight translation units that include the headers it changes (`dispatcher_{bf16,fp16}_{int2,int4,int8}.cu`, `moe_gemv.cu` and `qmoe_kernels.cu`) for the same three architectures: 24 / 24, with no warning. The Hopper dispatchers are built for `90a` only, as in the normal build. - The PTX guards (`cudaFuncGetAttributes(...).ptxVersion >= 80`) have no unit test, because they need an executable that carries only `compute_75` PTX. On the H200 (the driver JIT-compiles the PTX for SM 9.0) such an executable shows that the dense INT8 (four configurations), QMoE INT4 (both variants) and INT3 (FC1 and FC2) launchers decline without launching. - Offline tuning tool (`ORT_RUN_CUDA_TACTIC_CACHE_TESTS=1 python onnxruntime/test/python/onnxruntime_test_python_fpa_intb_tune.py` against the Python package of the 13-commit head `c7e138fe27`, one H200): 14 run, **14 passed**, including `test_fresh_process_reuse_repair_and_lock_contention`, which checks the cache header version and key columns end to end. The later commits do not touch the tool, its test, the profiler or the cache format. - CI: the SM 9.0-only `Int8TensorCoreGemvTest.*` tests and `MatMulNBitsFpAIntBLayout.TensorCoreSessionOptions` skip on `sm_86` runners. The QMoE parity and session-option tests need SM80+, so they run there. The quick-build legs (`onnxruntime_QUICK_BUILD=ON`: Windows CUDA and both CUDA Plugin CIs) do not build the native BF16 INT4/INT8 QMoE runner (`moe_quantization.cc`), so there `QMoETest_Bf16Int4TensorCoreGemvParity` is compiled out, and `QMoETensorCoreSessionOptions` and `QMoETest_Bf16TensorCoreGemvExtremeScales` cover INT3 only; the INT3 parity test still runs. The three new INT3 QMoE tests of 7 need SM80+ and are also built in the quick-build legs, because the INT3 runner is. `Int8LargeMDequantGemmCudaGraphCapture` needs SM80+ and `USE_FPA_INTB_GEMM`, and sets `ep.cuda.fpa_intb_dequant_gemm_min_m` to 512 itself, so it also runs on `sm_86`. `Int8LargeMDequantGemmSessionOption` skips in a build without `USE_FPA_INTB_GEMM`. #### Checklist - [x] Tests added/updated and passing - [x] No breaking API changes. Behavior changes are described under Gates and switches and Numerics: on SM 9.0 the dequantize + cuBLAS route replaces CUTLASS for `M >= 512` (`ORT_FPA_INTB_DEQUANT_GEMM_MIN_M=0` restores it), the INT3 tensor-core kernel is exact only for block scales in [2^-66, 2^55) and activations of magnitude below 2^68 (`ep.cuda.disable_qmoe_int3_tc=1` for others), an explicitly configured `qmoe_int3_tc_min_routes` above 1 keeps constant INT3 weights in the raw layout instead of the prepacked one, and persisted tactic caches of version 4 are no longer reused (`kTacticSelectionVersion` 5) - [x] Documentation updated (if applicable) ### Motivation and Context At `S = 8` with top-10 routing a Qwen3.8-Flash-Next decode step routes about 80 rows over 16-25 experts per layer. The scalar QMoE GEMV dequantized every weight once per routed row and was instruction-issue bound. The 97 dense INT8 projections of a decode cycle (block 32, symmetric) ran on the fpA_intB GEMV at about 11 thread instructions per weight, 62% FMA-pipe utilization and only 28-37% of the H200 DRAM bandwidth. The first versions of the tensor-core MoE kernels were latency bound: `ncu` showed the INT4 kernel waiting on `long_scoreboard` for 60% of its warp time (40% of the DRAM peak), and the INT3 FC1 kernel at 167 registers (2 CTAs per SM) with issue slots 38% busy, `long_scoreboard` 27%, `wait` 20% and the cross-warp K-split barrier 9%. In prefill, the CUTLASS mixed-input GEMM converts every weight again for each row chunk and reaches about 300 TFLOPS on H200, where a BF16 cuBLAS GEMM reaches about 730 TFLOPS: 826 `GemmFpAIntB` launches took 138.6 ms (22%) of an 8192-token chunk. #### Performance on NVIDIA H200 (SM 9.0a, CUDA 13.0) Measured on the Qwen3.8-Flash-Next packages (`hidden = 2560`, `intermediate = 640`, 512 experts, top-10; `all_int4` is the all-INT4 package, `top16` and `top08` are the two INT3-mixed packages), B = 1, DFlash2 `S = 8` decode, on a stack that also contains #32126 and #33650. Decode is in ms per cycle, an A/B of consecutive providers on the same packages (single rounds unless noted). | Change | `all_int4` | `top16` | `top08` | |---|---:|---:|---:| | INT3 bf16x2 dequantization (2K / 8K) | unchanged | -5.1% / -4.8% | -6.3% / -6.2% | | INT4 tensor-core MoE GEMV (2K) | -15% | -5% | -2.3% | | INT3 tensor-core MoE GEMV vs. the bf16x2 step (2K / 8K) | n/a | -5.8% / -5.1% | -6.8% / -6.7% | | Dense INT8 tensor-core GEMV (2K / 8K / 32K; mean of 2 rounds) | -3.8% / -4.4% / -3.8% | -5.6% / -5.5% / -4.6% | -5.1% / -4.9% / -6.9% | | Software-pipelined INT4 MoE GEMV (2K / 8K / 32K; mean of 2 rounds) | -1.8% / -2.0% / -1.9% | -0.5% / -0.9% / -1.1% | -0.6% / +0.3% / -0.4% | | Rolled-pipeline INT3 MoE GEMV (2K / 8K / 32K; mean of 2 rounds) | n/a | -4.3% / -4.4% / -4.7% | -5.6% / -5.4% / -5.8% | The gain of the pipelined INT4 kernel scales with the number of INT4 MoE launches per cycle (96 / 32 / 16), and the `+0.3%` cell is a noisy round (-0.5% and +1.1%). The dense INT8 numbers were measured with an earlier dequantization that folded a fixed power-of-two shift into the scale. The range-safe one in this PR (see Numerics) takes 0.6-2.4% longer in a standalone kernel benchmark (cold weights, `M = 1-8`, `N = 2560-16480`, `K = 2560-6144`, group size 32-128) and gave bit-identical outputs on realistic scales; the end-to-end numbers were not re-measured with it. The early PDL triggers (here and in #33707) shorten a cycle by 1.4-1.7%. Cumulative 2K decode of the first 13 commits, ms per cycle (`all_int4` / `top16` / `top08`): | Stack | 2K decode | |---|---| | Without this PR | 14.69 / 15.13 / 15.31 | | + bf16x2 INT3 dequantization, INT4 / INT3 tensor-core MoE GEMV, early PDL | 12.31 / 12.72 / 12.90 | | + dense INT8 tensor-core GEMV = **the first 13 commits of this PR** | **11.83 / 12.04 / 12.21** | The second row also changes the packages' prefill chunk size and QMoE row tile (package configuration, no code), which does not affect decode. It includes the early PDL trigger of the gate-mul kernel, which is in #33707 now. The INT8-KV DSA prefill step that was between these two rows (prefill only, decode 12.30 / 12.76 / 12.86) is not part of this PR. The three later commits were measured as separate steps (rows above), so their gains are not added to this cumulative figure. Kernel level (H200): | Kernel | Before | After | |---|---|---| | INT3 FC1 per layer, `S = 8`, first tensor-core version | 53.9 us (`FusedFc1FastVecKernel`) | 43.1 us | | INT3 FC2 per layer, `S = 8`, first tensor-core version | 33.6 us (`Fc2RouteFastVecKernel`) | 19.1 us | | INT3 FC1 (`n = 1280`, `k = 2560`), rolled-pipeline rewrite (standalone, 80 routes over 56 experts, L2 flushed, including the ~5.5 us launch floor) | 49.06 us | 30.37 us (-38.1%) | | INT3 FC2 (`n = 2560`, `k = 640`), rewrite (same harness) | 24.16 us | 22.37 us (-7.4%) | | INT4 FC1, pipelined K loop (standalone, 80 rows over 29 experts, L2 flushed, including the ~5.5 us launch floor) | 29.2 us | 23.2 us (-20%) | | INT4 FC2, pipelined K loop (same harness) | 19.0 us | 16.1 us (-15%) | | Dense INT8 `in_proj` (`N = 16480`, `K = 2560`, `M = 8`, cold weights) | 24.9 us | 13.0 us | | Dense INT8 `out_proj` (`N = 2560`, `K = 6144`) | 12.7 us | 9.7 us | | Dense INT8 `N = 13952`, `K = 2560` | about 20 us | 11.8 us | | All dense INT8 projections of a decode cycle | 1.84 ms | 1.20 ms | | Dense INT8 `in_proj`, `M = 512` / `M = 8192`, CUTLASS vs. dequantize + cuBLAS (standalone, BF16, group 32) | 128 us / 2050 us | 110 us / 944 us (the dequantization is 50 us of it) | | Dense INT8 `out_proj`, `M = 8192` (same setup) | 723 us | 350 us (dequantization 22 us) | TTFT of the large-`M` route in the model (A/B of this change alone, `all_int4` / `top16` / `top08`, median of 5 runs; decode is unchanged). The route is taken from `M = 512`, so the 128- and 256-token columns are run-to-run noise: | Package | 128 | 256 | 512 | 1024 | 2048 | 8192 | 32768 | |---|---:|---:|---:|---:|---:|---:|---:| | `all_int4` | -1.4% | 0.0% | -4.0% | -5.8% | -8.5% | **-11.8%** | **-11.3%** | | `top16` | -0.7% | 0.0% | -1.9% | -4.9% | -7.4% | **-11.1%** | **-10.7%** | | `top08` | -0.2% | -0.5% | -2.6% | -5.0% | -7.3% | **-11.5%** | **-10.6%** | `all_int4` absolute: 2048 tokens 188.6 -> 172.6 ms, 8192 tokens 650.9 -> 574.0 ms, 32768 tokens 2754.8 -> 2442.5 ms. In an 8192-token chunk (`nsys`, `all_int4`) the critical path drops from 628.8 to 551.6 ms: the 826 CUTLASS `GemmFpAIntB` launches (138.6 ms) are replaced by 97 `DequantizeInt8InterleavedKernel` launches (3.5 ms), cuBLAS GEMMs (31.5 ms for the dense layers) and 2.2 ms of remaining `GemmFpAIntB` (the `lm_head` and small layers). The threshold was swept over 128 / 256 / 512 on the 128-512 token prompts: 128 was 3-5% slower, 256 0-3% slower and 512 2.5-3.5% faster than CUTLASS, so 512 is the default (the extra pass reads 1 B and writes 2 B per weight and the GEMM reads 2 B, which only pays off once the GEMM is large). For the whole stack that contains this PR, #33707, #32126, #33650 and the package configuration, the median over 4 rounds against the previous release candidate is -32.8 / -36.0 / -36.9% TTFT at 8192 tokens, -32.3 / -35.2 / -36.1% at 32768 tokens, and -27.6 / -24.6 / -24.2% decode ms per cycle at 2K (`all_int4` / `top16` / `top08`); these numbers cannot be attributed to a single PR. ##### Final-head isolated benchmarks (`d6813ee3ab`, other PRs held constant) To isolate this PR's final diff from #33707, #33714, #33650, #32126, and the package configuration, two CUDA plugin providers were built from the same companion stack (`main` `0a12bc0480` + #33714 `ebde3f13e5` + #33707 `204d6593c0` + #33650 / #32126 `78d59a59c7`, run with GenAI #2808 `5d561a2f5` and `prefill.chunk = 8192` / `qmoe_row_tile_size = 8192`): `base` (without #33689, md5 `4dc3654493ea4cf8e7b847162d5ab968`) and `final` (with #33689 at `d6813ee3ab`, md5 `d8f2bdb16e4db21966f10a7822da054e`), and benchmarked over 2 rotated GPU rounds on H200 (5 runs per cell per round): | Package | Stack | TTFT 128 | TTFT 256 | TTFT 512 | TTFT 1024 | TTFT 2048 | TTFT 8192 | TTFT 32768 | Decode 2K / 8K / 32K (ms/cycle) | |---|---|---:|---:|---:|---:|---:|---:|---:|---| | `all_int4` | `base` (no #33689) | 0.0463 s | 0.0465 s | 0.0703 s | 0.1120 s | 0.1890 s | 0.6484 s | 2.7609 s | 13.82 / 13.96 / 14.02 | | | `final` (#33689 `d6813ee3ab`) | 0.0460 s | 0.0468 s | 0.0687 s | 0.1068 s | 0.1733 s | 0.5792 s | 2.4442 s | 10.85 / 10.99 / 11.24 | | | **delta (`final` vs `base`)** | -0.7% | +0.7% | **-2.3%** | **-4.6%** | **-8.3%** | **-10.7%** | **-11.5%** | **-21.4% / -21.3% / -19.9%** | | `top16` | `base` (no #33689) | 0.0509 s | 0.0582 s | 0.0840 s | 0.1277 s | 0.2063 s | 0.6682 s | 2.8386 s | 14.29 / 14.66 / 14.59 | | | `final` (#33689 `d6813ee3ab`) | 0.0510 s | 0.0584 s | 0.0817 s | 0.1212 s | 0.1889 s | 0.5931 s | 2.5220 s | 10.70 / 10.91 / 11.16 | | | **delta (`final` vs `base`)** | +0.3% | +0.4% | **-2.8%** | **-5.0%** | **-8.4%** | **-11.2%** | **-11.2%** | **-25.1% / -25.6% / -23.5%** | | `top08` | `base` (no #33689) | 0.0524 s | 0.0610 s | 0.0872 s | 0.1312 s | 0.2086 s | 0.6723 s | 2.8502 s | 14.42 / 14.51 / 14.92 | | | `final` (#33689 `d6813ee3ab`) | 0.0527 s | 0.0616 s | 0.0865 s | 0.1267 s | 0.1956 s | 0.5992 s | 2.5440 s | 10.71 / 10.87 / 11.23 | | | **delta (`final` vs `base`)** | +0.6% | +1.0% | **-0.7%** | **-3.4%** | **-6.2%** | **-10.9%** | **-10.7%** | **-25.7% / -25.1% / -24.7%** | Per-feature ablations on `final` (`d6813ee3ab`, 2 rotated GPU rounds): - **Dense INT8 tensor-core GEMV (`ORT_DISABLE_FPA_INTB_GEMV_TC=1` -> `0`, range-free `Dequant`)**: standalone H200 `M = 8` latencies are `24.9 -> 14.3 us` (`in_proj` `16480 x 2560`), `12.7 -> 10.6 us` (`out_proj` `2560 x 6144`), and `~20.0 -> 12.6 us` (`13952 x 2560`). In-model decode `ms/cycle` at 2K / 8K / 32K drops by **-4.5% / -4.7% / -4.3%** (`11.37 -> 10.85`, `11.54 -> 10.99`, `11.74 -> 11.24` ms) on `all_int4`, **-5.3% / -4.1% / -5.2%** (`11.29 -> 10.70`, `11.38 -> 10.91`, `11.78 -> 11.16` ms) on `top16`, and **-5.1% / -5.3% / -4.9%** (`11.29 -> 10.71`, `11.47 -> 10.87`, `11.81 -> 11.23` ms) on `top08`. - **Large-`M` INT8 dequantize + cuBLAS (`ORT_FPA_INTB_DEQUANT_GEMM_MIN_M=0` -> `512`)**: TTFT at `512 / 1024 / 2048 / 8192 / 32768` tokens drops by **-2.6% / -3.5% / -6.8% / -11.0% / -11.2%** (`all_int4`), **-3.8% / -5.4% / -7.4% / -11.8% / -11.2%** (`top16`), and **-0.6% / -3.7% / -6.5% / -11.2% / -10.8%** (`top08`). - **INT3 tensor-core decode (`bd53cf38e2` range-free `bf16x2` dequantization, H200 standalone, 80 routes over 56 experts, L2 flushed)**: FC1 (`1280 x 2560`) takes `33.02 us` raw (`Int3TcKernel`) and `30.96 us` prepacked (`Int3TcCoalescedKernel`); FC2 (`2560 x 640`) takes `23.76 us` raw and `21.78 us` prepacked. #### Accuracy MMLU-Pro, 800 questions x 3 packages = 2400 items, greedy decoding through the full engine, `max_new = 8192`. The kernels of this PR change only the accumulation precision or order: the dequantizations are bit-exact (INT3 bf16x2, dense INT8 GEMV, large-`M` dequantization), and the INT4 pipelining and the INT3 rewrite are bit-identical to the kernels they replace, and the coalesced INT3 prepack of #33740 is bit-exact with the INT3 kernel it feeds, so no systematic bias is expected. The measurements are of the accumulated stack, not of this PR alone: - The first-round stack (the 13-commit version of this PR, which included the INT8-KV prefill kernel that is now in #33707, plus other changes of the stack) against the previous release candidate: **-0.67 pp** pooled (95% CI -1.36 .. +0.03 pp, exact McNemar p = 0.076); no package differs at the 5% level. The last step alone, the dense INT8 GEMV against the stack before it: 23 / 18 discordant items (p = 0.53). - With the large-`M` route and the split-K DSA decode kernel of #33707 added: 668 / 678 / 667 correct answers against 665 / 676 / 670 for the previous release candidate, pooled **+0.08 pp** (95% CI -0.60 .. +0.77 pp, p = 0.91); no package differs at the 5% level. This run used the stack before the INT4 pipelining, the INT3 rewrite and the sigmoid change, which are bit-exact and were not run separately. - MMLU-Pro covers the large-`M` route only marginally: 27 of its 800 prompts have 512 or more tokens, and 6 of the 61 changed answer letters belong to them. - These MMLU-Pro runs predate the review follow-up of #33707, which keeps more FP16 bits of the INT8-KV attention probabilities below 2^-14 and is therefore not bit-exact. On the final heads of this PR, #33707, #33714, #32126 and #33650, `engine_behaviours.py` passes on the three packages, and TTFT and decode ms per cycle are within the noise of the stack before it (details in the Accuracy section of #33707). The generated tokens differ in near-tie cases (4 of 126 first tokens of the TTFT prompts), and MMLU-Pro was not repeated for it. A drift of up to about 1.4 pp cannot be excluded with 2400 items. ##### Final-head long-prompt accuracy verification for the large-`M` route (`d6813ee3ab`, `min_m = 512` vs `0`) 1. **Kernel execution confirmation (`ORT_FPA_INTB_DEBUG=1`)**: across prompt lengths `256, 512, 1024, 2048, 4096, 8192, 32768` on `final`, all **97** eligible dense INT8 `MatMulNBits` layers per prefill pass launch `kernel=dequant+cuBLAS` at `512, 1024, 2048, 4096, 8192` tokens (`388` launches across four 8192-token chunks at `32768`, and `0` launches at `256` or when `ep.cuda.fpa_intb_dequant_gemm_min_m = 0`, where all 97 run CUTLASS `GemmFpAIntB`). 2. **Node-level accuracy on real Qwen3.8-Flash-Next weights (`M = 512, 1024, 2048, 4096, 8192`)**: probing all `93.6M` weights of the `16480 x 2560`, `2560 x 6144`, and `13952 x 2560` nodes with `M = 512` identity slices gives `0` mismatches between `dequant+cuBLAS` and CUTLASS. Against an FP64 reference of those exact BF16 weights, `dequant+cuBLAS` has `1.11e-4 .. 1.87e-4` relative RMS error across all shapes and `M = 512..8192`: its outputs are **bitwise identical (`0.0` difference)** to CUTLASS whenever CUTLASS selects `split_k = 1` (`16480 x 2560` and `13952 x 2560` at all `M`, and `2560 x 6144` at `M = 4096, 8192`), and are closer to the FP64 reference than CUTLASS on `2560 x 6144` at `M = 512, 1024, 2048` (`1.83e-4` vs `2.15e-3 .. 2.98e-3` relative RMS, cosine similarity `>= 0.9999956`), where CUTLASS reduces split-K partial sums in BF16. 3. **End-to-end long-prompt comparisons (`ep.cuda.fpa_intb_dequant_gemm_min_m = 512` vs `0`, rest of stack unchanged)**: - **55 natural-text prompts per package (`10` each at `512, 1024, 2048, 4096, 8192` tokens and `5` at `32768` tokens = `165` pairs across `all_int4`, `top16`, `top08`)**: first generated token agrees in **`158 / 165` (`95.8%`)** runs (`53/55` `all_int4`, `53/55` `top16`, `52/55` `top08`; `100%` at `4096` and `32768` where all layers use `split_k = 1`), and full 32-token greedy sequences are bitwise identical in `104 / 165` (`63.0%`). - **27 natural MMLU-Pro prompts with `>= 512` tokens (`512..1150` tokens) x 3 packages (`81` evaluations)**: `min_m = 512` scores **`54 / 81` (`66.67%`: `17/27` `all_int4`, `17/27` `top16`, `20/27` `top08`)** vs **`51 / 81` (`62.96%`: `17/27`, `16/27`, `18/27`)** for `min_m = 0`, with **`80 / 81` (`98.8%`)** first-token agreement. - **30 long-context MMLU-Pro prompts per package padded with natural-text history to exact lengths `1024, 2048, 4096` (`10` per length = `90` evaluations)**: `min_m = 512` scores **`57 / 90` (`63.3%`: `21/30`, `18/30`, `18/30`)** vs **`56 / 90` (`62.2%`: `19/30`, `19/30`, `18/30`)** for `min_m = 0`, with **`81 / 90` (`90.0%`)** answer-letter agreement and **`84 / 90` (`93.3%`)** first-token agreement. - **`engine_behaviours.py`**: **27 / 27 PASS** on all three packages (`all_int4`, `top16`, `top08`) with `final` (`d6813ee3ab`). #### Related PRs and review guide - #33707 is the sibling PR (INT8-KV `DynamicSparseAttention` prefill and decode kernels, the gate-mul early PDL trigger and the hyper-connection sigmoid), split from this PR by operator. The changes have no kernel in common. `cuda_session_options.h` is added by both PRs with identical content, and the two edits of `onnxruntime_session_options_config_keys.h` are in different places, so the PRs merge in either order; a trial merge of both branches is clean. - Independent of #32126 (SwiGLU operator and fusion) and #33650 (RouterGateUpProj): no file overlap, and trial merges with both heads and with `main` are clean. The decode graph measured above contains both. - The PR has 43 files and 20 commits, above the usual review size. It is updated only by fast-forward pushes (the first review round saw the 13-commit head `c7e138fe27`), so that the existing review threads and commit hashes are kept. - Commits 1-13 are the previous head: commits 1-3 and 7 (QMoE), commit 4 (early PDL), commit 5 (INT8-KV DSA prefill), commits 6 and 8 (dense INT8 GEMV; commit 8 reads its environment variable with `_dupenv_s` on Windows), commit 9 (docs), commit 10 (review follow-up: range-safe dense GEMV dequantization, the FP32 fallback of the INT8 prefill kernel, the PTX guards and the session options with their tests and docs), commits 11-13 (CI follow-ups: an MSVC narrowing warning in a DSA test, the BF16 INT4 QMoE tests compiled out in quick builds, and `fpa_intb_tune.py`, its test and the cache layout doc for tactic-selection version 5). - Commit 14 (`b32d9eac54`) moves everything `DynamicSparseAttention`, hyper-connection and gate-mul out to #33707. It keeps the early PDL trigger of the fpA_intB GEMV kernels and `cuda_session_options.h`. The hunks of commits 4, 5, 9, 10 and 11 for the moved files are reverted, so the net diff of this PR contains none of them. The earlier review threads on the prefill kernel and its test were answered with the fixes of commit 3 of #33707. - Commits 15-17 are new: the large-`M` dequantize + cuBLAS route, the INT4 K-loop pipelining and the INT3 rewrite, each with its own measurements in its commit message. They can be reviewed separately. - Commits 18-20 come after that: `2b5ceeb045` relaxes the zero-point tolerance of the BF16 large-`M` `MatMulNBits` test (the zero-point variants run CUTLASS), `a8360fcbfe` is #33740 (coalesced INT3 prepack, the prefill repack that consumes it and the compact INT4 FC2 GEMV), and `97dcea1380` is the review follow-up under 7. - The session-option helper and the config keys are shared by the QMoE, dense GEMV and large-`M` groups, and `fpA_intB_gemv/dispatcher.h` is shared by the early-PDL and dense GEMV groups, in different hunks. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Tianlei Wu (tianleiwu)
added a commit
that referenced
this pull request
Nov 1, 2026
… branch-free hyper-connection sigmoid (#33707) ### Description This PR adds tensor-core CUDA kernels for the INT8-KV-cache `DynamicSparseAttention` of Qwen3.8-Flash-Next (12 full-attention layers, 24 query / 2 KV heads, head size 256): a prefill kernel and a split-K decode kernel. It also triggers the programmatic-dependent-launch (PDL) successor early in `ProjectedSigmoidGateMulKernel` and makes the sigmoid of `HyperConnectionLowRankMix` branch free. It is the second half of #33689, which was split by operator: #33689 keeps the tensor-core `MatMulNBits` (fpA_intB) and QMoE / MoE GEMM kernels, and this PR carries the DynamicSparseAttention, gate-mul and hyper-connection changes. The two PRs are independent. The only file both add is `cuda_session_options.h`, and this PR's version is a superset of the one in #33689 (see "Related PRs and review guide"). **Both INT8 tensor-core kernels are opt-in.** `ep.cuda.dsa_prefill_int8_tc` / `ORT_DSA_PREFILL_INT8_TC` and `ep.cuda.dsa_decode_int8_tc` / `ORT_DSA_DECODE_INT8_TC` default to `0`, which runs the previous kernels bit for bit, and the two switches are independent. The end-to-end accuracy and speed results were measured for an earlier revision of the kernels, before the numerics changes of the review, so they are preliminary. The kernels stay off by default until the final revision is qualified end to end (see "Accuracy"). At the operator level on an H200 (`benchmark_dynamic_sparse_attention.py`, added by this PR) the tensor-core decode kernel is 3.3x (1 row) to 5.6x (8 rows) faster than the CUDA-core kernel it replaces, and the prefill kernel is 2.7-3.0x faster than the BF16-dequantizing kernel. Against an FP64 reference the decode error stays within 3% of the previous kernel and the prefill error is 1.6-2.5 times smaller. A CTA whose inputs break the error bounds of the FP16 path is evaluated on a slower FP32 CUDA-core path, which can take longer than the previous kernels (see "Numerics"); the Qwen3.8-Flash-Next layers stay far from its conditions. In earlier-revision end-to-end runs the kernels lowered TTFT by 4.0-9.5% and the `S = 8` decode cycle by 4.6-6.8% (preliminary), and the branch-free sigmoid lowered TTFT by 0.3-1.1%, which is about the size of the repeat noise. #### Summary of Changes ##### 1. DynamicSparseAttention: tensor-core prefill for the INT8 KV cache | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh` (new) | `DynamicSparseAttentionInt8TcKernel<kSplitK>` (`kSplitK = false` for prefill): grid (rows, KV heads), 128 threads, 64-candidate tiles, GQA group <= 16 padded to the `m16n8k16` M dimension. The gathered K/V rows stay INT8 in shared memory (`cp.async`, double buffered, 3-deep index ring) and the `ldmatrix` fragments are converted exactly to FP16 in registers, for FP16 `mma` with FP32 accumulation. `key_scale` is folded into the query once per CTA (`Q' = fp16(Q * key_scale)`) and `value_scale` is applied to the output once, after the accumulator has been divided by 2^15 (so a tiny `value_scale` cannot underflow). Scores and the running maximum stay in the natural-log domain, and only the difference to the new maximum is scaled by log2(e) for `ex2`. The FP16 probability tile holds `p x 2^15` and the row sum stays unscaled FP32. A CTA that cannot bound the FP16 rounding evaluates its query heads in FP32 on the CUDA cores instead (`AttendRowFp32`; the decision is uniform per CTA and the FP32 path starts from the inputs, so nothing is clipped or rounded): an element of `Q * key_scale` is not finite or above 65504; for one head `eps = scale * 128 * sum_d abs(Q'_d - fp16(Q'_d))`, which bounds the shift of a logit by the rounding of `Q'`, exceeds 0.125 nats (`kInt8TcQueryErrorBudget`); a `value_scale` is not finite; or some channel has `abs(value_scale) > 2^16` and a valid candidate is more than 20 nats (`kInt8TcTailDepth`) below the running maximum (a probability below 2^-29 is subnormal in the `p x 2^15` tile). The kernel is selected only for `scale <= 1` (`kInt8TcMaxScale`; the op validates the scale as finite and non-negative, and the default `1 / sqrt(head_size)` is 0.0625 here), because the rounding bound grows with the scale; a larger scale takes the BF16-dequantizing kernel. | | `onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu` | Dispatch to the new kernel when its switch is on, the softmax `scale` is at most 1, the buffers are aligned and the loaded PTX target is at least SM80 (`cudaFuncGetAttributes`, otherwise an SM75 PTX image would run an empty kernel stub); `DynamicSparseAttentionPrefillKernel<true>` stays as the fallback. | | `onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.{h,cc}`, `dynamic_sparse_attention_impl.h` | The four `ep.cuda.dsa_*` options are resolved once at kernel construction (session option, else environment variable, else default) and passed to the dispatch as parameters. The two INT8 tensor-core switches default to off. | | `onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc` | The expected output of the INT8 cases is an FP64 reference computed from the inputs as the kernel receives them (BF16 query, INT8 codes, scales) with max-subtracted weights, rounded to BF16 once. `Int8CachePrefillTensorCoreKernel_CUDA` sweeps query heads {2, 12, 24, 32} x selected candidates {1, 64, 65, 509} x invalid indices with the tensor-core kernel on and off, on random INT8 rows with per-channel scales. A two-candidate fixture (two selected candidates of an otherwise empty cache, so that the output follows from a few numbers) checks the output to one BF16 ulp, each case once with both candidates in one 64-candidate tile and once with them in the first and the last of three tiles, which the decode kernel spreads over up to three splits (two of them active): `Int8TensorCoreQueryFoldingRoundingError_CUDA` (an in-range query whose FP16 rounding cancels the logit gap: 0.5 instead of 1), `Int8TensorCoreLargeCommonLogit_CUDA` (a large common logit, 8290560 at scale 1) and `Int8TensorCoreTinyValueScale_CUDA` (`value_scale = 2^-135`, a BF16 subnormal result). `Int8TensorCoreRandomQueryFolding_CUDA` uses random rows whose folded query products are not FP16 numbers, on both sides of the error budget, and `Int8TensorCoreLargeValueScaleRandom_CUDA` random rows with value scales of 2^14 to 2^17 and logit spreads of 0.7, 11 and 22 nats for the tail-depth check. The FP32 paths are also covered by `Int8PrefillCancellingLargeQuery_CUDA` and `Int8PrefillQueryBeyondHalfRange_CUDA` (four head counts and three candidate counts with invalid entries and the query scaled by 2^28, which sends every CTA to the FP32 path, plus a case where only some CTAs take it). `PrefillSessionOptions_CUDA` checks that options override the environment, that invalid values fail (the failure cases are not compiled with `ORT_NO_EXCEPTIONS`) and that only the strings `0` and `1` are accepted. Misaligned-buffer cases (8-byte K/V, 4-byte output) cover the fallback to the BF16-dequantizing kernel. `Int8TensorCoreLargeSoftmaxScale_CUDA` (`scale = 2^30` with `key_scale = 2^-30`: the `scale <= 1` gate) and `Int8TensorCoreTinyProbabilityLargeValueScale_CUDA` (`value_scale = 2^28` with a candidate of probability about 2^-29: the probability scaling) run prefill and then decode on a two-candidate case with an analytic result. | ##### 2. DynamicSparseAttention: split-K tensor-core decode for the INT8 KV cache | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh` | The kernel body of section 1 with `kSplitK = true`: the grid gets a third dimension, every CTA consumes a contiguous range of the row's 64-candidate tiles and writes an un-normalized partial result (running maximum in the natural-log domain, running sum, value-scaled accumulator) in the layout of the grouped decode kernel. It replaces the CUDA-core `GroupedDecodeDynamicSparseAttentionKernel`, which was issue bound on its per-candidate dot product, warp reduction and exponentials. Decode passes the probability tile as an FP16 high part plus an FP16 residual (about 22 significant bits), because the signed sum of large value codes can cancel to a value far below the rounding of one FP16 probability. A CTA that meets one of the conditions of section 1 evaluates its candidates in FP32 (`AttendRowFp32<true>`) and publishes an FP32 partial in the same layout, so the fallback holds for the splits too. A partial is written as `(acc * 2^-15) * value_scale`, so the merge kernel is unchanged. | | `onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu` | `GetInt8TcDecodeSplitCount`: the smallest power-of-two number of 64-candidate tiles per split for which all `ceil(max_selected / 64) * rows * kv_heads` tiles fit one wave of resident CTAs (SM count x `cudaOccupancyMaxActiveBlocksPerMultiprocessor` of the kernel's shared-memory and register footprint: 396 on an H200). The split count is `ceil(tiles / tiles_per_split)`, at most 32 and never more than the workspace was sized for. A positive `ep.cuda.dsa_decode_int8_tc_tiles_per_split` replaces the automatic value; it is a target, because the cap, the workspace size and the even distribution of a row's tiles still apply. The partial vectors come first in the workspace so that they stay 16-byte aligned. `MergeGroupedDecodeDynamicSparseAttentionKernel` combines the partials unchanged. Routing on top of `UseGroupedDecodeAttention` is listed under "Gates and switches"; everything else keeps the CUDA-core kernel. | | `onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc` | `Int8CacheDecodeTensorCoreKernel_CUDA`: 1, 3 and 8 query rows x 8, 12, 24 and 32 query heads (group sizes 4, 6, 12 and 16, padded to the 16-row MMA tile) x selected candidates {1, 65, 509} around the 64-candidate tile x invalid indices, with the automatic split count, 1 and 3 tiles per split (a split without tiles publishes an empty partial) and the kernel switched off. `Int8TensorCoreDecodeSignedValueCancellation_CUDA` (probabilities 0.5 +- 2^-14 with values +-4064: a single FP16 tile would cancel them to 0) and `Int8TensorCoreDecodeTilesPerSplitCaps_CUDA` (lists of 2048, 2112 and 8192 candidates, that is 32, 33 and 128 tiles, with the automatic split count and 1, 2, 3 and 64 tiles per split, 1 and 8 rows: whatever number of splits results, every candidate reaches the merge; 8192 candidates with 1 tile per split get 32 splits of 4 tiles). `Int8DecodeQueryBeyondHalfRange_CUDA` multiplies every query by 2^28 so that every CTA takes the FP32 path for any split count, plus a case where only some CTAs do. `DecodeSessionOptions_CUDA` checks the switch, the tuning knob, the precedence over the environment and the rejection of invalid values. `Int8CacheDecodeMisalignedBuffers_CUDA`: 16-byte-aligned buffers use the tensor-core kernel, misaligned ones fall back to the CUDA-core kernel. The two scale tests of section 1 run the decode kernel as well. | ##### 3. Early PDL trigger in `ProjectedSigmoidGateMulKernel` | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu` | `ProjectedSigmoidGateMulKernel` issues `griddepcontrol.launch_dependents` first (`__CUDA_ARCH__ >= 900`). | The kernel feeds the hyper-connection down kernel, which is launched with PDL and reads the gate-mul output only after `griddepcontrol.wait`; its constant-weight prefetch now overlaps with the gate-mul kernel. There is no effect for plain launches, no change of any result and no switch. The same trigger in the fpA_intB GEMV kernels is part of #33689; together they shorten an `S = 8` decode cycle by 1.4-1.7% (preliminary, see "Performance"). ##### 4. HyperConnectionLowRankMix: branch-free sigmoid | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu` | `SigmoidF32` evaluates `e = expf(-fabsf(v))` once and selects the numerator, `(v >= 0 ? 1 : e) / (1 + e)`, instead of branching between `1 / (1 + expf(-v))` and `expf(v) / (1 + expf(v))`. Both branches already pass `-fabsf(v)` to `expf`, so the result is bit-identical. In the prefill pre-mix neighbouring elements have different signs, so nearly every warp executed both sides of the old branch. | ##### 5. Session options (shared by 1 and 2) | File | Change | |------|--------| | `include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h` | Four `ep.cuda.*` keys: `dsa_prefill_int8_tc`, `dsa_prefill_kernel`, `dsa_decode_int8_tc` and `dsa_decode_int8_tc_tiles_per_split`. The descriptions state the opt-in defaults and the exact `"0"` / `"1"` grammar, and that tiles per split is a target subject to the split cap, the workspace size and the balancing of the row's tiles. | | `onnxruntime/contrib_ops/cuda/cuda_session_options.h` (new) | `ReadCudaSessionOption` (integer threshold) and `ReadCudaSessionSwitch` (exactly `"0"` or `"1"`: `"00"`, `"01"`, `"-0"` and `"+1"` are rejected): an explicit session option wins, otherwise the environment variable (an empty variable counts as unset), otherwise the default. The whole string is parsed and range-checked, and an invalid value fails kernel construction. #33689 adds the same file for its own switches (see "Related PRs and review guide"). | ##### 6. Documentation and benchmark | File | Change | |------|--------| | `docs/contrib_ops/cuda/dynamic_sparse_attention.md` | The `ep.cuda.dsa_*` options next to their environment variables with the opt-in defaults, the dispatch-table rows and alignment fallbacks of the INT8-cache prefill and split-K decode kernels, the tiles-per-split rule with its sweep and regret statistics, "Numerics and fallbacks" (the conditions of the FP32 path, what the model does, the measured cost of the path), "Rollout" and the operator-level benchmark. | | `docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md` | The early `griddepcontrol.launch_dependents` trigger of the kernel in front of `HcDownKernel` and the branch-free `SigmoidF32`. | | `onnxruntime/test/python/transformers/benchmark_dynamic_sparse_attention.py` (new) | Operator-level benchmark: a single-node `DynamicSparseAttention` model per case (the shapes of the 12 layers; `--rc-model` reads the scales of a model's node), one session per kernel configuration, the GPU time of the kernels of one `Run` from the CUPTI-based torch profiler, alternating rounds, and the RMS error against an FP64 reference. `--query-std`, `--key-scale-multiplier` and `--value-scale-multiplier` reach the FP32 paths, and `--tiles-per-split` sweeps the decode split rule. | ##### Gates and switches Every switch is available as an environment variable and as a CUDA session option (`ep.cuda.*`, declared in `onnxruntime_session_options_config_keys.h`). It is resolved once when the kernel is created, so a change needs a new session, and an explicit session option takes precedence over the environment variable. The 0/1 switches accept exactly the strings `0` and `1`, the tuning knob a non-negative decimal integer; any other value (including an empty session-option string) fails session initialization. `ORT_DSA_PREFILL_KERNEL` already exists on `main` (read on every call; any non-zero integer enabled it) and now follows the same rules. | Kernel | Used when | Environment variable / session option | |--------|-----------|--------| | INT8-KV DSA prefill | BF16 query / output, INT8 main cache, the prefill kernel path (`sequence_length >= 8`, not the grouped decode path), `head_size = 256`, at most 16 query heads per KV head, selected-only from the main cache (no `local_plus_selected`), no smooth softmax, `scale <= 1`, `max_selected > 0`, SM80+, 16-byte-aligned `present_key`, `present_value` and scales, 8-byte-aligned `output`, enough opt-in shared memory (72,960 B) | `ORT_DSA_PREFILL_INT8_TC=1` / `ep.cuda.dsa_prefill_int8_tc` turns it on (default `0`: the BF16-dequantizing kernel, which has no FP16 range limit); `ORT_DSA_PREFILL_KERNEL=0` / `ep.cuda.dsa_prefill_kernel` turns off the whole tensor-core prefill family (default `1`) | | INT8-KV DSA decode | BF16, INT8 main cache, grouped decode path (`1 <= sequence_length <= 8`, 4 to 16 query heads per KV head, `max_selected > 256`), `head_size = 256`, selected-only from the main cache, no smooth softmax, `scale <= 1`, 16-byte-aligned query, `present_key`, `present_value`, scales and attention workspace, SM80+, enough opt-in shared memory (75,264 B) | `ORT_DSA_DECODE_INT8_TC=1` / `ep.cuda.dsa_decode_int8_tc` turns it on (default `0`: the CUDA-core grouped decode kernel, bit for bit as before); `ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT` / `ep.cuda.dsa_decode_int8_tc_tiles_per_split` (default 0 = automatic; a requested tiles-per-split target, see section 2) | | Early PDL trigger | SM 9.0+, PDL successor | none (no effect for plain launches) | | Branch-free sigmoid | always | none (bit-identical) | The tensor-core kernels also require a loaded PTX target of at least SM80. If only older PTX is available (for example a `compute_75` image that the driver JIT-compiles for a newer GPU), their `__CUDA_ARCH__ >= 800` bodies are empty stubs that would silently leave the output unwritten, so the dispatch checks `cudaFuncGetAttributes(...).ptxVersion >= 80` and otherwise uses the generic kernels. ##### Numerics - The branch-free sigmoid is bit-identical to the branchy one for all 2^32 `float` bit patterns (exhaustive check in a standalone harness, same nvcc flags, no fast-math). - The early PDL trigger changes no result (generated tokens identical in 45 / 45 decode runs). - The INT8 tensor-core kernels are close to, but not bit-identical with, the previous kernels: they round `Q'` and the probability tile to FP16 and sum in a different FP32 order. The switches' default `0` keeps the previous kernels bit for bit. Against an FP64 reference (the BF16 rounding of the output alone gives about 1.5e-3 relative RMS), decode is at 1.6-1.7e-3 with both kernels (the tensor-core kernel is up to 3% higher), and prefill at 1.7e-3 against 2.7-4.3e-3 for the BF16-dequantizing kernel (tables under "Performance"). - The cases found in review, and how the kernels handle them: | Case | Handling | Test | |---|---|---| | A finite `Q * key_scale` whose FP16 rounding changes the logits (products 32768 and -32766 round to 32768 and -32768, the gap of 16 becomes 0: 0.5 instead of 1) | The per-head bound `eps` above 0.125 nats sends the CTA to the FP32 path | `Int8TensorCoreQueryFoldingRoundingError_CUDA`, `Int8TensorCoreRandomQueryFolding_CUDA` | | Decode: two probabilities 0.5 +- 2^-14 round to the same FP16 value and the signed values cancel (the output 127/256 = 0.496 became 0) | The decode tile is an FP16 high part plus an FP16 residual (prefill keeps one tile, as the BF16 kernel does with a coarser one) | `Int8TensorCoreDecodeSignedValueCancellation_CUDA` | | Scores of 8290560 and 8290560.5 (a large common logit) converted to the log2 domain before the maximum is subtracted lose their gap to float rounding (0.5 instead of 0.62 at scale 1 and 0.508 at scale 1/16) | Natural-log scores and maximum; only the difference is scaled by log2(e) | `Int8TensorCoreLargeCommonLogit_CUDA` | | `value_scale = 2^-135`: `value_scale * 2^-15` underflows before the accumulator is applied (BF16 subnormal result became 0) | Divide the accumulator by 2^15 first, then multiply by `value_scale` | `Int8TensorCoreTinyValueScale_CUDA` | | `abs(value_scale) > 2^16` with a probability below 2^-29 (subnormal in the FP16 tile) | The tail-depth check decides per CTA, instead of sending every CTA to the FP32 path as soon as one channel exceeds 2^16 | `Int8TensorCoreLargeValueScaleRandom_CUDA` | - Cost of the FP32 path (H200, `benchmark_dynamic_sparse_attention.py`; 24 query / 2 KV heads, 2048 selected candidates, context 8192; attention-kernel time, previous kernel against tensor-core kernel): | Data | Rows | Previous | Tensor core | Tensor core vs previous | |---|---|---|---|---| | Queries like the model's | 128 / 8192 | 0.173 / 7.82 ms | 0.063 / 2.83 ms | 2.7x / 2.8x faster | | Queries 8 times larger (`--query-std 11.2`): every CTA is over the error budget | 128 / 8192 / 8 (decode) | 0.172 ms / 7.82 ms / 73.6 us | 3.68 ms / 135.6 ms / 240 us | 21x / 17x / 3.3x slower | | `abs(value_scale)` 24k to 360k (`--value-scale-multiplier 1.2e7`), logit standard deviation about 2.4 nats: 10 to 12% of the CTAs have a candidate deeper than 20 nats | 128 / 512 / 2048 | 0.172 / 0.577 / 2.288 ms | 3.61 / 3.90 / 6.17 ms | 21x / 6.8x / 2.7x slower | | The same value scales, logit standard deviation about 1.3 nats: no CTA is deeper than 20 nats | 8192 / 8 (decode) | 7.78 ms / 73.6 us | 2.86 ms / 13.3 us | 2.7x / 5.6x faster | A CTA on the FP32 path needs about 3.6 ms for 2048 candidates, so the time of a launch with one wave of CTAs is set by its slowest CTA. The fallback keeps the accuracy (relative RMS 1.7e-3 against FP64 with queries 8 times larger; the previous prefill kernel: 8.5e-3). The first group of rows is the Qwen3.8-Flash-Next case: the largest `value_scale` of its 12 layers is 0.285 (the tail-depth check is never armed) and a Monte Carlo over their `q_norm` weights and key scales gives a per-head `eps` of 0.015 on average and 0.022 at most, against the budget of 0.125, which corresponds to queries about 8 times larger than the model produces. Workloads with such queries, or with `abs(value_scale) > 2^16` and wide logit spreads, should leave the switches at `0`. The FP32 path of a CTA is documented in the new "Numerics and fallbacks" section of `docs/contrib_ops/cuda/dynamic_sparse_attention.md`. - Cost of the numerics fixes on the fast path (H200, the same script; separate processes against the previous head of this PR, `ceed60271c`, not paired; attention-kernel time): prefill at 128 / 2048 / 8192 rows 0.062 / 0.815 / 2.760 ms against 0.063 / 0.832 / 2.830 ms (+1.6% to +2.5%); decode with 8 rows 12.23 us against 13.2-13.3 us (+8-9%) and with 1 row 8.84 us against 9.46 us (+7%). The previous head sent every CTA to the FP32 path when one channel had `abs(value_scale) > 2^16` (8192 rows: 2.76 ms became 136 ms). #### Testing In-tree CUDA 13.0 Release build, `sm_90a`, one H200 (SM 9.0), on this branch (`onnxruntime_provider_test`; this checkout needs `-Donnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=OFF` because `gemm_tactic_cache_test.cc` does not link otherwise (`undefined reference to onnxruntime::g_host`), which is unrelated to this PR): | Command (`--gtest_filter=`) | Result | |---|---| | `*DynamicSparseAttention*:*HyperConnection*` | 103 run, **103 passed**, 0 failed (73 `DynamicSparseAttentionTest`, 13 `HyperConnectionLowRankMixTest`, 17 `HyperConnectionOpsTest`); about 100 s | | `*GatedAdd*:*ProjectedSigmoid*:*SigmoidGate*:*GateMul*` | 14 run, **14 passed**, 0 failed (12 `ContribOpGatedAddTest`, including `ProjectedGateBFloat16` and `ProjectedGateK1Regression`, and 2 `EngramOpsTest`) | - The review round adds 7 DSA tests (`Int8TensorCoreQueryFoldingRoundingError_CUDA`, `Int8TensorCoreDecodeSignedValueCancellation_CUDA`, `Int8TensorCoreLargeCommonLogit_CUDA`, `Int8TensorCoreTinyValueScale_CUDA`, `Int8TensorCoreDecodeTilesPerSplitCaps_CUDA`, `Int8TensorCoreLargeValueScaleRandom_CUDA`, `Int8TensorCoreRandomQueryFolding_CUDA`), moves the INT8 reference to FP64 with max-subtracted weights, and tightens the two session-option tests (`"00"`, `"01"`, `"-0"`, `"000001"`, `"+1"`, `" 1"`, `"1 "`, `""`, `"2"`, `"-1"`, `"1x"` and `"true"` are rejected for the switches, and `"01"`, `"-0"` and `"+1"` for the environment variables). - Red-green against the previous head of this PR (`ceed60271c`): the final test binary was run against the provider built from that head (`ORT_UNIT_TEST_MAIN_DYNAMIC_PLUGIN_EP_CONFIG_JSON` selects the library) for the 13 tests that concern the review. 7 fail on it: the five numerics tests `Int8TensorCoreQueryFoldingRoundingError_CUDA`, `Int8TensorCoreDecodeSignedValueCancellation_CUDA`, `Int8TensorCoreLargeCommonLogit_CUDA`, `Int8TensorCoreTinyValueScale_CUDA` and `Int8TensorCoreRandomQueryFolding_CUDA`, and `PrefillSessionOptions_CUDA` and `DecodeSessionOptions_CUDA` (the stricter parser). The other 6 pass on it: `Int8TensorCoreDecodeTilesPerSplitCaps_CUDA` and `Int8TensorCoreLargeValueScaleRandom_CUDA` are coverage tests (the split caps; the tail-depth check against FP64 on random data), and `Int8TensorCoreTinyProbabilityLargeValueScale_CUDA`, `Int8TensorCoreLargeSoftmaxScale_CUDA`, `Int8PrefillQueryBeyondHalfRange_CUDA` and `Int8DecodeQueryBeyondHalfRange_CUDA` are the guards of the earlier rounds. All 73 pass with the final provider. - Under `nsys` the session-option tests launch the intended kernels on the final tree. `PrefillSessionOptions_CUDA`: the split / merge kernels for `dsa_prefill_kernel=0`, `DynamicSparseAttentionPrefillKernel<true>` for `dsa_prefill_int8_tc=0` and `DynamicSparseAttentionInt8TcKernel<false>` for `=1` (two launches of each). `DecodeSessionOptions_CUDA`: `GroupedDecodeDynamicSparseAttentionKernel` for `dsa_decode_int8_tc=0` and `DynamicSparseAttentionInt8TcKernel<true>` for `=1` (three launches of each), each followed by `MergeGroupedDecodeDynamicSparseAttentionKernel`. - `lintrunner` (clang-format, ruff, ruff-format) is clean on all 13 changed files. `dynamic_sparse_attention_impl.cu`, `gated_add_impl.cu` and `hyper_connection_low_rank_mix_impl.cu` also compile for `sm_75`, `sm_80`, `sm_86` and `sm_120` with the flags of the build. - The PTX guards (`cudaFuncGetAttributes(...).ptxVersion >= 80`) have no unit test, because they need an executable that carries only `compute_75` PTX. On the H200 (the driver JIT-compiles the PTX for SM 9.0) such an executable reports `ptxVersion = 75` for `DynamicSparseAttentionPrefillKernel<true>` and `<false>`, `DynamicSparseAttentionInt8TcKernel<false>` and `DynamicSparseAttentionInt8TcKernel<true>`, so dispatch rejects them; built for `sm_90a` the same kernels report `ptxVersion = 90`. - CI: the new DSA tests need a BF16-capable device (SM80+) and skip otherwise. They need 72,960 B (prefill) and 75,264 B (decode) of opt-in shared memory for the tensor-core kernels, which `sm_86` has, so they run the tensor-core paths there. #### Checklist - [x] Tests added/updated and passing - [x] No breaking API changes. By default the INT8-KV attention behaves as before (both tensor-core kernels are opt-in; the CUDA-core decode kernel and the BF16-dequantizing prefill kernel run bit for bit as before). Behavior changes: `ORT_DSA_PREFILL_KERNEL` is now resolved once per kernel and strictly 0/1, and the tensor-core kernels, when enabled, are not bit-identical to the previous kernels - [x] Documentation updated (if applicable) ### Motivation and Context At 8192 tokens per prefill chunk the INT8-KV prefill kernel took 12-13% of the chunk with HMMA at only 2.8% of its instructions: it dequantized every gathered K/V tile to BF16 in shared memory before the BF16 MMAs. In decode (`S = 8`, 12 query heads per KV head) the CUDA-core INT8 kernel took 70.3 us plus 5.3 us for the merge per layer at 2K to 32K context, 6.9% of a decode cycle (ncu: issue and shared-load bound on the per-candidate serial chain). In the prefill pre-mix of the hyper-connections, the data-dependent branch of `SigmoidF32` made nearly every warp execute both of its paths. #### Performance on NVIDIA H200 (SM 9.0a, CUDA 13.0) ##### Operator level (final revision) `benchmark_dynamic_sparse_attention.py` builds a single-node `DynamicSparseAttention` model per case (24 query / 2 KV heads, head size 256, INT8 cache, scale 1/16, 2048 selected candidates per row: the shapes of the 12 layers) and one session per kernel configuration, because the `ep.cuda.dsa_*` options are read when the kernel is created. The time is the GPU time of the kernels that one `Run` launches (CUPTI-based torch profiler, so host launch overhead does not enter): "attention" sums the `DynamicSparseAttention` kernels including the decode merge kernel, "node" all GPU activity of the node (cache append, query preparation, ...). Each round measures every configuration once, the order alternates between rounds, the tables show the median of 7 rounds, and the speedup is the median of the paired ratios of the rounds. 3 warmup runs per measurement, 200 profiled runs per round for decode and 5 for prefill. The error is the RMS error of the BF16 output relative to the RMS of an FP64 reference over 64 evenly spaced rows. The data are synthetic: BF16 queries with standard deviation 1.4 (the queries of the model after `q_norm` have a per-head RMS of 1.3 to 1.6), INT8 cache codes with standard deviation 40, and key and value scales drawn from the ranges of that model. "Previous" is the kernel that the switch replaces (the CUDA-core grouped decode, the BF16-dequantizing prefill): ```bash python benchmark_dynamic_sparse_attention.py --rows 1 8 --context 2048 8192 32768 --tiles-per-split 0 1 2 4 8 16 python benchmark_dynamic_sparse_attention.py --rows 128 2048 8192 --context 8192 32768 ``` Decode (`sequence_length <= 8`): | Rows | Context | Attention, previous | Attention, tensor core | Speedup | Node, previous | Node, tensor core | RMS error, previous / tensor core | |---|---|---|---|---|---|---|---| | 1 | 2048 | 31.1 us | 9.5 us | 3.30x | 34.1 us | 12.5 us | 1.59e-3 / 1.62e-3 | | 1 | 8192 | 31.2 us | 9.5 us | 3.30x | 34.2 us | 12.5 us | 1.55e-3 / 1.57e-3 | | 1 | 32768 | 31.2 us | 9.5 us | 3.30x | 34.1 us | 12.5 us | 1.67e-3 / 1.72e-3 | | 8 | 2048 | 73.6 us | 13.2 us | 5.58x | 76.7 us | 16.5 us | 1.64e-3 / 1.69e-3 | | 8 | 8192 | 73.6 us | 13.3 us | 5.57x | 76.8 us | 16.6 us | 1.62e-3 / 1.66e-3 | | 8 | 32768 | 73.7 us | 13.2 us | 5.57x | 77.0 us | 16.6 us | 1.68e-3 / 1.73e-3 | Prefill: | Rows | Context | Attention, previous | Attention, tensor core | Speedup | Node, previous | Node, tensor core | RMS error, previous / tensor core | |---|---|---|---|---|---|---|---| | 128 | 8192 | 0.173 ms | 0.063 ms | 2.72x | 0.179 ms | 0.070 ms | 4.19e-3 / 1.71e-3 | | 128 | 32768 | 0.191 ms | 0.064 ms | 2.99x | 0.199 ms | 0.072 ms | 4.25e-3 / 1.70e-3 | | 2048 | 8192 | 2.287 ms | 0.832 ms | 2.75x | 2.350 ms | 0.894 ms | 4.25e-3 / 1.71e-3 | | 2048 | 32768 | 2.294 ms | 0.832 ms | 2.75x | 2.356 ms | 0.895 ms | 4.21e-3 / 1.70e-3 | | 8192 | 8192 | 7.818 ms | 2.830 ms | 2.76x | 8.057 ms | 3.069 ms | 2.69e-3 / 1.66e-3 | | 8192 | 32768 | 8.877 ms | 3.200 ms | 2.77x | 9.116 ms | 3.439 ms | 4.29e-3 / 1.71e-3 | The decode time does not depend on the context length (the selected candidates are gathered). The error includes the rounding of the BF16 output, whose floor is about 1.5e-3 relative to the RMS: both decode kernels and the tensor-core prefill are near it, and the BF16-dequantizing prefill, which rounds the dequantized keys and values and the probabilities to BF16, has 1.6 to 2.5 times the error of the tensor-core prefill. Decode time against the tiles per split (8 rows, 2048 candidates, kernel plus merge, `--tiles-per-split`): | Tiles per split | 1 | 2 (automatic) | 4 | 8 | 16 | |---|---|---|---|---|---| | Splits, CTAs | 32, 512 | 16, 256 | 8, 128 | 4, 64 | 2, 32 | | Time | 18.8 us | 13.2 us | 12.7 us | 17.7 us | 29.2 us | The automatic rule is a heuristic, not the optimum. Over 20 shapes on the H200 (1, 2, 4 and 8 rows, 512 to 8192 candidates, context 8192, each swept over 1, 2, 3, 4, 6, 8 and 16 tiles per split) it is on average 1.8% and at most 14.8% slower than the best swept value (4 rows and 2048 candidates: 11.9 us against 10.4 us); the previous target of about 384 CTAs is 2.7% and 14.2% slower on the same points. It matters when fewer CTAs are resident: with the kernel's shared memory padded on the H200 so that one or two CTAs fit an SM (132 or 264 resident CTAs; an ad hoc harness that is not part of this change) the rule is on average 0.4% (at most 2.4%) and 3.1% (13.4%) slower than the best swept value, where the 384-CTA target is 23.8% (44.9%) and 14.4% (38.4%) slower. An A100 (108 SMs, two resident CTAs per SM by the 75,264 B of shared memory) is expected to choose 4 tiles per split, 8 splits, for 8 rows and 2048 candidates, the faster setting in the A100 measurement reported in the review (22.1 us against 25.9 us at 16 splits). That choice is derived from the rule, not measured: the A100 has not been run with this script. `ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT` tunes other devices. The other kernels of this PR, per launch (H200, `nsys` / CUDA events; unchanged since the earlier revision): | Kernel | Before | After | |---|---|---| | `PreMixVec8Kernel` per launch (hidden 2560, 4 branches, 8192-token chunk; 96 launches per chunk) | 112.1 us | 82.7 us | | `PreMixVec8Kernel`, the same launch in the model (`nsys`, 8192-token chunk, `all_int4`) | 161.2 us | 119.2 us | | `ScaledSiluKernel` | 7.3 us | 6.1 us | ##### End to end (preliminary) These end-to-end numbers were measured for an earlier revision of the kernels, before the numerics changes of the review, and have not been reproduced for the final revision: the stack that they need (the model packages and runtime, and #33689 together with this PR on top of a `main` that contains #32126, #33650 and #33714) is not available for a final-head run. Each row is an A/B of consecutive providers on the same packages, one round unless noted (not alternated), without raw timings and without confidence intervals. They show the direction of the effect, not a qualified gain. A shorter cycle is not the same as better user-visible throughput, and accepted tokens per cycle and time per accepted output token were not measured for any of these runs. Measured on the Qwen3.8-Flash-Next packages (`hidden = 2560`, `intermediate = 640`, 512 experts, top-10; `all_int4` is the all-INT4 package, `top16` and `top08` are the two INT3-mixed packages), B = 1, DFlash2 `S = 8` decode. Decode is in ms per cycle. TTFT of the INT8-KV prefill kernel in seconds (one round; decode is unchanged): | Package | 2048 | 8192 | 32768 | |---|---|---|---| | `all_int4` | 0.1985 -> 0.1889 (-4.8%) | 0.7122 -> 0.6524 (-8.4%) | 3.0210 -> 2.7782 (-8.0%) | | `top16` | 0.2155 -> 0.2068 (-4.0%) | 0.7339 -> 0.6714 (-8.5%) | 3.0969 -> 2.8608 (-7.6%) | | `top08` | 0.2216 -> 0.2121 (-4.3%) | 0.7478 -> 0.6768 (-9.5%) | 3.1504 -> 2.8540 (-9.4%) | Decode of the split-K INT8-KV kernel, ms per cycle at 2K / 8K / 32K context: | Package | 2K | 8K | 32K | |---|---|---|---| | `all_int4` | 11.672 -> 11.017 (-5.6%) | 11.962 -> 11.143 (-6.8%) | 12.157 -> 11.391 (-6.3%) | | `top16` | 12.148 -> 11.381 (-6.3%) | 12.353 -> 11.608 (-6.0%) | 12.629 -> 11.953 (-5.4%) | | `top08` | 12.216 -> 11.505 (-5.8%) | 12.474 -> 11.675 (-6.4%) | 12.661 -> 12.082 (-4.6%) | Branch-free sigmoid, TTFT in ms (median of 5 repetitions; 9 / 9 cells faster, first tokens identical). Decode is within +-0.3%, because the fused decode kernels call the function for at most 8 rows. The per-cell gains of 0.3-1.1% are about the size of the repeat noise between repeats of one stack (about 1%); the kernel-level `PreMixVec8Kernel` time above is the stronger evidence: | Package | 2048 | 8192 | 32768 | |---|---|---|---| | `all_int4` | 174.1 -> 173.3 (-0.45%) | 574.7 -> 572.9 (-0.31%) | 2450.1 -> 2439.2 (-0.45%) | | `top16` | 191.3 -> 190.5 (-0.42%) | 596.3 -> 592.1 (-0.70%) | 2560.2 -> 2531.7 (-1.11%) | | `top08` | 196.2 -> 194.7 (-0.79%) | 601.1 -> 597.2 (-0.65%) | 2573.5 -> 2554.2 (-0.75%) | The early PDL trigger of the gate-mul kernel and the one of the fpA_intB GEMV kernels (#33689) together shorten a cycle by 1.4-1.7% (a joint result of two PRs). For the whole stack that contains both PRs (plus #32126, #33650 and the package configuration), the median over 4 rounds against the previous release candidate is -32.8 / -36.0 / -36.9% TTFT at 8192 tokens, -32.3 / -35.2 / -36.1% at 32768 tokens, and -27.6 / -24.6 / -24.2% decode ms per cycle at 2K (`all_int4` / `top16` / `top08`); these combined-stack numbers cannot be attributed to a single PR. #### Accuracy The sigmoid change is bit-exact and the early PDL trigger changes no result. At the operator level (final revision) the error against an FP64 reference is in the tables above: decode 1.6-1.7e-3 with both kernels (the tensor-core kernel is up to 3% higher) and prefill 1.7e-3 against 2.7-4.3e-3 for the BF16-dequantizing kernel, with the BF16 rounding of the output alone at about 1.5e-3. The unit tests compare against an FP64 reference with max-subtracted weights and fix the numerical corners found in review to one BF16 ulp. End to end, the final revision has not been qualified, and that is why both tensor-core kernels are opt-in. The one end-to-end accuracy gate is for an earlier revision. The split-K decode kernel changes the generated tokens (a different FP32 accumulation order; for 263 of 800 `all_int4` MMLU-Pro prompts), so it was gated with MMLU-Pro: 800 questions x 3 packages = 2400 items, greedy decoding through the full engine, `max_new = 8192`. Together with the dequantize + cuBLAS route of #33689, it gave 668 / 678 / 667 correct answers against 665 / 676 / 670 for the previous stack, pooled **+0.08 pp** (95% CI -0.60 .. +0.77 pp, exact McNemar p = 0.91); no package differs at the 5% level. The first-round stack, which contains the INT8-KV prefill kernel, was -0.67 pp pooled against the previous stack (95% CI -1.36 .. +0.03 pp, p = 0.076). These runs measure the accumulated effect of all kernels of the stack, not of one PR, and a drift of up to about 1.4 pp cannot be excluded with 2400 items; a non-significant difference does not establish non-inferiority. They predate the numerics of the final revision (the natural-log domain, the `Q'` rounding bound, the decode probability residual, the value-scale ordering and the tail-depth check), and so does a later check of commit 8 (`engine_behaviours.py` 27 / 27 on the three packages; TTFT and decode within the noise against the stack before commit 8). They do not cover the final revision. Making a switch the default needs the final head qualified with the rest of the stack held fixed and prefill and decode toggled separately: per-package paired accuracy with confidence intervals against a stated degradation bound, including long contexts, exact build and model identities, and for speculative decode the accepted tokens per cycle and the time per accepted token. Packages that qualify the kernels enable them in their session options (`ep.cuda.dsa_prefill_int8_tc` and `ep.cuda.dsa_decode_int8_tc` set to `1`). #### Related PRs and review guide - #33689 is the sibling PR (tensor-core `MatMulNBits` / fpA_intB and QMoE / MoE GEMM kernels). The changes are divided by operator and have no kernel in common. Both PRs add `cuda_session_options.h`, and the two versions now differ: this PR's is a superset. It adds `ResolveCudaSessionOptionText`, `ResolveCudaSessionSwitch` (exactly `"0"` or `"1"`) and `ReadCudaSessionSwitch`, and `ReadCudaSessionOption` (the integer parser, with unchanged behavior) now calls the text helper. A trial merge of the two heads (`git merge-tree`) has this add/add conflict in that one file and no other, and the two edits of `onnxruntime_session_options_config_keys.h` are in different places. Whichever PR lands second takes this PR's version of the file; the switches of #33689 can then adopt `ReadCudaSessionSwitch`, which is a change in #33689. - #32126 (SwiGLU operator and fusion), #33650 (RouterGateUpProj) and #33714 (QMoE INT3 prefill tile) have merged into `main` since this PR was opened. This branch does not touch their files, and a trial merge of this head with the current `main` (`git merge-tree`) is clean. The decode graph measured above contains all three. - The PR has 13 files and 14 commits: commit 1 (early PDL trigger of the gate-mul kernel), commit 2 (INT8-KV prefill kernel), commit 3 (review follow-up: the FP32 fallback for out-of-range queries, the `ep.cuda.dsa_*` session options with their tests), commit 4 (fixes MSVC narrowing warnings in the DSA tests), commit 5 (split-K decode kernel and its options and tests), commit 6 (branch-free sigmoid), commit 7 (docs), commit 8 (review follow-up: the `scale <= 1` condition, the 2^15 probability scaling, `ORT_NO_EXCEPTIONS` guards in the option tests and two regression tests), commit 9 (`ceed60271c`, value-scale gate with an FP32 fallback for all CTAs, replaced by the tail-depth check of commit 10), commit 10 (`1bf2eb17cb`, review fixes of the kernels, the split count, the options and the opt-in defaults), commit 11 (`bd24bb8503`, tests), commit 12 (`2dedd0c4ff`, docs and the benchmark script), commit 13 (`e8b3ad2232`, the end-to-end timings in the hyper-connection doc are labeled preliminary) and commit 14 (`d60b0f4b63`, the DSA doc states that the end-to-end timings have no acceptance measurements). `dynamic_sparse_attention_prefill_int8.cuh` is the main file to review (812 lines). - The review threads of #33689 on the prefill kernel and its test (clamping of the scaled query, identical K/V rows in the sweep) were answered there with the fixes of commit 3; its other threads concern the QMoE and dense INT8 kernels that stay in #33689. The review threads of this PR on the numerics (the rounding of the folded query, the decode probability rounding, the log2 domain, the tiny value scale and the value-scale cliff), the tests (FP64 reference, active splits, exact checks), the split target, the switch grammar, the tiles-per-split wording and the rollout and end-to-end claims are answered by commits 10 to 14. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
RunPackedGroupedGemminmoe_int2_prefill.curuns the CUTLASS INT4 grouped GEMM of the QMoE INT2/INT3 prefill path: raw INT2/INT3 weights are repacked to the SM80 INT4 layout first, then FC1 and FC2 go through this GEMM. It hard-coded the 32-row tileCtaShape32x128x64_WarpShape32x32x64. That tile is too small once the experts average about 24 or more rows; Qwen3.8-Flash-Next (hidden 2560, intermediate 640, 512 experts, top-10, block size 64) has 40 rows per expert at a 2048-token prefill chunk.This PR picks the tile from
rows_per_expert = expanded_rows / num_experts:CtaShape64x128x64_WarpShape64x32x64for 24 or more,CtaShape32x128x64_WarpShape32x32x64below. It is the crossover thatPickPrefillTacticinmoe_quantization.ccalready uses for the INT4 layers. Chunks with fewer than 24 rows per expert (1024 tokens and less in the model above) keep the 32-row tile, so nothing changes for them.Only the CTA and warp tiling changes. The K loop (64-wide tiles, the group-wise dequantization, the
m16n8k16MMA) and the epilogue are the same, so every output element accumulates in the same order.Summary of Changes
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_int2_prefill.cuRunPackedGroupedGemm(theweight_bits == 4CUTLASS branch) selects the CTA/warp tile from the average rows per expert instead of always using 32x128x64.onnxruntime/test/contrib_ops/moe_test.ccQMoETest_Int3AndMixedWidthCudaPrefillParitygets three 64-row cases,{3, 3, 64, 64, 0},{3, 4, 32, 64, 0}and{4, 3, 128, 64, 0}. 64 rows with top_k 2 over 4 experts are 32 rows per expert, so the 64-row tile runs. The existing cases (at most 19 rows over 4 experts) stay on the 32-row tile.Motivation and Context
Measured on H200 (SM90, batch 1) with the two mixed INT4/INT3 Qwen3.8-Flash-Next packages, against a build that selects the 64-row tile at 24 or more rows per expert (an earlier version of this path did):
nsysof a 2048-token prefill chunk, layer 14: the INT4 grouped GEMMs take 899 us (FC1) and 533 us (FC2) on the 64x128x64 tile and 966 us and 584 us on 32x128x64, i.e. 7.5% and 9.6% longer. Over the last prefill step the MoE GEMM time is 67.88 ms against 73.04 ms (+5.16 ms), which is about 89% of the 5.8 ms TTFT difference at 2048 tokens.Testing
Validated on H200 with an in-tree CUDA EP build (
onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON, SM90), on this commit together with the tips of #33689 and #33707. The change does not depend on them: trial merges of this branch with #33689, #33707, #33650 and #32126 are conflict-free in any order (#33689 also editsmoe_test.cc, in other hunks).MoETest.QMoETest_Int3AndMixedWidthCudaPrefillParity(FP16 and BF16, 13 cases each) passes.nsysof this test shows the new rows reach the new tile: per data type, 22 launches ofMoeFCGemm<DqMmaMultistage<GemmShape<32, 128, 64>>>and 6 ofGemmShape<64, 128, 64>(3 cases, FC1 and FC2). Before this change the test cannot launch the 64-row kernel.MoETest.QMoETest_Int3*:MoETest.QMoETest_Int2Cuda*:MoETest.QMoETest_MixedWidthCuda*:*QMoECudaInt3*: 45 run, 45 passed.MoETest.*:*QMoE*: 131 run, 118 passed, 13 skipped (the CPU QMoE cases and the two Mixtral cases, which these CUDA builds skip, and the routing-statistics test for the plugin EP), 0 failed.lintrunneron the two files: no issues.