[Model] Add Hopper FA4 relative attention for Inkling (#48858)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Jul 16, 2026 at 11:26 UTC
f61163e6c736ba2660982769c1d729411b44490e
2 files changed
+109
-5
tests/models/inkling/test_fa4_rel_attention.py
+47
@@ -13,6 +13,8 @@ guide::
13
with causal (and optionally sliding-window) masking handled by the backend.
14
"""
15
16
+import importlib
17
+
18
import pytest
19
import torch
20
@@ -21,6 +23,7 @@ from vllm.models.inkling.nvidia.attention import (
23
compute_log_scaling_tau,
24
)
25
from vllm.models.inkling.nvidia.ops.fa4_rel_attention import (
26
+ _use_sheared_bias,
27
inkling_fa4_num_splits,
28
inkling_fa4_rel_attention,
29
)
@@ -78,6 +81,23 @@ def test_num_splits_hopper_is_unsplit(monkeypatch):
81
)
82
83
84
+@pytest.mark.parametrize(
85
+ ("major", "expected"),
86
+ [(9, False), (10, True), (11, True), (12, False)],
87
+)
88
+def test_sheared_bias_architecture_selection(monkeypatch, major, expected):
89
+ monkeypatch.setattr(
90
+ current_platform,
91
+ "get_device_capability",
92
+ lambda: DeviceCapability(major=major, minor=0),
93
+ )
94
+ _use_sheared_bias.cache_clear()
95
+ try:
96
+ assert _use_sheared_bias() is expected
97
+ finally:
98
+ _use_sheared_bias.cache_clear()
99
+
100
+
101
@pytest.fixture
102
def blackwell_platform(monkeypatch):
103
monkeypatch.setattr(
@@ -251,6 +271,14 @@ def _run_case(seq_lens, num_heads, num_kv_heads, rel_extent, window_left, seed=0
271
window_size = (-1, -1) if window_left is None else (window_left, 0)
272
273
preallocated_out = torch.empty_like(q)
274
+ num_splits = inkling_fa4_num_splits(
275
+ is_local=window_left is not None,
276
+ batch_size=num_seqs,
277
+ max_query_len=max(q_lens),
278
+ num_heads=num_heads,
279
+ num_kv_heads=num_kv_heads,
280
+ max_kv_len=max(kv_lens),
281
+ )
282
out = inkling_fa4_rel_attention(
283
q,
284
key_cache,
@@ -264,6 +292,7 @@ def _run_case(seq_lens, num_heads, num_kv_heads, rel_extent, window_left, seed=0
292
window_size=window_size,
293
rel_extent=rel_extent,
294
rel_logits=rel_logits,
295
+ num_splits=num_splits,
296
out=preallocated_out,
297
)
298
assert out.data_ptr() == preallocated_out.data_ptr()
@@ -285,6 +314,24 @@ def _run_case(seq_lens, num_heads, num_kv_heads, rel_extent, window_left, seed=0
314
torch.testing.assert_close(out.float(), ref.float(), atol=2e-2, rtol=2e-2)
315
316
317
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA")
318
+@pytest.mark.skipif(
319
+ _cap is None or _cap.major < 9,
320
+ reason="FA4 score-mod requires Hopper+ (SM90+)",
321
+)
322
+@torch.inference_mode()
323
+def test_score_mod_relative_attention(monkeypatch):
324
+ module = importlib.import_module("vllm.models.inkling.nvidia.ops.fa4_rel_attention")
325
+ monkeypatch.setattr(module, "_use_sheared_bias", lambda: False)
326
+ _run_case(
327
+ [(64, 64), (1, 80)],
328
+ num_heads=4,
329
+ num_kv_heads=4,
330
+ rel_extent=128,
331
+ window_left=None,
332
+ )
333
+
334
+
335
@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA")
336
@pytest.mark.skipif(
337
_cap is None or _cap.major < 9,
vllm/models/inkling/nvidia/ops/fa4_rel_attention.py
+62
-5
@@ -2,6 +2,10 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from __future__ import annotations
4
5
+from collections.abc import Callable
6
+from functools import cache
7
+from typing import Any
8
+
9
import torch
10
11
from vllm.platforms import current_platform
@@ -12,6 +16,47 @@ def bucket_max_seqlen_q(max_seqlen_q: int) -> int:
16
return 1 << max(0, max_seqlen_q - 1).bit_length()
17
18
19
+@cache
20
+def _use_sheared_bias() -> bool:
21
+ capability = current_platform.get_device_capability()
22
+ return capability is not None and capability.major in (10, 11)
23
+
24
+
25
+@cache
26
+def _get_score_mod(rel_extent: int) -> Callable:
27
+ """Return the score modification that adds Inkling relative bias."""
28
+ import cutlass.cute as cute
29
+ from cutlass.cute import Float32
30
+
31
+ from vllm.vllm_flash_attn.cute.seqlen_info import SeqlenInfoQK
32
+
33
+ @cute.jit
34
+ def score_mod_rel_bias(
35
+ scores: cute.TensorSSA,
36
+ b_idx: cute.TensorSSA,
37
+ h_idx: cute.TensorSSA,
38
+ q_idx: cute.TensorSSA,
39
+ kv_idx: cute.TensorSSA,
40
+ seqlen_info: SeqlenInfoQK,
41
+ aux_tensors: list[cute.Tensor],
42
+ ) -> cute.TensorSSA:
43
+ rel_logits = aux_tensors[0]
44
+
45
+ seqlen_local_offset = seqlen_info.seqlen_k - seqlen_info.seqlen_q
46
+ rel_dist = (q_idx + seqlen_local_offset) - kv_idx
47
+ global_q_idx = seqlen_info.offset_q + q_idx
48
+
49
+ rel_dist_0 = rel_dist[0]
50
+ rel_idx = rel_dist_0 if rel_dist_0 >= 0 else 0
51
+ rel_idx = rel_idx if rel_idx < rel_extent else (rel_extent - 1)
52
+
53
+ rel_bias = rel_logits[global_q_idx[0], h_idx[0], rel_idx]
54
+ rel_bias = Float32(rel_bias) if rel_dist_0 == rel_idx else Float32(0.0)
55
+ return scores + rel_bias
56
+
57
+ return score_mod_rel_bias
58
+
59
+
60
def inkling_fa4_num_splits(
61
*,
62
is_local: bool,
@@ -21,7 +66,7 @@ def inkling_fa4_num_splits(
66
num_kv_heads: int,
67
max_kv_len: int,
68
) -> int:
24
- """Return the split-KV cap for FA4 with sheared relative bias."""
69
+ """Return the split-KV cap for Inkling relative attention."""
70
capability = current_platform.get_device_capability()
71
if capability is not None and capability.major == 9:
72
return 1
@@ -78,13 +123,25 @@ def inkling_fa4_rel_attention(
123
per-request KV lengths (``seqused_k``). ``rel_logits`` is
124
``(num_tokens, num_heads, rel_extent)``.
125
81
- The bias uses tml-fa4's sheared relative-bias layout.
126
+ Hopper uses standard FA4's score-mod gather. Blackwell uses tml-fa4's
127
+ sheared relative-bias layout.
128
"""
83
- from vllm.third_party.tml_fa4 import flash_attn_varlen_func
84
-
129
# cute uses (None, None) to mean "no window".
130
cute_window = (None, None) if window_size == (-1, -1) else window_size
131
132
+ rel_logits = rel_logits.contiguous()
133
+ if _use_sheared_bias():
134
+ from vllm.third_party.tml_fa4 import flash_attn_varlen_func
135
+
136
+ bias_kwargs: dict[str, Any] = {"rel_bias": rel_logits}
137
+ else:
138
+ from vllm.vllm_flash_attn.cute import flash_attn_varlen_func
139
+
140
+ bias_kwargs = {
141
+ "score_mod": _get_score_mod(rel_extent),
142
+ "aux_tensors": [rel_logits],
143
+ }
144
+
145
ret = flash_attn_varlen_func(
146
q=q,
147
k=key_cache,
@@ -99,7 +156,7 @@ def inkling_fa4_rel_attention(
156
num_splits=num_splits,
157
return_lse=False,
158
out=out,
102
- rel_bias=rel_logits.contiguous(),
159
+ **bias_kwargs,
160
)
161
if isinstance(ret, tuple):
162
return ret[0]