| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Compatibility shim for the ROCm relative-attention implementation. |
| 4 | |
| 5 | ROCm uses a Triton kernel which is compiled by vLLM's normal model warmup. The |
| 6 | NVIDIA path registers ahead-of-time CuTeDSL units; importing that provider on |
| 7 | ROCm would pull in CUDA-only tml-fa4 code. |
| 8 | """ |
| 9 | |
| 10 | from __future__ import annotations |
| 11 | |
| 12 | from dataclasses import dataclass |
| 13 | |
| 14 | import torch |
| 15 | |
| 16 | |
| 17 | @dataclass(frozen=True) |
| 18 | class InklingFA4WarmupConfig: |
| 19 | num_heads: int |
| 20 | num_kv_heads: int |
| 21 | head_dim: int |
| 22 | rel_extent: int |
| 23 | window_size: tuple[int, int] |
| 24 | is_local: bool |
| 25 | max_kv_len: int |
| 26 | dtype: torch.dtype |
| 27 | kv_dtype: torch.dtype |
| 28 | block_size: int |
| 29 | max_num_reqs: int |
| 30 | max_num_batched_tokens: int |
| 31 | |
| 32 | |
| 33 | def register_fa4_warmup(config: InklingFA4WarmupConfig) -> None: |
| 34 | """Triton compilation is triggered by the ordinary vLLM warmup forward.""" |
| 35 | del config |