iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf][Model] Qwen4Exp QSA: sm_90 tuning table for _select_config (#57273)

Browse Source 
Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com>
Signed-off-by: Shuolei Wang <948904026@qq.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
This commit is contained in:
Shuolei Wang
2026-09-17 07:02:24 +00:00
committed by GitHub
co-authored by Roger Wang
parent bd2d7ae7c1
commit 7bbce752b8
1 changed files with 39 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/models/qwen4_exp/nvidia/ops/qsa.py
+39 -1
View File
@@ -22,6 +22,12 @@ def _is_sm120() -> bool:
return current_platform.get_device_capability() == (12, 0)
@lru_cache(maxsize=1)
def _is_sm90() -> bool:
"""True on sm_90 (H100/H200/H20): selects the sm_90 tuning table."""
return current_platform.get_device_capability() == (9, 0)
@triton.jit(do_not_specialize=["num_rows", "num_requests"])
def _qsa_sparse_paged_gqa_splitk_kernel(
q_ptr,
@@ -509,6 +515,33 @@ def _select_sm120_config(
return 32, 4, 1
def _select_sm90_config(
base_programs: int, use_prefill_config: bool, is_fp8: bool
) -> tuple[int, int, int]:
"""(block_n, target_splits, num_warps) tuned on sm_90 (H20)."""
if base_programs > 2048:
return (32, 1, 1) if use_prefill_config else (64, 1, 2)
if is_fp8:
if base_programs <= 24:
return 64, 64, 2
if base_programs <= 32:
return 32, 16, 1
if base_programs <= 64:
return 32, 8, 1
if base_programs <= 128:
return 32, 4, 1
if base_programs <= 256:
return 32, 8, 1
return 32, 4, 1
if base_programs <= 32:
return 64, 64, 2
if base_programs <= 64:
return 32, 16, 1
if base_programs <= 256:
return 32, 8, 1
return 32, 4, 1
def _select_config(
num_rows: int,
num_kv_heads: int,
@@ -521,13 +554,18 @@ def _select_config(
Keyed on base_programs = num_rows * num_kv_heads. The bp > 2048 region splits
on use_prefill_config (capture-stable: at FULL-graph capture max_query_len is
the uniform decode/verify length). This default table was tuned on GB300;
sm_120 (RTX PRO 6000 Blackwell) dispatches to _select_sm120_config instead.
sm_120 (RTX PRO 6000 Blackwell) dispatches to _select_sm120_config instead,
and sm_90 (Hopper) to _select_sm90_config.
"""
base_programs = num_rows * num_kv_heads
if _is_sm120():
BLOCK_N, target_splits, num_warps = _select_sm120_config(
base_programs, use_prefill_config, is_fp8
)
elif _is_sm90():
BLOCK_N, target_splits, num_warps = _select_sm90_config(
base_programs, use_prefill_config, is_fp8
)
elif base_programs > 2048:
BLOCK_N, target_splits, num_warps = (
(32, 1, 1) if use_prefill_config else (64, 1, 2)