[Perf][Model] Qwen4Exp QSA: sm_90 tuning table for _select_config (#57273)
Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com> Signed-off-by: Shuolei Wang <948904026@qq.com> Co-authored-by: Roger Wang <hey@rogerw.io>
Shuolei Wang committed
Sep 17, 2026 at 15:02 UTC
7bbce752b85f183eb34e27ea6df9647f80986060
1 file changed
+39
-1
vllm/models/qwen4_exp/nvidia/ops/qsa.py
+39
-1
@@ -22,6 +22,12 @@ def _is_sm120() -> bool:
22
return current_platform.get_device_capability() == (12, 0)
23
24
25
+@lru_cache(maxsize=1)
26
+def _is_sm90() -> bool:
27
+ """True on sm_90 (H100/H200/H20): selects the sm_90 tuning table."""
28
+ return current_platform.get_device_capability() == (9, 0)
29
+
30
+
31
@triton.jit(do_not_specialize=["num_rows", "num_requests"])
32
def _qsa_sparse_paged_gqa_splitk_kernel(
33
q_ptr,
@@ -509,6 +515,33 @@ def _select_sm120_config(
515
return 32, 4, 1
516
517
518
+def _select_sm90_config(
519
+ base_programs: int, use_prefill_config: bool, is_fp8: bool
520
+) -> tuple[int, int, int]:
521
+ """(block_n, target_splits, num_warps) tuned on sm_90 (H20)."""
522
+ if base_programs > 2048:
523
+ return (32, 1, 1) if use_prefill_config else (64, 1, 2)
524
+ if is_fp8:
525
+ if base_programs <= 24:
526
+ return 64, 64, 2
527
+ if base_programs <= 32:
528
+ return 32, 16, 1
529
+ if base_programs <= 64:
530
+ return 32, 8, 1
531
+ if base_programs <= 128:
532
+ return 32, 4, 1
533
+ if base_programs <= 256:
534
+ return 32, 8, 1
535
+ return 32, 4, 1
536
+ if base_programs <= 32:
537
+ return 64, 64, 2
538
+ if base_programs <= 64:
539
+ return 32, 16, 1
540
+ if base_programs <= 256:
541
+ return 32, 8, 1
542
+ return 32, 4, 1
543
+
544
+
545
def _select_config(
546
num_rows: int,
547
num_kv_heads: int,
@@ -521,13 +554,18 @@ def _select_config(
554
Keyed on base_programs = num_rows * num_kv_heads. The bp > 2048 region splits
555
on use_prefill_config (capture-stable: at FULL-graph capture max_query_len is
556
the uniform decode/verify length). This default table was tuned on GB300;
524
- sm_120 (RTX PRO 6000 Blackwell) dispatches to _select_sm120_config instead.
557
+ sm_120 (RTX PRO 6000 Blackwell) dispatches to _select_sm120_config instead,
558
+ and sm_90 (Hopper) to _select_sm90_config.
559
"""
560
base_programs = num_rows * num_kv_heads
561
if _is_sm120():
562
BLOCK_N, target_splits, num_warps = _select_sm120_config(
563
base_programs, use_prefill_config, is_fp8
564
)
565
+ elif _is_sm90():
566
+ BLOCK_N, target_splits, num_warps = _select_sm90_config(
567
+ base_programs, use_prefill_config, is_fp8
568
+ )
569
elif base_programs > 2048:
570
BLOCK_N, target_splits, num_warps = (
571
(32, 1, 1) if use_prefill_config else (64, 1, 2)