| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | import vllm.models.inkling.nvidia.ops.fa4_rel_attention as fa4_rel_attention |
| 9 | from vllm.models.inkling.configs import InklingMMConfig, InklingModelConfig |
| 10 | from vllm.models.inkling.nvidia.ops.fa4_rel_attention import ( |
| 11 | InklingFA4RelAttentionKernel, |
| 12 | _num_warps_bucket, |
| 13 | bucket_max_seqlen_q, |
| 14 | inkling_fa4_num_splits, |
| 15 | ) |
| 16 | |
| 17 | |
| 18 | def _vllm_config_from_reference_config(config: SimpleNamespace) -> SimpleNamespace: |
| 19 | return SimpleNamespace( |
| 20 | model_config=SimpleNamespace( |
| 21 | hf_config=InklingMMConfig( |
| 22 | text_config=InklingModelConfig( |
| 23 | num_attention_heads=config.num_heads, |
| 24 | num_key_value_heads=config.num_kv_heads, |
| 25 | head_dim=config.head_dim, |
| 26 | rel_extent=config.rel_extent, |
| 27 | swa_num_attention_heads=config.num_heads, |
| 28 | swa_num_key_value_heads=config.num_kv_heads, |
| 29 | swa_head_dim=config.head_dim, |
| 30 | sliding_window_size=config.rel_extent, |
| 31 | ) |
| 32 | ), |
| 33 | max_model_len=config.max_kv_len, |
| 34 | dtype=config.dtype, |
| 35 | ), |
| 36 | cache_config=SimpleNamespace( |
| 37 | cache_dtype="auto", |
| 38 | block_size=config.block_size, |
| 39 | ), |
| 40 | scheduler_config=SimpleNamespace( |
| 41 | max_num_seqs=config.max_num_reqs, |
| 42 | max_num_batched_tokens=config.max_num_batched_tokens, |
| 43 | ), |
| 44 | ) |
| 45 | |
| 46 | |
| 47 | def test_bucket_max_seqlen_q(): |
| 48 | assert [bucket_max_seqlen_q(n) for n in range(1, 10)] == [ |
| 49 | 1, |
| 50 | 2, |
| 51 | 4, |
| 52 | 4, |
| 53 | 8, |
| 54 | 8, |
| 55 | 8, |
| 56 | 8, |
| 57 | 16, |
| 58 | ] |
| 59 | |
| 60 | |
| 61 | def test_warmup_enumerates_every_runtime_compile_class(monkeypatch): |
| 62 | monkeypatch.setattr( |
| 63 | fa4_rel_attention, |
| 64 | "get_tensor_model_parallel_world_size", |
| 65 | lambda: 1, |
| 66 | ) |
| 67 | config = SimpleNamespace( |
| 68 | num_heads=16, |
| 69 | num_kv_heads=2, |
| 70 | head_dim=128, |
| 71 | rel_extent=1024, |
| 72 | window_size=(-1, -1), |
| 73 | is_local=False, |
| 74 | max_kv_len=65536, |
| 75 | dtype=torch.bfloat16, |
| 76 | kv_dtype=torch.bfloat16, |
| 77 | block_size=16, |
| 78 | max_num_reqs=64, |
| 79 | max_num_batched_tokens=192, |
| 80 | ) |
| 81 | vllm_config = _vllm_config_from_reference_config(config) |
| 82 | kernel = InklingFA4RelAttentionKernel() |
| 83 | warmed_keys = set(kernel.get_warmup_keys(vllm_config)) |
| 84 | expected_keys: set[InklingFA4RelAttentionKernel.CompileKey] = set() |
| 85 | |
| 86 | for is_local in (False, True): |
| 87 | max_kv_len = config.rel_extent if is_local else config.max_kv_len |
| 88 | window_size = (config.rel_extent - 1, 0) if is_local else (-1, -1) |
| 89 | for query_len in range(1, config.max_num_batched_tokens + 1): |
| 90 | max_seqlen_q = bucket_max_seqlen_q(query_len) |
| 91 | max_num_reqs = min( |
| 92 | config.max_num_reqs, |
| 93 | config.max_num_batched_tokens - query_len + 1, |
| 94 | ) |
| 95 | for num_reqs in range(1, max_num_reqs + 1): |
| 96 | num_splits = inkling_fa4_num_splits( |
| 97 | is_local=is_local, |
| 98 | batch_size=num_reqs, |
| 99 | max_query_len=max_seqlen_q, |
| 100 | num_heads=config.num_heads, |
| 101 | num_kv_heads=config.num_kv_heads, |
| 102 | max_kv_len=max_kv_len, |
| 103 | ) |
| 104 | expected_keys.add( |
| 105 | InklingFA4RelAttentionKernel.CompileKey( |
| 106 | is_local=is_local, |
| 107 | num_heads=config.num_heads, |
| 108 | num_kv_heads=config.num_kv_heads, |
| 109 | head_dim=config.head_dim, |
| 110 | rel_extent=config.rel_extent, |
| 111 | dtype=config.dtype, |
| 112 | kv_dtype=config.kv_dtype, |
| 113 | block_size=config.block_size, |
| 114 | window_size=window_size, |
| 115 | max_seqlen_q=max_seqlen_q, |
| 116 | num_splits=num_splits, |
| 117 | num_warps_bucket=( |
| 118 | _num_warps_bucket(num_reqs) if num_splits > 1 else None |
| 119 | ), |
| 120 | large_num_reqs=num_reqs > 1024, |
| 121 | ) |
| 122 | ) |
| 123 | |
| 124 | assert warmed_keys == expected_keys |