| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | from vllm.model_executor.warmup.qwen_triton_warmup import ( |
| 8 | _FLA_POST_CONV_WARMUP_LENGTHS, |
| 9 | _QwenGDNWarmupConfig, |
| 10 | _warm_causal_conv1d_fwd_kernel, |
| 11 | _warm_fused_post_conv_kernel, |
| 12 | _warm_gated_rms_norm_kernel, |
| 13 | ) |
| 14 | from vllm.platforms import current_platform |
| 15 | |
| 16 | |
| 17 | def _cuda_gdn_config() -> _QwenGDNWarmupConfig: |
| 18 | h, hv, k, v = 2, 2, 16, 16 |
| 19 | conv_kernel_size = 4 |
| 20 | conv_dim = 2 * h * k + hv * v |
| 21 | device = torch.device("cuda") |
| 22 | conv_state = torch.empty( |
| 23 | (8, conv_dim, conv_kernel_size - 1), |
| 24 | dtype=torch.bfloat16, |
| 25 | device=device, |
| 26 | ) |
| 27 | return _QwenGDNWarmupConfig( |
| 28 | h=h, |
| 29 | hv=hv, |
| 30 | k=k, |
| 31 | v=v, |
| 32 | conv_kernel_size=conv_kernel_size, |
| 33 | conv_state=conv_state, |
| 34 | conv_dtype=conv_state.dtype, |
| 35 | norm_weight_dtype=torch.bfloat16, |
| 36 | norm_before_gate=True, |
| 37 | norm_activation="silu", |
| 38 | a_log=torch.zeros(hv, dtype=torch.float32, device=device), |
| 39 | dt_bias=torch.zeros(hv, dtype=torch.float32, device=device), |
| 40 | state_stride_token=hv * v * k, |
| 41 | state_dtype=torch.float32, |
| 42 | ) |
| 43 | |
| 44 | |
| 45 | @pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="CUDA is required") |
| 46 | def test_qwen_gdn_prefill_warmup_kernels_compile_on_gpu() -> None: |
| 47 | config = _cuda_gdn_config() |
| 48 | device = torch.device("cuda") |
| 49 | _warm_gated_rms_norm_kernel( |
| 50 | device, config, max_num_tokens=16, x_dtype=config.conv_dtype |
| 51 | ) |
| 52 | _warm_causal_conv1d_fwd_kernel(device, config) |
| 53 | _warm_fused_post_conv_kernel(device, config) |
| 54 | assert _FLA_POST_CONV_WARMUP_LENGTHS == (1, 2, 16) |
| 55 | torch.accelerator.synchronize(device) |