[CUDA] Remove memset nodes from decode CUDA graphs and prefer GEMV at small M

[Tianlei Wu (tianleiwu)]

Description

Speculative decoding on a hybrid (gated delta net + full attention) decoder replays one CUDA graph per verify step, and that graph carried up to about 130 memset nodes: serial split-K semaphores of the fpA_intB CUTLASS GEMM, paged XQA semaphores, and the VarlenCausalConvWithState state_update output. Inside a graph each memset node adds several microseconds of dependency latency, while a small kernel adds well under one. This PR clears those buffers with kernels and makes the weight-only tactic profiler stop picking CUTLASS over the CUDA GEMV on near ties at small M, which it does because it times L2-resident synthetic weights.

Summary of Changes

Replace in-graph memsets with kernels

File	Change
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/device/gemm_universal_base_compat.h	ClearSplitKSemaphoresKernel replaces cudaMemsetAsync for the serial split-K semaphore workspace (int-sized workspaces up to 4 MiB; anything else keeps the memset).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_impl_gen.cuh, xqa_paged_loader_impl.cuh	Paged XQA semaphores are cleared with onnxruntime::cuda::Fill<int32_t>.
onnxruntime/contrib_ops/cuda/bert/varlen_causal_conv_with_state.cc	The state_update output is zeroed with Fill<int32_t> when its size is a multiple of 4 bytes, and with cudaMemsetAsync otherwise.

fpA_intB tactic selection at small M

File	Change
onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h	New virtual getSelectionTime(m, tactic, time) hook. The default returns the measured time unchanged, so other profilers are unaffected.
onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h, .cc	For M < 16 (the range where the CUDA GEMV is a candidate), a CUTLASS tactic's time is multiplied by 1.1, so CUTLASS has to be more than 10% faster to replace the GEMV.

Testing

• Synthetic graph on RTX 4090 (sm_89, CUDA 13.3, 200 tiny kernels): 1.0 us per kernel alone, 12.1 us with a memset node after each kernel, and 1.8 us with a one-block zeroing kernel after each kernel.
• The profiler's own M = 8 timings for the Qwen3.8-27B projections put CUTLASS within about 1% of the GEMV on some shapes, for example q_proj 5120 -> 12288: 23.4 us CUTLASS vs 23.7 us GEMV. In the model, where weights stream from DRAM, the same shapes run 45.7 us on CUTLASS and 43.2 us on the GEMV. Across the projections, the GEMV was 4-30% faster in-model at M = 8.
• End-to-end on RTX 4090, ORT GenAI CUDA plugin EP, Qwen3.8-27B INT4 + INT8 paged KV + DFlash2 width 7, greedy:
  • Forcing the GEMV for every M < 16 bucket (a temporary experiment) cut the mean decode step from 27.28 to 26.75 ms over 16 prompts, in two interleaved repeats. With this PR the profiler picks the GEMV for every M = 8 projection.
  • nsys node-level traces of the target graph, per verify step: before, the memset nodes and their neighbouring gaps cost about 725 us (small-N GEMV counters 326 us, CUTLASS split-K semaphores 296 us, XQA 58 us, causal conv 45 us), and the idle time between graph nodes was 723 us. With this PR and the small-N counter change in [CUDA] Auto-tune small-M decode MatMul (cuBLAS vs small-N GEMV) and speed up MatMulNBits GEMV #32876, no memset nodes remain and the idle time between nodes is 306 us.
• Correctness through the plugin EP against a float64 reference, MatMulNBits 4-bit block 32 with M in {4, 8, 12, 16, 24, 32, 48, 64} on the split-K shapes 5120x6144, 1024x5120, 5120x17408 and 10240x5120: worst relative error 4.3e-3, identical to the build without this change. The trace confirms ClearSplitKSemaphoresKernel ran for every split-K launch.
• MMLU-Pro (100 items, same protocol as [CUDA] Auto-tune small-M decode MatMul (cuBLAS vs small-N GEMV) and speed up MatMulNBits GEMV #32876): 86/100 vs 88/100. The kernel selection change alters rounding, so 46 of the 100 generations diverge at some token. Two answers flipped from correct to incorrect, and all outputs were well formed (no unextracted answers or speculative-decoding failures).
• Suggested CI coverage: the MatMulNBits CUDA tests, PagedAttention / XQA tests, and the VarlenCausalConvWithState tests.

Motivation and Context

• The memsets sit in hot decode paths. Per verify step on the model above: 64 CUTLASS split-K clears (when CUTLASS is selected), 16 XQA clears, and 48 causal-conv clears.
• Timing the profiler with its synthetic weights evicted from L2 was tried in [CUDA] Auto-tune small-M decode MatMul (cuBLAS vs small-N GEMV) and speed up MatMulNBits GEMV #32876 and did not change end-to-end time, so this PR adds a selection margin rather than changing how tactics are timed.
• Complements [CUDA] Auto-tune small-M decode MatMul (cuBLAS vs small-N GEMV) and speed up MatMulNBits GEMV #32876, which makes the same change for the small-N fp16 GEMV counters.

Checklist

• Tests added/updated (behavior-preserving; covered by existing op tests)
• No breaking changes
• Documentation updated (not applicable)

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Baiju Meswani (baijumeswani)]

In varlen_causal_conv_with_state.cc:136-139, Fill<int32_t> uses 32-bit indexing. For a state update larger than 2^31 four-byte values, the index can wrap and write outside the buffer. Please use cudaMemsetAsync for larger outputs, or make Fill safe for 64-bit counts. A boundary test would help. Also, please check the fixed 10% CUTLASS timing penalty with weights larger than L2; it could choose a slower kernel. This is a comment review, not a request for changes.

[Baiju Meswani (baijumeswani)]

Non-blocking suggestion: The kernel choice now depends on GPU L2 cache size, but the saved choice is shared across GPUs. Two GPUs with the same SM and different L2 sizes could reuse a slower choice. Please save choices per GPU, or include L2 size in the cache key. This can be a follow-up.