Add FP8 GEMV KSplit32 scheduling for the 48-SM SM121 GPU

[Baiju Meswani (baijumeswani)]

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.

[Copilot]

🟡 Changes recommended

The KSplit32 dispatch may instantiate a kernel exceeding CUDA’s static shared-memory limit.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds FP8 GEMV KSplit32 scheduling for low-SM-count SM12x GPUs.

Changes:

• Adds device- and shape-aware K-split selection.
• Adds KSplit32 kernel dispatch.
• Tests scheduling decisions.

Findings:

• Critical: Dispatch instantiates an unreachable <32, 4> kernel requiring 64 KiB static shared memory. Launch only the reachable <32, 1> specialization directly.
• Nit: Add a functional test that forces KSplit32 rather than testing only heuristic selection.

File summaries

File	Description
onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc	Tests K-split selection.
onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu	Integrates selection and KSplit32 dispatch.
onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h	Defines scheduling heuristics.

Review details

Suppressed comments (1)

onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc:69

• This test validates only the host-side selection result. None of the existing functional cases execute the new KSplit=32 specialization on the usual SM8/SM9 test devices, so the 1024-thread launch and 32-way reduction can fail without this test noticing. Please add a functional result test with a test seam that forces KSplit 32 (covering FP16/BF16 as appropriate), rather than testing only the heuristic.

    EXPECT_EQ(onnxruntime::contrib::cuda::PickFp8MmaKSplit(
                  c.n, c.m, c.windows, c.sm_count, c.compute_capability_major),
              c.expected);

• Files reviewed: 3/3 changed files
• Comments generated: 1
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

The tests are not hermetic and may bypass KSplit32 when the existing kill switch is disabled.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc:134

• CUDA availability is checked only in the child, but a skipped GoogleTest process exits successfully. On a CUDA build without a usable device, the parent therefore records this test as passed instead of skipped. Check availability before spawning so the test result accurately reports that the KSplit32 kernel was not exercised.

  if (!is_child_process) {
    const std::string command =
        "\"" + CurrentExecutablePath() +
        "\" --gtest_filter=MatMulBlockQuantizedFp8WeightOpTest.GemvTensorCoreForcedKSplit32 --gtest_color=no";
    ASSERT_EQ(std::system(command.c_str()), 0);

• Files reviewed: 3/3 changed files
• Comments generated: 1
• Review effort level: Balanced

[Tianlei Wu (tianleiwu)]

Reviewed the current head. The KSplit32 schedule is narrowly gated to the qualified 48-SM SM121 configuration, recursive tiles preserve the generic schedule, and the launch instantiates only the valid <32, 1> kernel. The selector boundary cases and forced FP16/BF16 test cover the new behavior, including a ragged 33-window reduction. The prior shared-memory and test-isolation concerns are addressed; I found no remaining actionable issues.