| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Connect loaded Qwen4Exp QSA modules to their kernel-owned warmup.""" |
| 4 | |
| 5 | import sys |
| 6 | from typing import TYPE_CHECKING, cast |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | from vllm.logger import init_logger |
| 11 | |
| 12 | if TYPE_CHECKING: |
| 13 | from vllm.v1.worker.gpu.model_runner import GPUModelRunner as GPUModelRunnerV2 |
| 14 | from vllm.v1.worker.gpu_worker import Worker |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | |
| 19 | def qwen4_exp_qsa_triton_warmup(worker: "Worker") -> None: |
| 20 | """Warm every reachable QSA specialization: indexer decode-query-length |
| 21 | profiles plus the sparse attention split-K/merge configs.""" |
| 22 | qsa_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.indexer_qsa") |
| 23 | attn_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.qsa") |
| 24 | if qsa_module is None or attn_module is None: |
| 25 | return |
| 26 | indexer = None |
| 27 | owner = None |
| 28 | for layer in worker.get_model().modules(): |
| 29 | if indexer is None and isinstance(layer, qsa_module.QSAIndexer): |
| 30 | indexer = layer |
| 31 | elif owner is None and isinstance(layer, attn_module.Qwen4ExpQSAAttention): |
| 32 | owner = layer |
| 33 | if indexer is None or owner is None: |
| 34 | return |
| 35 | |
| 36 | runner = worker.model_runner |
| 37 | |
| 38 | def block_table_for(prefix: str) -> torch.Tensor: |
| 39 | group_id = next( |
| 40 | i |
| 41 | for i, group in enumerate(runner.kv_cache_config.kv_cache_groups) |
| 42 | if prefix in group.layer_names |
| 43 | ) |
| 44 | if worker.use_v2_model_runner: |
| 45 | runner_v2 = cast("GPUModelRunnerV2", runner) |
| 46 | return runner_v2.block_tables.input_block_tables[group_id] |
| 47 | return runner.input_batch.block_table[group_id].get_device_tensor( |
| 48 | runner.max_num_reqs |
| 49 | ) |
| 50 | |
| 51 | if worker.use_v2_model_runner: |
| 52 | max_decode_query_len = cast("GPUModelRunnerV2", runner).decode_query_len |
| 53 | else: |
| 54 | max_decode_query_len = runner.uniform_decode_query_len |
| 55 | |
| 56 | from vllm.models.qwen4_exp.nvidia.ops.qsa_indexer import ( |
| 57 | warmup_qsa_mqa_paged_decode, |
| 58 | ) |
| 59 | |
| 60 | k_cache = indexer.compressed_key_cache.kv_cache |
| 61 | assert k_cache.numel() |
| 62 | profiles = warmup_qsa_mqa_paged_decode( |
| 63 | k_cache, |
| 64 | block_table_for(indexer.compressed_key_cache.prefix), |
| 65 | num_heads=indexer.index_n_heads, |
| 66 | head_dim=indexer.index_head_dim, |
| 67 | max_decode_query_len=max_decode_query_len, |
| 68 | max_num_reqs=runner.max_num_reqs, |
| 69 | max_num_batched_tokens=runner.max_num_tokens, |
| 70 | ) |
| 71 | logger.info("Warmed up Qwen4Exp QSA decode kernels: %s.", profiles) |
| 72 | |
| 73 | from vllm.models.qwen4_exp.nvidia.ops.qsa import ( |
| 74 | warmup_qsa_sparse_paged_attention, |
| 75 | ) |
| 76 | |
| 77 | kv_cache = owner.kv_cache |
| 78 | assert kv_cache.numel() |
| 79 | attention_profiles = warmup_qsa_sparse_paged_attention( |
| 80 | kv_cache, |
| 81 | block_table_for(owner.layer_name), |
| 82 | num_query_heads=owner.num_heads, |
| 83 | selection_width=indexer.output_width, |
| 84 | ) |
| 85 | logger.info( |
| 86 | "Warmed up Qwen4Exp QSA sparse attention kernels: %s.", |
| 87 | attention_profiles, |
| 88 | ) |