From 4a69044529f850f21766f338c36b9d7f525c82ac Mon Sep 17 00:00:00 2001 From: Peuqui Date: Thu, 3 Sep 2026 11:39:29 +0200 Subject: [PATCH] [Bugfix][Perf][SM70/SM75] Widen QSA sparse launch profile to pre-Ampere The GB300 prefill profiles select a 64-column tile that cannot launch on SM75 at D=256 (Triton OutOfResources against Turing's 64 KiB shared memory), and the existing SM70 retune does not apply there. Gate the branch on compute capability < 8.0 and pick the 16-column four-warp tile, which launches on both pre-Ampere architectures and measures 1.16-2.6x faster than the best previously runnable profile across the 64..2048-row prefill regimes (V100-PCIE-32GB, Quadro RTX 8000; #441). Small-batch decode profiles and the sm80+ table are unchanged. Co-authored-by: Claude Signed-off-by: Peuqui --- .../qwen4_exp/test_qsa_launch_profile.py | 21 ++++++++++--------- tests/models/qwen4_exp/test_qsa_ops.py | 10 ++++----- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 20 ++++++++++-------- 3 files changed, 27 insertions(+), 24 deletions(-) diff --git a/tests/models/qwen4_exp/test_qsa_launch_profile.py b/tests/models/qwen4_exp/test_qsa_launch_profile.py index f27083726f..bedba47f10 100644 --- a/tests/models/qwen4_exp/test_qsa_launch_profile.py +++ b/tests/models/qwen4_exp/test_qsa_launch_profile.py @@ -9,23 +9,24 @@ @pytest.mark.parametrize( - ("block_m", "base_programs", "is_sm70", "expected"), + ("block_m", "base_programs", "is_pre_ampere", "expected"), [ - pytest.param(8, 8, True, (16, 64, 4), id="sm70_small_block_m8"), - pytest.param(16, 4, True, (16, 64, 4), id="sm70_small_block_m16"), - pytest.param(16, 5, True, (16, 32, 4), id="sm70_narrow"), - pytest.param(8, 256, True, (64, 8, 4), id="sm70_split8"), - pytest.param(8, 512, True, (32, 4, 4), id="sm70_split4"), - pytest.param(8, 513, False, (64, 1, 2), id="non_sm70_split1"), - pytest.param(8, 513, True, (32, 1, 4), id="sm70_split1"), + pytest.param(8, 8, True, (16, 64, 4), id="pre_ampere_small_block_m8"), + pytest.param(16, 4, True, (16, 64, 4), id="pre_ampere_small_block_m16"), + pytest.param(16, 5, True, (16, 32, 4), id="pre_ampere_narrow"), + pytest.param(8, 256, True, (16, 8, 4), id="pre_ampere_split8"), + pytest.param(8, 512, True, (16, 4, 4), id="pre_ampere_split4"), + pytest.param(8, 513, False, (64, 1, 2), id="ampere_split1"), + pytest.param(8, 513, True, (16, 1, 4), id="pre_ampere_split1"), ], ) def test_qsa_sparse_launch_profile( block_m: int, base_programs: int, - is_sm70: bool, + is_pre_ampere: bool, expected: tuple[int, int, int], ) -> None: assert ( - qsa_ops._qsa_sparse_launch_profile(base_programs, block_m, is_sm70) == expected + qsa_ops._qsa_sparse_launch_profile(base_programs, block_m, is_pre_ampere) + == expected ) diff --git a/tests/models/qwen4_exp/test_qsa_ops.py b/tests/models/qwen4_exp/test_qsa_ops.py index b1e8651d8f..0ff5f00316 100644 --- a/tests/models/qwen4_exp/test_qsa_ops.py +++ b/tests/models/qwen4_exp/test_qsa_ops.py @@ -19,13 +19,13 @@ pytestmark = pytest.mark.skip_global_cleanup -def test_sm70_qsa_prefill_uses_narrow_tiles_and_four_warps(): - assert _qsa_sparse_launch_profile(511, 8, True) == (64, 4, 4) - assert _qsa_sparse_launch_profile(512, 8, True) == (32, 4, 4) - assert _qsa_sparse_launch_profile(8192, 8, True) == (32, 1, 4) +def test_pre_ampere_qsa_prefill_uses_narrow_tiles_and_four_warps(): + assert _qsa_sparse_launch_profile(511, 8, True) == (16, 4, 4) + assert _qsa_sparse_launch_profile(512, 8, True) == (16, 4, 4) + assert _qsa_sparse_launch_profile(8192, 8, True) == (16, 1, 4) -def test_non_sm70_qsa_prefill_keeps_gb300_profile(): +def test_ampere_qsa_prefill_keeps_gb300_profile(): assert _qsa_sparse_launch_profile(512, 8, False) == (64, 4, 2) assert _qsa_sparse_launch_profile(8192, 8, False) == (64, 1, 2) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index f00236eb0a..7407ef4ded 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -1900,7 +1900,7 @@ def qsa_sparse_paged_attention( block_n, target_splits, partial_warps = _qsa_sparse_launch_profile( base_programs, block_m, - current_platform.is_device_capability(70), + not current_platform.has_device_capability(80), ) if ( @@ -2001,7 +2001,7 @@ def qsa_sparse_paged_attention( def _qsa_sparse_launch_profile( base_programs: int, block_m: int, - is_sm70: bool, + is_pre_ampere: bool, ) -> tuple[int, int, int]: """Return BLOCK_N, target splits, and warps for sparse QSA.""" small_profile_limit = 8 if block_m <= 8 else 4 @@ -2018,14 +2018,16 @@ def _qsa_sparse_launch_profile( block_n, target_splits, partial_warps = 64, 4, 2 else: block_n, target_splits, partial_warps = 64, 1, 2 - if is_sm70 and block_n == 64: - # Two warps serialize the D=256 tensor-core work on V100. Four warps - # restore warp-level parallelism for split and non-split prefill. + if is_pre_ampere and block_n == 64: + # Pre-Ampere: the 64-column tile at D=256 does not fit Turing's + # 64 KiB shared-memory limit (Triton OutOfResources -- the kernel + # cannot launch on SM75 at all), and two warps serialize the D=256 + # tensor-core work on V100. A 16-column tile with four warps + # launches on both and measured 1.16-2.6x faster than the best + # previously runnable profile across the 64..2048-row prefill + # regimes (V100-PCIE-32GB and Quadro RTX 8000, see #441). partial_warps = 4 - if base_programs >= 512: - # A 32-column tile improves the exact 512-row and 8192-row Qwen4Exp - # prefill shapes without changing small-batch or non-SM70 routes. - block_n = 32 + block_n = 16 return block_n, target_splits, partial_warps