[Tianlei Wu (tianleiwu)]

Summary

Stacked on #33786 ( lwu/20261031/qmoe-cuda-int2-tc).

This PR optimizes CUDA prefill and decode kernels for Qwen3.8-Flash-Next (RC5) on Blackwell SM 12.x (100 KB shared memory/SM, RTX Spark N1X) from short/medium contexts (2K–32K) up to 256K (262,144) context length:

1. DynamicSparseAttention (dynamic_sparse_attention.cc, dynamic_sparse_attention_impl.cu, dynamic_sparse_attention_prefill_int8.cuh, dynamic_sparse_attention_helper.h):
   • Enable DynamicSparseAttentionInt8TcKernel prefill & split-K decode by default on SM >= 8.0.
   • Add a 2-stage shared-memory pipeline (kTwoBufferSmem, 38.1 KB / 40.4 KB smem) on 100 KB/SM architectures (SM 8.6/8.9/12.0/12.1) to double occupancy from 1 to 2 CTAs/SM (2.83x faster DSA prefill, 7.5x faster DSA decode).
   • Clamp otal_sequence_length to past_cache_capacity when a smaller runtime past_key cache is bound via
     untime_profiles.
2. PackedSparseAttentionIndexer (packed_sparse_attention_indexer.cc, packed_sparse_attention_indexer_impl.cu):
   • Tile 8 query tokens per CTA (4 in registers + 4 in shared memory) in QsaRaggedPrefillScoreKernel<T, TState, true> and replace shared-memory reduction barriers with a bit-exact 19-shuffle 16-head butterfly reduction, cutting present_key_state DRAM reads by 2x and speeding up 128K prefill scoring by 1.96x (11.67 s -> 5.94 s).
   • Avoid zero-filling inactive 2048-key chunks in QsaScoreTileTopKKernel beyond element-0 sentinels, eliminating the 256K state_capacity padding tax on short/medium prompts.
   • Aggregate radix histogram updates in QsaPartialTopKKernel with warp-level __match_any_sync and extend single-launch QsaFinalMergeAndEmitKernel up to 64 chunks (256K context).
   • Allow runtime past_key_state capacity to be sized dynamically from dynamic_batching /
     untime_profiles.
3. HyperConnectionLowRankMix (hyper_connection_low_rank_mix_impl.cu):
   • Tune kWeightSlots to 5 on CUDA_ARCH >= 1000 (30.8 KB in HcDownKernel, 23.6 KB in HcUpPreKernel), increasing 100 KB/SM occupancy from 2 to 3–4 CTAs/SM (-34% to -43% kernel time).
4. QMoE Decode & Prefill (moe_gemv_tc.cuh, qmoe_int3_tc.cuh, qmoe_kernels.cu, int8_tc_gemv.cuh):
   • Increase kRouteChunks to 3 (38.7 KB smem -> 2 CTAs/SM) for otal_routes <= 96 and enable kM = 2 2-route FC2 tile sharing in MoeGemvTcInt2Kernel.
   • Enable kRemMode = 1 FC2 remainder prefetch in Int3TcCoalescedKernel.

Validation

• End-to-end smoke.py and 10-item MMLU-Pro smoke evaluation across all 3 Qwen3.8-Flash-Next RC5 variants (2b375, 3b3_16, all_int4).
• 2K / 8K / 32K and 8K / 128K / 255K + 1K (262,144 total tokens) TTFT & decode benchmarks on RTX Spark N1X (SM 12.1):
  • 2b375: 32K TTFT 24.85s -> 19.63s (+26.6%), 128K TTFT 85.96s -> 75.06s (-10.91s), 255K TTFT 205.73s -> 164.07s (-41.66s, +25.4%), 128K decode 97.13 tok/s.
  • 3b3_16: 32K TTFT 29.96s -> 24.47s (+22.4%), 128K TTFT 101.07s -> 86.50s (-14.57s), 255K TTFT 234.31s -> 197.38s (-36.94s, +18.7%).

[Copilot]

🟡 Changes recommended

Unresolved critical issues can discard valid top-k candidates, admit stale workspace data, and break builds targeting pre-SM70 GPUs and tests on SM90.

5 open issues

• Update GEMV launch-configuration expectations and test table · New
• Initialize the full suffix for noncompact top-k mode · New
• Round up active chunk boundary to prevent sentinel/live-key overlap · New
• Guard __match_any_sync for targets below SM70 · New
• Use runtime tensor capacity when validating reuse · New

What changed in this PR

This PR builds on #33786 to optimize ONNX Runtime’s CUDA kernels for Qwen3.8-Flash-Next, targeting SM12x GPUs and contexts up to 256K tokens.

Changes:

• Enables INT8 tensor-core sparse attention by default and adds a lower-shared-memory pipeline.
• Optimizes sparse-indexer scoring and top-k selection, while allowing runtime-sized caches.
• Tunes HyperConnection and quantized MoE kernel dispatch, buffering, and prefetching.

File	Description
onnxruntime/​contrib_ops/​cuda/​sparse/​packed_sparse_attention_indexer.cc	Allows runtime cache capacities in QSA and CSA refresh paths.
onnxruntime/​contrib_ops/​cuda/​sparse/​packed_sparse_attention_indexer_impl.cu	Tiles query scoring and optimizes top-k initialization, histograms, and merging.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_kernels.cu	Makes scalar fallback launch-grid conversions explicit.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_int3_tc.cuh	Enables FC2 remainder prefetch and adjusts kernel configurations.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemv_tc.cuh	Specializes small-route discovery and widens eligible INT2 FC2 tiles.
onnxruntime/​contrib_ops/​cuda/​llm/​fpA_intB_gemv/​int8_tc_gemv.cuh	Simplifies load addressing and retunes INT8 GEMV dispatch.
onnxruntime/​contrib_ops/​cuda/​bert/​hyper_connection_low_rank_mix_impl.cu	Reduces shared-memory weight slots on selected architectures.
onnxruntime/​contrib_ops/​cuda/​bert/​dynamic_sparse_attention.cc	Enables INT8 tensor-core prefill and decode by default.
onnxruntime/​contrib_ops/​cuda/​bert/​dynamic_sparse_attention_prefill_int8.cuh	Adds a two-buffer shared-memory attention pipeline.
onnxruntime/​contrib_ops/​cuda/​bert/​dynamic_sparse_attention_impl.cu	Selects and launches attention variants based on shared-memory capacity.
onnxruntime/​contrib_ops/​cpu/​bert/​dynamic_sparse_attention_helper.h	Clamps total sequence length to the bound past-cache capacity.

_{Review effort: Balanced}