onnxruntime

e9157780 - Add FP8 GEMV KSplit32 scheduling for the 48-SM SM121 GPU (#32409)

Add FP8 GEMV KSplit32 scheduling for the 48-SM SM121 GPU (#32409) Add a KSplit32 tensor-core GEMV specialization for `MatMulBlockQuantizedFp8Weight` and select it only for qualified low-M decode shapes on the 48-SM SM121 GB10 GPU. The final selector preserves the generic schedule unless all of these conditions hold: - compute capability is exactly 12.1; - the device has exactly 48 SMs; - the original request has M at most 8; - the shape is either: - at least 1,024 output blocks and 80 64-element K windows; or - at least 320 output blocks and 128 K windows. Each FP8 GEMV output block owns 16 columns, so the thresholds use `ceil(N / 16)` rather than raw N. This gives shapes with identical launch geometry the same SM121 override: the wide boundary starts at N=16369, and the long-reduction boundary starts at N=5105. The generic selector is unchanged, preserving behavior on other devices. SM120, other SM12x minors, other SM counts, M above 8, attention-sized projections, small outputs, short reductions, and shapes immediately below both output-block boundaries retain the generic selector. KSplit32 dispatches directly to the only qualified launch, `<KSplit=32, MTiles=1>`. This avoids instantiating unreachable `<32,2>` and `<32,4>` kernels; `<32,4>` would require 64 KiB of static shared memory. Requests above M=8 are recursively tiled with GB10 KSplit32 tuning disabled, so the original request size cannot be lost during dispatch. ## Kernel performance Measured on RTX Spark (SM121, 48 SMs) as end-to-end ORT `Run()` latency with an explicit shared CUDA stream. Each final-selector cell used three alternating process rounds, with 100 warmups and 1,000 measured calls per arm per round. | Shape | M | Generic | KSplit32 | Latency reduction | |---|---:|---:|---:|---:| | 17408x5120 gate/up | 1 | 0.3674 ms | 0.3611 ms | 1.7% | | 17408x5120 gate/up | 8 | 0.4901 ms | 0.4792 ms | 2.3% | | 5120x17408 down | 1 | 0.3771 ms | 0.3684 ms | 2.3% | | 5120x17408 down | 8 | 0.3771 ms | 0.3696 ms | 2.0% | | 32768x5120 wide output | 1 | 0.6704 ms | 0.6552 ms | 2.2% | | 32768x5120 wide output | 8 | 0.7986 ms | 0.7729 ms | 3.2% | The review specifically raised the possibility that KSplit32 could regress already-saturated, vocabulary-sized output grids. Additional three-round GB10 negative-control measurements instead showed that KSplit32 remained beneficial as N increased through the Qwen 3.8 lm-head dimension: | Shape (M=1, K=5120) | Output blocks | Output waves on 48 SMs | KSplit8 | KSplit32 | Latency reduction | |---|---:|---:|---:|---:|---:| | N=32769 | 2,049 | 42.7 | 0.6836 ms | 0.6559 ms | 3.95% | | N=65536 | 4,096 | 85.3 | 1.3174 ms | 1.2693 ms | 3.63% | | N=131072 | 8,192 | 170.7 | 2.6234 ms | 2.5102 ms | 4.32% | | N=248320 lm head | 15,520 | 323.3 | 5.2056 ms | 4.8773 ms | 6.54% | Every cell used three alternating process rounds, 100 warmups, and 1,000 CUDA-event-timed ORT calls per arm. Full outputs matched with `atol=0.5` and `rtol=1e-3`. A final-selector smoke run confirmed that the natural selector matched forced KSplit32 at N=248320 (4.8339 ms natural versus 4.8683 ms forced). Because the proposed regression did not reproduce and the benefit persisted through more than 323 output waves, no arbitrary upper N cap is applied. Exact boundary cases were also positive across M=1, 2, 4, and 8: | Boundary | Latency reduction | |---|---:| | 1,024 output blocks (measured at `N=16384`), 80 K windows | 1.2% to 2.5% | | 320 output blocks (measured at `N=5120`), 128 K windows | 3.2% to 5.9% | The broad initial heuristic was rejected after finding regressions at M=16/32, on attention and small-output shapes, at 32 K windows, and around N=4096. The final selector excludes all of those regimes. Negative controls at 1,023 output blocks (`N=16368`), 79 windows, 319 output blocks (`N=5104`), and 127 windows retain the generic schedule. No-override runs followed the expected schedule for every selected and excluded case. Randomized activations and varying block scales matched the generic output in every compared full tensor. ## Whole-model performance The exact final CUDA plugin was compared with GB10 tuning disabled versus enabled. Each value is the median paired gain from three alternating process pairs. The code-copy workload keeps generated tokens, speculative acceptance, tokens per target forward, and target-forward counts identical between arms. | Qwen 3.8 27B DFlash2 workload | Decode improvement | |---|---:| | FP16 KV, batch 1, 4K | 1.19% | | FP16 KV, batch 1, 16K | 1.03% | | INT8 KV, batch 1, 4K | 1.73% | | INT8 KV, batch 1, 16K | 1.05% | | INT8 KV, batch 4, 4K | 0.16% | Batch 4 is effectively neutral, with paired changes from -0.17% to +0.41%, while every batch-1 pair improved. Median TTFT did not regress in these matched-acceptance comparisons. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>