[ROCm] [Model] Enable TML inkling (#48841)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: OpenAI Codex <codex@openai.com>
TJian committed
Jul 27, 2026 at 18:05 UTC
0906123953a5a253cc0cc4f26e548ca08ece102a
33 files changed
+9213
-9
benchmarks/kernels/benchmark_inkling_qkvr_prep.py
new
+176
@@ -0,0 +1,176 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import statistics
5
+
6
+import torch
7
+from tabulate import tabulate
8
+
9
+from vllm.models.inkling.nvidia.ops import qkvr_prep
10
+from vllm.utils.argparse_utils import FlexibleArgumentParser
11
+
12
+
13
+def make_inputs(tokens: int, tp_size: int, is_local: bool):
14
+ torch.manual_seed(0)
15
+ num_q_heads = 64 // tp_size
16
+ num_kv_heads = (16 if is_local else 8) // tp_size
17
+ head_dim = 128
18
+ d_rel = 16
19
+ rel_extent = 512 if is_local else 1024
20
+ page_size = 16
21
+ num_blocks = (tokens + page_size - 1) // page_size
22
+ q_width = num_q_heads * head_dim
23
+ kv_width = num_kv_heads * head_dim
24
+ r_width = num_q_heads * d_rel
25
+ device = "cuda"
26
+
27
+ qkvr = torch.randn(
28
+ tokens,
29
+ q_width + 2 * kv_width + r_width,
30
+ device=device,
31
+ dtype=torch.bfloat16,
32
+ )
33
+ k_weight = torch.randn(kv_width, 4, device=device, dtype=torch.bfloat16)
34
+ v_weight = torch.randn_like(k_weight)
35
+ q_norm_weight = torch.randn(head_dim, device=device, dtype=torch.bfloat16)
36
+ k_norm_weight = torch.randn_like(q_norm_weight)
37
+ rel_proj = torch.randn(d_rel, rel_extent, device=device, dtype=torch.bfloat16)
38
+ conv_cache = torch.zeros(
39
+ num_blocks,
40
+ num_kv_heads,
41
+ page_size,
42
+ 2 * head_dim,
43
+ device=device,
44
+ dtype=torch.bfloat16,
45
+ )
46
+ key_cache = torch.empty(
47
+ num_blocks,
48
+ page_size,
49
+ num_kv_heads,
50
+ head_dim,
51
+ device=device,
52
+ dtype=torch.bfloat16,
53
+ )
54
+ value_cache = torch.empty_like(key_cache)
55
+ positions = torch.arange(tokens, device=device, dtype=torch.int64)
56
+ block_table = torch.arange(num_blocks, device=device, dtype=torch.int32)[None]
57
+ seq_idx = torch.zeros(tokens, device=device, dtype=torch.int32)
58
+ slots = torch.arange(tokens, device=device, dtype=torch.int64)
59
+ query_start = torch.zeros(tokens, device=device, dtype=torch.int32)
60
+ log_scaling = None
61
+ if not is_local:
62
+ effective_n = (positions + 1).to(torch.float32)
63
+ log_scaling = 1.0 + 0.1 * torch.log(torch.clamp(effective_n / 128000, min=1.0))
64
+ return (
65
+ qkvr,
66
+ k_weight,
67
+ v_weight,
68
+ q_norm_weight,
69
+ k_norm_weight,
70
+ rel_proj,
71
+ 1e-6,
72
+ num_q_heads,
73
+ num_kv_heads,
74
+ head_dim,
75
+ d_rel,
76
+ conv_cache,
77
+ key_cache,
78
+ value_cache,
79
+ positions,
80
+ block_table,
81
+ seq_idx,
82
+ slots,
83
+ query_start,
84
+ slots,
85
+ 0,
86
+ head_dim,
87
+ page_size,
88
+ log_scaling,
89
+ )
90
+
91
+
92
+def capture(implementation, inputs):
93
+ outputs = []
94
+
95
+ def run():
96
+ outputs[:] = implementation.fused_qkvr_prep(*inputs)
97
+
98
+ stream = torch.cuda.Stream()
99
+ stream.wait_stream(torch.cuda.current_stream())
100
+ with torch.cuda.stream(stream):
101
+ for _ in range(3):
102
+ run()
103
+ torch.cuda.current_stream().wait_stream(stream)
104
+ torch.accelerator.synchronize()
105
+ graph = torch.cuda.CUDAGraph()
106
+ with torch.cuda.graph(graph):
107
+ run()
108
+ torch.accelerator.synchronize()
109
+ return graph, outputs
110
+
111
+
112
+def time_graph(graph: torch.cuda.CUDAGraph, warmup: int, repeats: int) -> float:
113
+ for _ in range(warmup):
114
+ graph.replay()
115
+ torch.accelerator.synchronize()
116
+ start = torch.cuda.Event(enable_timing=True)
117
+ end = torch.cuda.Event(enable_timing=True)
118
+ start.record()
119
+ for _ in range(repeats):
120
+ graph.replay()
121
+ end.record()
122
+ end.synchronize()
123
+ return start.elapsed_time(end) * 1000 / repeats
124
+
125
+
126
+def benchmark(inputs, args) -> float:
127
+ graph, _ = capture(qkvr_prep, inputs)
128
+ return statistics.median(
129
+ time_graph(graph, args.warmup, args.repeats) for _ in range(args.trials)
130
+ )
131
+
132
+
133
+@torch.inference_mode()
134
+def main(args):
135
+ rows = []
136
+ for tp_size in args.tp_sizes:
137
+ for tokens in args.tokens:
138
+ for is_local in (True, False):
139
+ triton_us = benchmark(make_inputs(tokens, tp_size, is_local), args)
140
+ rows.append(
141
+ [
142
+ tp_size,
143
+ tokens,
144
+ "local" if is_local else "global",
145
+ triton_us,
146
+ ]
147
+ )
148
+
149
+ print("Inkling QKVR prep (CUDA graph, median latency)")
150
+ print(
151
+ tabulate(
152
+ rows,
153
+ headers=[
154
+ "TP",
155
+ "tokens",
156
+ "scope",
157
+ "Triton (us)",
158
+ ],
159
+ floatfmt=("d", "d", "", ".2f"),
160
+ )
161
+ )
162
+
163
+
164
+if __name__ == "__main__":
165
+ parser = FlexibleArgumentParser()
166
+ parser.add_argument(
167
+ "--tokens",
168
+ type=int,
169
+ nargs="+",
170
+ default=[1 << power for power in range(15)],
171
+ )
172
+ parser.add_argument("--tp-sizes", type=int, nargs="+", default=[4, 8])
173
+ parser.add_argument("--warmup", type=int, default=20)
174
+ parser.add_argument("--repeats", type=int, default=200)
175
+ parser.add_argument("--trials", type=int, default=5)
176
+ main(parser.parse_args())
tests/models/inkling/rocm/conftest.py
new
+17
@@ -0,0 +1,17 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from vllm.platforms import current_platform
7
+
8
+
9
+@pytest.fixture(autouse=True)
10
+def require_rocm_cdna4() -> None:
11
+ if not current_platform.is_rocm():
12
+ pytest.skip("requires ROCm")
13
+
14
+ from vllm.platforms.rocm import on_gfx950
15
+
16
+ if not on_gfx950():
17
+ pytest.skip("requires CDNA4")
tests/models/inkling/rocm/test_model_alignment.py
new
+59
@@ -0,0 +1,59 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ROCm guards for shared Inkling model-definition contracts."""
4
+
5
+from pathlib import Path
6
+
7
+import torch
8
+
9
+import vllm.models.inkling as inkling
10
+from vllm.lora.utils import get_supported_lora_modules
11
+from vllm.model_executor.models.interfaces import supports_lora
12
+from vllm.models.inkling.amd import model as amd_model
13
+from vllm.models.inkling.amd.mtp import InklingMTP
14
+
15
+
16
+def test_backend_neutral_definitions_match_nvidia_reference() -> None:
17
+ inkling_dir = Path(amd_model.__file__).parents[1]
18
+ for filename in ("logits_processor.py", "mtp.py"):
19
+ assert (inkling_dir / "amd" / filename).read_bytes() == (
20
+ inkling_dir / "nvidia" / filename
21
+ ).read_bytes()
22
+
23
+
24
+def test_rocm_exports_its_aligned_mtp_implementation() -> None:
25
+ assert inkling.InklingMTP is InklingMTP
26
+
27
+
28
+def test_rocm_model_exposes_the_upstream_lora_contract() -> None:
29
+ model_cls = amd_model._TmlForCausalLMBase
30
+
31
+ assert supports_lora(model_cls)
32
+ assert model_cls.packed_modules_mapping == {
33
+ "qkvr": ["wq_du", "wk_dv", "wv_dv", "wr_du"],
34
+ "w13": ["w1", "w3"],
35
+ }
36
+ assert model_cls.embedding_modules == {"lm_head": "output_embeddings"}
37
+
38
+ model = torch.nn.Module()
39
+ model.embedding_modules = model_cls.embedding_modules
40
+ supported = get_supported_lora_modules(model)
41
+ assert "embed_tokens" not in supported
42
+ assert "lm_head" in supported
43
+
44
+
45
+def test_lightseek_bundled_adapter_weights_remain_opt_in() -> None:
46
+ mapper = amd_model._TmlForCausalLMBase.hf_to_vllm_mapper
47
+ weight = torch.empty(1)
48
+ name, mapped_weight = next(
49
+ iter(
50
+ mapper.apply([("language_model.layers.3.attn.wq_du.lora_A.weight", weight)])
51
+ )
52
+ )
53
+
54
+ assert name == "model.layers.3.attn.qkvr.lora_A.weight"
55
+ assert mapped_weight.shard_id == 0
56
+ assert amd_model._is_peft_adapter_weight(name)
57
+
58
+ lm_head_name = mapper.apply_list(["language_model.lm_head.lora_B.weight"])
59
+ assert lm_head_name == ["lm_head.lora_B.weight"]
tests/models/inkling/rocm/test_mxfp4_load.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for Quark OCP MXFP4 weights loaded into AITER MoE layouts."""
4
+
5
+from types import SimpleNamespace
6
+
7
+import torch
8
+from torch import nn
9
+
10
+from vllm.models.inkling.amd.moe import InklingMoE
11
+
12
+
13
+def _fake_moe() -> tuple[InklingMoE, SimpleNamespace]:
14
+ moe = InklingMoE.__new__(InklingMoE)
15
+ nn.Module.__init__(moe)
16
+ moe.n_routed_experts = 4
17
+ routed = SimpleNamespace(
18
+ moe_config=SimpleNamespace(
19
+ moe_parallel_config=SimpleNamespace(tp_rank=1, tp_size=2)
20
+ ),
21
+ # Logical intermediate per rank is 3; AITER pads it to 4.
22
+ w13_weight=nn.Parameter(torch.zeros(4, 8, 2, dtype=torch.uint8), False),
23
+ w2_weight=nn.Parameter(torch.zeros(4, 5, 3, dtype=torch.uint8), False),
24
+ w13_weight_scale=nn.Parameter(torch.ones(4, 8, 2, dtype=torch.uint8), False),
25
+ w2_weight_scale=nn.Parameter(torch.ones(4, 5, 3, dtype=torch.uint8), False),
26
+ )
27
+ moe.experts = SimpleNamespace(routed_experts=routed)
28
+ moe._local_expert_slots = ( # type: ignore[method-assign]
29
+ lambda: dict(enumerate(range(4)))
30
+ )
31
+ return moe, routed
32
+
33
+
34
+def test_loads_logical_tp_shards_and_preserves_aiter_padding():
35
+ moe, routed = _fake_moe()
36
+
37
+ w13 = torch.arange(4 * 12 * 2, dtype=torch.uint8).view(4, 12, 2)
38
+ w2 = torch.arange(4 * 5 * 4, dtype=torch.uint8).view(4, 5, 4)
39
+ moe.load_expert_weight("experts.w13_weight", w13)
40
+ moe.load_expert_weight("experts.w2_weight", w2)
41
+
42
+ # TP rank 1 consumes the second six interleaved gate/up rows. AITER keeps
43
+ # one padding row in each destination half and one padding w2 column.
44
+ torch.testing.assert_close(routed.w13_weight[:, :3], w13[:, 6:12:2])
45
+ torch.testing.assert_close(routed.w13_weight[:, 4:7], w13[:, 7:12:2])
46
+ assert torch.count_nonzero(routed.w13_weight[:, 3]) == 0
47
+ assert torch.count_nonzero(routed.w13_weight[:, 7]) == 0
48
+ torch.testing.assert_close(routed.w2_weight[:, :, :2], w2[:, :, 2:4])
49
+ assert torch.count_nonzero(routed.w2_weight[:, :, 2]) == 0
50
+
51
+
52
+def test_unflattens_quark_scales_and_preserves_scale_padding():
53
+ moe, routed = _fake_moe()
54
+
55
+ flat_w13_scale = torch.arange(4 * 12 * 2, dtype=torch.uint8).view(4 * 12, 2)
56
+ flat_w2_scale = torch.arange(4 * 5 * 4, dtype=torch.uint8).view(4 * 5, 4)
57
+ moe.load_expert_weight("experts.w13_weight_scale", flat_w13_scale)
58
+ moe.load_expert_weight("experts.w2_weight_scale", flat_w2_scale)
59
+
60
+ w13_scale = flat_w13_scale.view(4, 12, 2)
61
+ w2_scale = flat_w2_scale.view(4, 5, 4)
62
+ torch.testing.assert_close(routed.w13_weight_scale[:, :3], w13_scale[:, 6:12:2])
63
+ torch.testing.assert_close(routed.w13_weight_scale[:, 4:7], w13_scale[:, 7:12:2])
64
+ assert torch.all(routed.w13_weight_scale[:, 3] == 1)
65
+ assert torch.all(routed.w13_weight_scale[:, 7] == 1)
66
+ torch.testing.assert_close(routed.w2_weight_scale[:, :, :2], w2_scale[:, :, 2:4])
67
+ assert torch.all(routed.w2_weight_scale[:, :, 2] == 1)
tests/models/inkling/rocm/test_rel_attention.py
new
+425
@@ -0,0 +1,425 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Correctness tests for Inkling's ROCm relative-bias paged attention."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.models.inkling.amd.ops.fa4_rel_attention import (
9
+ bucket_max_seqlen_q,
10
+ inkling_fa4_rel_attention,
11
+ use_gfx950_gluon_decode,
12
+ use_gfx950_gluon_extend,
13
+)
14
+from vllm.models.inkling.amd.ops.rel_attention_decode import (
15
+ decode_split_count,
16
+ use_split_kv_decode,
17
+)
18
+
19
+HEAD_DIM = 128
20
+DTYPE = torch.bfloat16
21
+
22
+
23
+def _reference(
24
+ q: torch.Tensor,
25
+ key_cache: torch.Tensor,
26
+ value_cache: torch.Tensor,
27
+ rel_logits: torch.Tensor,
28
+ block_table: torch.Tensor,
29
+ q_lens: list[int],
30
+ kv_lens: list[int],
31
+ rel_extent: int,
32
+ window_left: int | None,
33
+) -> torch.Tensor:
34
+ num_heads = q.shape[1]
35
+ num_kv_heads = key_cache.shape[2]
36
+ gqa_group = num_heads // num_kv_heads
37
+ out: list[torch.Tensor] = []
38
+ q_start = 0
39
+
40
+ for req, (q_len, kv_len) in enumerate(zip(q_lens, kv_lens)):
41
+ page_size = key_cache.shape[1]
42
+ num_pages = (kv_len + page_size - 1) // page_size
43
+ page_ids = block_table[req, :num_pages].long()
44
+ k = key_cache[page_ids].reshape(-1, num_kv_heads, HEAD_DIM)[:kv_len]
45
+ v = value_cache[page_ids].reshape(-1, num_kv_heads, HEAD_DIM)[:kv_len]
46
+ k = k.float().repeat_interleave(gqa_group, dim=1)
47
+ v = v.float().repeat_interleave(gqa_group, dim=1)
48
+
49
+ q_req = q[q_start : q_start + q_len].float()
50
+ rel_req = rel_logits[q_start : q_start + q_len].float()
51
+ scores = torch.einsum("qhd,khd->hqk", q_req, k) / HEAD_DIM
52
+
53
+ q_pos = torch.arange(q_len, device=q.device)[:, None] + kv_len - q_len
54
+ k_pos = torch.arange(kv_len, device=q.device)[None, :]
55
+ distance = q_pos - k_pos
56
+ rel_idx = distance.clamp(0, rel_extent - 1)
57
+ bias = rel_req.permute(1, 0, 2).gather(
58
+ 2, rel_idx[None].expand(num_heads, -1, -1)
59
+ )
60
+ in_rel_extent = (distance >= 0) & (distance < rel_extent)
61
+ scores += torch.where(in_rel_extent[None], bias, 0.0)
62
+
63
+ masked = distance < 0
64
+ if window_left is not None:
65
+ masked |= distance > window_left
66
+ scores.masked_fill_(masked[None], float("-inf"))
67
+ out.append(torch.einsum("hqk,khd->qhd", scores.softmax(-1), v))
68
+ q_start += q_len
69
+
70
+ return torch.cat(out).to(DTYPE)
71
+
72
+
73
+def test_query_length_bucket():
74
+ assert bucket_max_seqlen_q(1) == 1
75
+ assert bucket_max_seqlen_q(17) == 32
76
+ assert bucket_max_seqlen_q(33) == 64
77
+
78
+
79
+def test_split_kv_dispatch_thresholds():
80
+ assert not use_split_kv_decode(
81
+ max_query_len=4,
82
+ max_kv_len=65536,
83
+ page_size=128,
84
+ window_left=-1,
85
+ )
86
+ assert not use_split_kv_decode(
87
+ max_query_len=1,
88
+ max_kv_len=4096,
89
+ page_size=16,
90
+ window_left=-1,
91
+ )
92
+ assert use_split_kv_decode(
93
+ max_query_len=1,
94
+ max_kv_len=8192,
95
+ page_size=16,
96
+ window_left=-1,
97
+ )
98
+ assert use_split_kv_decode(
99
+ max_query_len=1,
100
+ max_kv_len=512,
101
+ page_size=128,
102
+ window_left=511,
103
+ )
104
+ assert decode_split_count(512, 511) == 4
105
+ assert decode_split_count(65536, -1) == 32
106
+
107
+
108
+def test_gfx950_gluon_dispatch_is_strict(monkeypatch: pytest.MonkeyPatch):
109
+ import vllm.models.inkling.amd.ops.fa4_rel_attention as rel_attention
110
+
111
+ monkeypatch.setattr(rel_attention, "on_gfx950", lambda: True)
112
+ assert use_gfx950_gluon_decode(max_query_len=1, page_size=128, head_dim=128)
113
+ assert not use_gfx950_gluon_decode(max_query_len=1, page_size=16, head_dim=128)
114
+ assert not use_gfx950_gluon_decode(max_query_len=4, page_size=128, head_dim=128)
115
+ assert not use_gfx950_gluon_decode(max_query_len=1, page_size=128, head_dim=256)
116
+
117
+ monkeypatch.setattr(rel_attention, "on_gfx950", lambda: False)
118
+ assert not use_gfx950_gluon_decode(max_query_len=1, page_size=128, head_dim=128)
119
+
120
+ monkeypatch.setattr(rel_attention, "on_gfx950", lambda: True)
121
+ assert use_gfx950_gluon_extend(
122
+ max_query_len=4,
123
+ max_kv_len=8192,
124
+ page_size=128,
125
+ head_dim=128,
126
+ window_left=-1,
127
+ )
128
+ assert not use_gfx950_gluon_extend(
129
+ max_query_len=4,
130
+ max_kv_len=8192,
131
+ page_size=128,
132
+ head_dim=128,
133
+ window_left=511,
134
+ )
135
+ assert not use_gfx950_gluon_extend(
136
+ max_query_len=4,
137
+ max_kv_len=8192,
138
+ page_size=16,
139
+ head_dim=128,
140
+ window_left=-1,
141
+ )
142
+
143
+
144
+@pytest.mark.parametrize(
145
+ ("page_size", "window_left", "max_kv_len"),
146
+ [
147
+ (16, None, None),
148
+ (16, 15, None),
149
+ (128, None, 8192),
150
+ ],
151
+)
152
+@torch.inference_mode()
153
+def test_ragged_multi_page_relative_attention(
154
+ page_size: int,
155
+ window_left: int | None,
156
+ max_kv_len: int | None,
157
+):
158
+ """Covers full/chunked prefill, decode, GQA, and the local window."""
159
+ torch.manual_seed(19 + int(window_left is not None))
160
+ device = "cuda"
161
+ q_lens = [17, 1]
162
+ kv_lens = [35, 80]
163
+ num_heads = 8
164
+ num_kv_heads = 2
165
+ rel_extent = 16
166
+ max_pages = max((length + page_size - 1) // page_size for length in kv_lens)
167
+
168
+ q = torch.randn(sum(q_lens), num_heads, HEAD_DIM, device=device)
169
+ q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE)
170
+ key_cache = torch.randn(
171
+ 1 + len(q_lens) * max_pages,
172
+ page_size,
173
+ num_kv_heads,
174
+ HEAD_DIM,
175
+ device=device,
176
+ )
177
+ key_cache = torch.nn.functional.normalize(key_cache.float(), dim=-1).to(DTYPE)
178
+ value_cache = torch.randn_like(key_cache)
179
+ rel_logits = torch.randn(
180
+ sum(q_lens), num_heads, rel_extent, device=device, dtype=DTYPE
181
+ )
182
+
183
+ block_table = torch.stack(
184
+ [
185
+ torch.arange(1 + req * max_pages, 1 + (req + 1) * max_pages)
186
+ for req in range(len(q_lens))
187
+ ]
188
+ ).to(device=device, dtype=torch.int32)
189
+ cu_seqlens_q = torch.tensor(
190
+ [0, *torch.tensor(q_lens).cumsum(0).tolist()],
191
+ device=device,
192
+ dtype=torch.int32,
193
+ )
194
+ cache_seqlens = torch.tensor(kv_lens, device=device, dtype=torch.int32)
195
+ window_size = (-1, -1) if window_left is None else (window_left, 0)
196
+
197
+ preallocated = torch.empty_like(q)
198
+ actual = inkling_fa4_rel_attention(
199
+ q,
200
+ key_cache,
201
+ value_cache,
202
+ block_table=block_table,
203
+ cache_seqlens=cache_seqlens,
204
+ cu_seqlens_q=cu_seqlens_q,
205
+ max_seqlen_q=bucket_max_seqlen_q(max(q_lens)),
206
+ softmax_scale=1.0 / HEAD_DIM,
207
+ causal=True,
208
+ window_size=window_size,
209
+ rel_extent=rel_extent,
210
+ rel_logits=rel_logits,
211
+ max_kv_len=max_kv_len,
212
+ out=preallocated,
213
+ )
214
+ assert actual.data_ptr() == preallocated.data_ptr()
215
+
216
+ expected = _reference(
217
+ q,
218
+ key_cache,
219
+ value_cache,
220
+ rel_logits,
221
+ block_table,
222
+ q_lens,
223
+ kv_lens,
224
+ rel_extent,
225
+ window_left,
226
+ )
227
+ torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=3e-2)
228
+
229
+
230
+@pytest.mark.parametrize(
231
+ ("page_size", "kv_lens", "num_kv_heads", "rel_extent", "window_left"),
232
+ [
233
+ (16, [257, 513], 1, 64, None),
234
+ (128, [257, 512], 2, 512, 511),
235
+ ],
236
+)
237
+@torch.inference_mode()
238
+def test_split_kv_decode_matches_reference(
239
+ page_size: int,
240
+ kv_lens: list[int],
241
+ num_kv_heads: int,
242
+ rel_extent: int,
243
+ window_left: int | None,
244
+):
245
+ """Covers page-16 split-KV and page-128 gfx950 Gluon decode."""
246
+ torch.manual_seed(41 + page_size)
247
+ device = "cuda"
248
+ batch_size = len(kv_lens)
249
+ q_lens = [1] * batch_size
250
+ num_heads = 8
251
+ max_pages = max((length + page_size - 1) // page_size for length in kv_lens)
252
+
253
+ q = torch.randn(batch_size, num_heads, HEAD_DIM, device=device)
254
+ q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE)
255
+ key_cache = torch.randn(
256
+ batch_size * max_pages,
257
+ page_size,
258
+ num_kv_heads,
259
+ HEAD_DIM,
260
+ device=device,
261
+ )
262
+ key_cache = torch.nn.functional.normalize(key_cache.float(), dim=-1).to(DTYPE)
263
+ value_cache = torch.randn_like(key_cache)
264
+ rel_logits = torch.randn(
265
+ batch_size,
266
+ num_heads,
267
+ rel_extent,
268
+ device=device,
269
+ dtype=DTYPE,
270
+ )
271
+ block_table = torch.arange(
272
+ batch_size * max_pages,
273
+ device=device,
274
+ dtype=torch.int32,
275
+ ).view(batch_size, max_pages)
276
+ cache_seqlens = torch.tensor(kv_lens, device=device, dtype=torch.int32)
277
+ cu_seqlens_q = torch.arange(
278
+ batch_size + 1,
279
+ device=device,
280
+ dtype=torch.int32,
281
+ )
282
+ window_size = (-1, -1) if window_left is None else (window_left, 0)
283
+ max_kv_len = 8192 if page_size == 16 else max(kv_lens)
284
+
285
+ actual = inkling_fa4_rel_attention(
286
+ q,
287
+ key_cache,
288
+ value_cache,
289
+ block_table=block_table,
290
+ cache_seqlens=cache_seqlens,
291
+ cu_seqlens_q=cu_seqlens_q,
292
+ max_seqlen_q=1,
293
+ softmax_scale=1.0 / HEAD_DIM,
294
+ causal=True,
295
+ window_size=window_size,
296
+ rel_extent=rel_extent,
297
+ rel_logits=rel_logits,
298
+ max_kv_len=max_kv_len,
299
+ )
300
+ expected = _reference(
301
+ q,
302
+ key_cache,
303
+ value_cache,
304
+ rel_logits,
305
+ block_table,
306
+ q_lens,
307
+ kv_lens,
308
+ rel_extent,
309
+ window_left,
310
+ )
311
+ torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=3e-2)
312
+
313
+
314
+@pytest.mark.parametrize(
315
+ (
316
+ "q_len",
317
+ "kv_len",
318
+ "num_kv_heads",
319
+ "rel_extent",
320
+ "window_left",
321
+ "gluon_enabled",
322
+ ),
323
+ [
324
+ (1, 513, 2, 512, 511, True),
325
+ (1, 513, 2, 512, 511, False),
326
+ (4, 8192, 1, 1024, None, True),
327
+ ],
328
+)
329
+@torch.inference_mode()
330
+def test_gfx950_page128_packed_kv_views_match_reference(
331
+ monkeypatch: pytest.MonkeyPatch,
332
+ q_len: int,
333
+ kv_len: int,
334
+ num_kv_heads: int,
335
+ rel_extent: int,
336
+ window_left: int | None,
337
+ gluon_enabled: bool,
338
+):
339
+ """Cover both Gluon and split-KV against vLLM's packed KV allocation."""
340
+ monkeypatch.setenv(
341
+ "INKLING_GFX950_GLUON",
342
+ "1" if gluon_enabled else "0",
343
+ )
344
+ torch.manual_seed(71 + q_len)
345
+ device = "cuda"
346
+ page_size = 128
347
+ batch_size = 2
348
+ num_heads = 8
349
+ q_lens = [q_len] * batch_size
350
+ kv_lens = [kv_len] * batch_size
351
+ pages_per_req = (kv_len + page_size - 1) // page_size
352
+
353
+ q = torch.randn(
354
+ batch_size * q_len,
355
+ num_heads,
356
+ HEAD_DIM,
357
+ device=device,
358
+ dtype=DTYPE,
359
+ )
360
+ q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE)
361
+
362
+ # FlashAttentionBackend allocates logical [block, head, page, 2 * dim].
363
+ # Inkling transposes to [block, page, head, 2 * dim] and splits K/V,
364
+ # producing non-contiguous views whose page/head strides include both.
365
+ packed_kv = torch.randn(
366
+ batch_size * pages_per_req,
367
+ num_kv_heads,
368
+ page_size,
369
+ 2 * HEAD_DIM,
370
+ device=device,
371
+ dtype=DTYPE,
372
+ )
373
+ key_cache, value_cache = packed_kv.transpose(1, 2).split(HEAD_DIM, dim=-1)
374
+ assert not key_cache.is_contiguous()
375
+ assert key_cache.stride() == value_cache.stride()
376
+
377
+ rel_logits = torch.randn(
378
+ batch_size * q_len,
379
+ num_heads,
380
+ rel_extent,
381
+ device=device,
382
+ dtype=DTYPE,
383
+ )
384
+ block_table = torch.arange(
385
+ batch_size * pages_per_req,
386
+ device=device,
387
+ dtype=torch.int32,
388
+ ).view(batch_size, pages_per_req)
389
+ cache_seqlens = torch.tensor(kv_lens, device=device, dtype=torch.int32)
390
+ cu_seqlens_q = torch.arange(
391
+ 0,
392
+ batch_size * q_len + 1,
393
+ q_len,
394
+ device=device,
395
+ dtype=torch.int32,
396
+ )
397
+ window_size = (-1, -1) if window_left is None else (window_left, 0)
398
+
399
+ actual = inkling_fa4_rel_attention(
400
+ q,
401
+ key_cache,
402
+ value_cache,
403
+ block_table=block_table,
404
+ cache_seqlens=cache_seqlens,
405
+ cu_seqlens_q=cu_seqlens_q,
406
+ max_seqlen_q=q_len,
407
+ softmax_scale=1.0 / HEAD_DIM,
408
+ causal=True,
409
+ window_size=window_size,
410
+ rel_extent=rel_extent,
411
+ rel_logits=rel_logits,
412
+ max_kv_len=kv_len,
413
+ )
414
+ expected = _reference(
415
+ q,
416
+ key_cache,
417
+ value_cache,
418
+ rel_logits,
419
+ block_table,
420
+ q_lens,
421
+ kv_lens,
422
+ rel_extent,
423
+ window_left,
424
+ )
425
+ torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=3e-2)
tests/models/inkling/rocm/test_rocm_mtp_input_fusion.py
new
+116
@@ -0,0 +1,116 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ROCm numerical-equivalence tests for the fused MTP depth-layer input kernel.
4
+
5
+``embed_dual_rmsnorm_cat`` must match the unfused module sequence numerically:
6
+each rmsnorm computes in fp32 and rounds to bf16 at the same points as the
7
+vendored ``rmsnorm`` kernel (including the bf16 round-trip between the
8
+chained backbone embed_norm and the depth embed_norm), and the fused row
9
+gather matches ``F.embedding``.
10
+"""
11
+
12
+from typing import cast
13
+
14
+import pytest
15
+import torch
16
+
17
+from vllm.models.inkling.amd.ops.norm import (
18
+ embed_dual_rmsnorm_cat,
19
+ embed_rmsnorm,
20
+ rmsnorm,
21
+)
22
+
23
+EPS = 1e-6
24
+VOCAB = 4096
25
+BF16_ATOL = torch.finfo(torch.bfloat16).eps
26
+
27
+
28
+def _assert_bf16_close(actual: torch.Tensor, expected: torch.Tensor) -> None:
29
+ # Different Triton reduction schedules can move a small number of outputs
30
+ # by one bf16 epsilon while remaining numerically equivalent.
31
+ torch.testing.assert_close(actual, expected, rtol=0, atol=BF16_ATOL)
32
+
33
+
34
+def _ref(hidden, w_h, w_e, emb, w_pre=None):
35
+ if w_pre is not None:
36
+ emb = rmsnorm(emb, w_pre, EPS)
37
+ return torch.cat([rmsnorm(hidden, w_h, EPS), rmsnorm(emb, w_e, EPS)], dim=-1)
38
+
39
+
40
+@pytest.mark.parametrize("n", [1536, 6144])
41
+@pytest.mark.parametrize("t", [0, 1, 7, 256])
42
+@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64])
43
+def test_embed_dual_rmsnorm_cat(n: int, t: int, ids_dtype: torch.dtype) -> None:
44
+ torch.manual_seed(0)
45
+ dev = "cuda"
46
+ table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16)
47
+ w_h = torch.randn(n, device=dev).to(torch.bfloat16)
48
+ w_e = (1 + 0.01 * torch.randn(n, device=dev)).to(torch.bfloat16)
49
+ w_pre = torch.randn(n, device=dev).to(torch.bfloat16)
50
+ ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype)
51
+ hidden = (torch.randn(t, n, device=dev) * 2).to(torch.bfloat16)
52
+ emb = table[ids.long()]
53
+
54
+ # Fused gather + chained backbone pre-norm (the decode draft-step path).
55
+ out = embed_dual_rmsnorm_cat(
56
+ hidden,
57
+ w_h,
58
+ w_e,
59
+ EPS,
60
+ input_ids=ids,
61
+ embed_table=table,
62
+ pre_norm_weight=w_pre,
63
+ )
64
+ assert out.shape == (t, 2 * n)
65
+ _assert_bf16_close(out, _ref(hidden, w_h, w_e, emb, w_pre))
66
+
67
+ # Precomputed embeds, no pre-norm (draft prefill with target-merged MM
68
+ # embeddings, already backbone-normed).
69
+ out = embed_dual_rmsnorm_cat(hidden, w_h, w_e, EPS, embeds=emb)
70
+ _assert_bf16_close(out, _ref(hidden, w_h, w_e, emb))
71
+
72
+ # Fused gather, no pre-norm (use_embed_norm=False).
73
+ out = embed_dual_rmsnorm_cat(
74
+ hidden, w_h, w_e, EPS, input_ids=ids, embed_table=table
75
+ )
76
+ _assert_bf16_close(out, _ref(hidden, w_h, w_e, emb))
77
+
78
+
79
+@pytest.mark.parametrize("n", [1536, 6144])
80
+@pytest.mark.parametrize("t", [0, 1, 7, 256])
81
+@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64])
82
+def test_embed_rmsnorm(n: int, t: int, ids_dtype: torch.dtype) -> None:
83
+ torch.manual_seed(0)
84
+ dev = "cuda"
85
+ table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16)
86
+ w = torch.randn(n, device=dev).to(torch.bfloat16)
87
+ ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype)
88
+ ref_emb = table[ids.long()]
89
+
90
+ # Gather + embed_norm (base model / MTP prefill embed path).
91
+ out = cast(torch.Tensor, embed_rmsnorm(ids, table, w, EPS))
92
+ assert out.shape == (t, n)
93
+ _assert_bf16_close(out, rmsnorm(ref_emb, w, EPS) if t else ref_emb)
94
+
95
+ # Pure gather (use_embed_norm=False / replicated module forward).
96
+ out = cast(torch.Tensor, embed_rmsnorm(ids, table, None, EPS))
97
+ assert torch.equal(out, ref_emb)
98
+
99
+ # Chained first-layer attn_norm (the target text-path forward): one launch
100
+ # emits both the residual and layer 0's normed attention input.
101
+ w_chain = (1 + 0.05 * torch.randn(n, device=dev)).to(torch.bfloat16)
102
+ res, attn_in = cast(
103
+ tuple[torch.Tensor, torch.Tensor],
104
+ embed_rmsnorm(ids, table, w, EPS, chain_weight=w_chain),
105
+ )
106
+ ref_res = rmsnorm(ref_emb, w, EPS) if t else ref_emb
107
+ _assert_bf16_close(res, ref_res)
108
+ _assert_bf16_close(attn_in, rmsnorm(ref_res, w_chain, EPS) if t else ref_res)
109
+
110
+ # Chained without embed_norm (use_embed_norm=False).
111
+ res, attn_in = cast(
112
+ tuple[torch.Tensor, torch.Tensor],
113
+ embed_rmsnorm(ids, table, None, EPS, chain_weight=w_chain),
114
+ )
115
+ assert torch.equal(res, ref_emb)
116
+ _assert_bf16_close(attn_in, rmsnorm(ref_emb, w_chain, EPS) if t else ref_emb)
tests/models/inkling/rocm/test_sconv_cache_layout.py
new
+20
@@ -0,0 +1,20 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+from torch import nn
6
+
7
+from vllm.models.inkling.amd.sconv_swa_attn import InklingConvState
8
+
9
+
10
+def test_runtime_sconv_block_size_tracks_unified_cache_page():
11
+ """The cache planner may enlarge W=4 blocks to match attention pages."""
12
+ owner = InklingConvState.__new__(InklingConvState)
13
+ nn.Module.__init__(owner)
14
+ owner.block_size = 4
15
+
16
+ owner.kv_cache = torch.tensor([])
17
+ assert owner.cache_block_size == 4
18
+
19
+ owner.kv_cache = torch.empty(2, 1, 32, 1024)
20
+ assert owner.cache_block_size == 32
vllm/models/inkling/__init__.py
+30
-9
@@ -2,12 +2,21 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from typing import TYPE_CHECKING
4
5
+from vllm.platforms import current_platform
6
+
7
if TYPE_CHECKING:
6
- from .nvidia.model import (
7
- InklingForCausalLM,
8
- InklingForConditionalGeneration,
9
- )
10
- from .nvidia.mtp import InklingMTP
8
+ if current_platform.is_rocm():
9
+ from .amd.model import (
10
+ InklingForCausalLM,
11
+ InklingForConditionalGeneration,
12
+ )
13
+ from .amd.mtp import InklingMTP as InklingMTP
14
+ else:
15
+ from .nvidia.model import (
16
+ InklingForCausalLM,
17
+ InklingForConditionalGeneration,
18
+ )
19
+ from .nvidia.mtp import InklingMTP as InklingMTP
20
21
__all__ = [
22
"InklingForConditionalGeneration",
@@ -18,11 +27,23 @@ __all__ = [
27
28
def __getattr__(name: str):
29
if name == "InklingMTP":
21
- from .nvidia import mtp
30
+ if current_platform.is_rocm():
31
+ from .amd import mtp as amd_mtp
32
+
33
+ return amd_mtp.InklingMTP
34
+
35
+ from .nvidia import mtp as nvidia_mtp
36
+
37
+ return nvidia_mtp.InklingMTP
38
23
- return mtp.InklingMTP
39
if name in __all__:
25
- from .nvidia import model
40
+ if current_platform.is_rocm():
41
+ from .amd import model as amd_model
42
+
43
+ return getattr(amd_model, name)
44
+
45
+ from .nvidia import model as nvidia_model
46
+
47
+ return getattr(nvidia_model, name)
48
27
- return getattr(model, name)
49
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
vllm/models/inkling/amd/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/inkling/amd/attention.py
new
+328
@@ -0,0 +1,328 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from __future__ import annotations
4
+
5
+from typing import cast
6
+
7
+import torch
8
+from torch import nn
9
+
10
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
11
+from vllm.config import VllmConfig, get_current_vllm_config
12
+from vllm.distributed import get_tensor_model_parallel_world_size
13
+from vllm.forward_context import get_forward_context
14
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
15
+from vllm.model_executor.layers.linear import (
16
+ MergedColumnParallelLinear,
17
+ RowParallelLinear,
18
+)
19
+from vllm.model_executor.layers.quantization import QuantizationConfig
20
+from vllm.utils.torch_utils import (
21
+ canonicalize_singleton_dim_strides,
22
+ kv_cache_dtype_str_to_dtype,
23
+)
24
+from vllm.v1.attention.backend import AttentionBackend
25
+from vllm.v1.attention.backends.flash_attn import (
26
+ FlashAttentionBackend,
27
+ FlashAttentionMetadata,
28
+)
29
+from vllm.v1.kv_cache_interface import (
30
+ FullAttentionSpec,
31
+ KVCacheSpec,
32
+ SlidingWindowSpec,
33
+)
34
+
35
+from ..configs import InklingModelConfig
36
+from .layernorm import InklingRMSNorm
37
+from .ops.fa4_rel_attention import (
38
+ bucket_max_seqlen_q,
39
+ inkling_fa4_num_splits,
40
+ inkling_fa4_rel_attention,
41
+)
42
+from .ops.fa4_warmup import InklingFA4WarmupConfig, register_fa4_warmup
43
+from .ops.qkvr_prep import fused_qkvr_prep
44
+from .sconv_swa_attn import _K, _V, InklingConvState, InklingSconvMetadata
45
+from .short_conv import InklingShortConv
46
+
47
+
48
+def compute_log_scaling_tau(
49
+ positions: torch.Tensor, n_floor: int, alpha: float
50
+) -> torch.Tensor:
51
+ effective_n = (positions + 1).to(torch.float32)
52
+ return 1.0 + alpha * torch.log(torch.clamp(effective_n / float(n_floor), min=1.0))
53
+
54
+
55
+class RelLogitsProj(nn.Module):
56
+ """Project the per-head relative branch ``r`` to per-distance logits."""
57
+
58
+ def __init__(self, d_rel: int, rel_extent: int) -> None:
59
+ super().__init__()
60
+ self.d_rel = d_rel
61
+ self.rel_extent = rel_extent
62
+ self.proj = nn.Parameter(torch.empty(d_rel, rel_extent), requires_grad=False)
63
+
64
+ def forward(self, r_out: torch.Tensor) -> torch.Tensor:
65
+ # r_out: (T, num_heads, d_rel) -> (T, num_heads, rel_extent)
66
+ return torch.einsum("thd,de->the", r_out, self.proj)
67
+
68
+
69
+class InklingAttention(nn.Module, AttentionLayerBase):
70
+ def __init__(
71
+ self,
72
+ config: InklingModelConfig,
73
+ *,
74
+ num_heads: int,
75
+ num_kv_heads: int,
76
+ head_dim: int,
77
+ rel_extent: int,
78
+ local_extent: int,
79
+ is_local: bool,
80
+ prefix: str,
81
+ quant_config: QuantizationConfig | None = None,
82
+ conv_owner: InklingConvState,
83
+ ) -> None:
84
+ super().__init__()
85
+ self.prefix = prefix
86
+ self.is_local = is_local
87
+ self.hidden_size = config.hidden_size
88
+ self.head_dim = head_dim
89
+ self.d_rel = config.d_rel
90
+ self.log_scaling_n_floor = config.log_scaling_n_floor
91
+ self.log_scaling_alpha = config.log_scaling_alpha
92
+ # q/k are per-head RMS-normed (unit norm), so Inkling scales by 1/head_dim.
93
+ self.scaling = 1.0 / head_dim
94
+
95
+ tp_size = get_tensor_model_parallel_world_size()
96
+ self.num_total_heads = num_heads
97
+ self.num_total_kv_heads = num_kv_heads
98
+ assert self.num_total_heads % tp_size == 0
99
+ self.num_heads = self.num_total_heads // tp_size
100
+ if self.num_total_kv_heads >= tp_size:
101
+ assert self.num_total_kv_heads % tp_size == 0
102
+ else:
103
+ assert tp_size % self.num_total_kv_heads == 0
104
+ self.num_kv_heads = max(1, self.num_total_kv_heads // tp_size)
105
+ # When tp_size > num_kv_heads the K/V projections are padded up to
106
+ # tp_size heads so each rank gets at least one (GQA replication).
107
+ kv_total_for_sizing = max(self.num_total_kv_heads, tp_size)
108
+
109
+ self.qkvr = MergedColumnParallelLinear(
110
+ input_size=config.hidden_size,
111
+ output_sizes=[
112
+ head_dim * self.num_total_heads,
113
+ head_dim * kv_total_for_sizing,
114
+ head_dim * kv_total_for_sizing,
115
+ self.d_rel * self.num_total_heads,
116
+ ],
117
+ bias=config.q_bias,
118
+ quant_config=quant_config,
119
+ prefix=f"{prefix}.qkvr",
120
+ )
121
+ self.wo_ud = RowParallelLinear(
122
+ input_size=head_dim * self.num_total_heads,
123
+ output_size=config.hidden_size,
124
+ bias=config.o_bias,
125
+ quant_config=quant_config,
126
+ # reduce_results=False: the partial output is all-reduced below
127
+ # (one-shot custom AR) so the attention-output sconv can run on the
128
+ # full hidden width fused with the residual add + rmsnorm.
129
+ reduce_results=False,
130
+ prefix=f"{prefix}.wo_ud",
131
+ )
132
+ self.rel_extent = local_extent if is_local else rel_extent
133
+ self.local_extent = local_extent if is_local else None
134
+ self.rel_logits_proj = RelLogitsProj(self.d_rel, self.rel_extent)
135
+ self.q_norm = InklingRMSNorm(head_dim, eps=config.rms_norm_eps)
136
+ self.k_norm = InklingRMSNorm(head_dim, eps=config.rms_norm_eps)
137
+
138
+ # Short convolution on the K/V streams (per-head-width, TP sharded),
139
+ # applied after the qkvr projection and before q/k norm.
140
+ kv_conv_dim = self.num_kv_heads * head_dim
141
+ self.conv_owner = conv_owner
142
+ self.k_sconv = InklingShortConv(
143
+ kv_conv_dim, config.sconv_kernel_size, owner=conv_owner, stream_idx=_K
144
+ )
145
+ self.v_sconv = InklingShortConv(
146
+ kv_conv_dim, config.sconv_kernel_size, owner=conv_owner, stream_idx=_V
147
+ )
148
+
149
+ # FA4 left/right window; right=0 keeps it causal. local_extent-1 mirrors
150
+ # the source (sliding_window_size - 1).
151
+ self.window_size: tuple[int, int] = (
152
+ (local_extent - 1, 0) if is_local else (-1, -1)
153
+ )
154
+ # Static per-layer-type KV length bound for the split heuristic: local
155
+ # layers never see more than the sliding window.
156
+ vllm_config = get_current_vllm_config()
157
+ self._max_kv_len = (
158
+ local_extent if is_local else vllm_config.model_config.max_model_len
159
+ )
160
+
161
+ # ---- KV-cache wiring (reuse FlashAttentionBackend for metadata) ----
162
+ cache_config = vllm_config.cache_config
163
+ self.kv_cache_dtype = (
164
+ cache_config.cache_dtype if cache_config is not None else "auto"
165
+ )
166
+ self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
167
+ self.kv_cache_dtype, vllm_config.model_config
168
+ )
169
+ self.register_buffer("k_scale", torch.ones((), dtype=torch.float32))
170
+ self.register_buffer("v_scale", torch.ones((), dtype=torch.float32))
171
+
172
+ compilation_config = vllm_config.compilation_config
173
+ if prefix in compilation_config.static_forward_context:
174
+ raise ValueError(f"Duplicate layer name: {prefix}")
175
+ compilation_config.static_forward_context[prefix] = self
176
+ self.kv_cache = torch.tensor([]) # replaced by bind_kv_cache
177
+
178
+ register_fa4_warmup(
179
+ InklingFA4WarmupConfig(
180
+ num_heads=self.num_heads,
181
+ num_kv_heads=self.num_kv_heads,
182
+ head_dim=self.head_dim,
183
+ rel_extent=self.rel_extent,
184
+ window_size=self.window_size,
185
+ is_local=self.is_local,
186
+ max_kv_len=self._max_kv_len,
187
+ dtype=vllm_config.model_config.dtype,
188
+ kv_dtype=self.kv_cache_torch_dtype,
189
+ block_size=vllm_config.cache_config.block_size,
190
+ max_num_reqs=vllm_config.scheduler_config.max_num_seqs,
191
+ max_num_batched_tokens=(
192
+ vllm_config.scheduler_config.max_num_batched_tokens
193
+ ),
194
+ )
195
+ )
196
+
197
+ def get_attn_backend(self) -> type[AttentionBackend]:
198
+ return FlashAttentionBackend
199
+
200
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
201
+ block_size = vllm_config.cache_config.block_size
202
+ if self.is_local:
203
+ assert self.local_extent is not None
204
+ return SlidingWindowSpec(
205
+ block_size=block_size,
206
+ num_kv_heads=self.num_kv_heads,
207
+ head_size=self.head_dim,
208
+ dtype=self.kv_cache_torch_dtype,
209
+ sliding_window=self.local_extent,
210
+ )
211
+ return FullAttentionSpec(
212
+ block_size=block_size,
213
+ num_kv_heads=self.num_kv_heads,
214
+ head_size=self.head_dim,
215
+ dtype=self.kv_cache_torch_dtype,
216
+ )
217
+
218
+ def _split_kv_cache(self) -> tuple[torch.Tensor, torch.Tensor]:
219
+ key_cache, value_cache = self.kv_cache.transpose(1, 2).split(
220
+ self.head_dim, dim=-1
221
+ )
222
+ return (
223
+ canonicalize_singleton_dim_strides(key_cache),
224
+ canonicalize_singleton_dim_strides(value_cache),
225
+ )
226
+
227
+ def forward(
228
+ self,
229
+ positions: torch.Tensor,
230
+ hidden_states: torch.Tensor,
231
+ log_scaling: torch.Tensor | None = None,
232
+ ) -> torch.Tensor:
233
+ num_tokens = hidden_states.shape[0]
234
+ qkvr, _ = self.qkvr(hidden_states)
235
+
236
+ attn_metadata = get_forward_context().attn_metadata
237
+ attn_output = torch.empty(
238
+ (num_tokens, self.num_heads, self.head_dim),
239
+ dtype=qkvr.dtype,
240
+ device=qkvr.device,
241
+ )
242
+ if not isinstance(attn_metadata, dict):
243
+ attn_output.zero_()
244
+ else:
245
+ conv_meta = attn_metadata[self.conv_owner.prefix]
246
+ md = attn_metadata[self.prefix]
247
+ assert isinstance(conv_meta, InklingSconvMetadata)
248
+ fa_md = cast(FlashAttentionMetadata, md)
249
+ assert self.kv_cache.numel() > 0
250
+ assert self.conv_owner.kv_cache.numel() > 0
251
+ # One launch: K/V sconv (conv-cache insert + conv + residual),
252
+ # Q/K per-head rmsnorm, and the attention KV-cache write. K/V are
253
+ # consumed via the KV cache; only normed q is materialized.
254
+ key_cache, value_cache = self._split_kv_cache()
255
+ off_k, _ = self.conv_owner.stream_ranges[_K]
256
+ off_v, _ = self.conv_owner.stream_ranges[_V]
257
+ q, rel_logits = fused_qkvr_prep(
258
+ qkvr,
259
+ self.k_sconv.weight.squeeze(1),
260
+ self.v_sconv.weight.squeeze(1),
261
+ self.q_norm.weight,
262
+ self.k_norm.weight,
263
+ self.rel_logits_proj.proj,
264
+ self.q_norm.variance_epsilon,
265
+ self.num_heads,
266
+ self.num_kv_heads,
267
+ self.head_dim,
268
+ self.d_rel,
269
+ self.conv_owner.kv_cache,
270
+ key_cache,
271
+ value_cache,
272
+ positions,
273
+ conv_meta.block_table,
274
+ conv_meta.seq_idx,
275
+ conv_meta.slot_mapping,
276
+ conv_meta.query_start,
277
+ fa_md.slot_mapping,
278
+ off_k,
279
+ off_v,
280
+ self.conv_owner.cache_block_size,
281
+ log_scaling if not self.is_local else None,
282
+ )
283
+ q = q.view(num_tokens, self.num_heads, self.head_dim)
284
+ self._attention(q, rel_logits, attn_output)
285
+
286
+ flat = attn_output.view(num_tokens, -1)
287
+ output, _ = self.wo_ud(flat)
288
+ return output
289
+
290
+ @eager_break_during_capture
291
+ def _attention(
292
+ self,
293
+ q: torch.Tensor,
294
+ rel_logits: torch.Tensor,
295
+ output: torch.Tensor,
296
+ ) -> None:
297
+ attn_metadata = get_forward_context().attn_metadata
298
+ assert isinstance(attn_metadata, dict)
299
+ md = cast(FlashAttentionMetadata, attn_metadata[self.prefix])
300
+
301
+ nt = md.num_actual_tokens
302
+ key_cache, value_cache = self._split_kv_cache()
303
+ max_seqlen_q = bucket_max_seqlen_q(md.max_query_len)
304
+ num_splits = inkling_fa4_num_splits(
305
+ is_local=self.is_local,
306
+ batch_size=md.seq_lens.shape[0],
307
+ max_query_len=max_seqlen_q,
308
+ num_heads=self.num_heads,
309
+ num_kv_heads=self.num_kv_heads,
310
+ max_kv_len=md.max_seq_len,
311
+ )
312
+ inkling_fa4_rel_attention(
313
+ q[:nt],
314
+ key_cache,
315
+ value_cache,
316
+ block_table=md.block_table,
317
+ cache_seqlens=md.seq_lens,
318
+ cu_seqlens_q=md.query_start_loc,
319
+ max_seqlen_q=max_seqlen_q,
320
+ softmax_scale=self.scaling,
321
+ causal=True,
322
+ window_size=self.window_size,
323
+ rel_extent=self.rel_extent,
324
+ rel_logits=rel_logits[:nt],
325
+ num_splits=num_splits,
326
+ max_kv_len=md.max_seq_len,
327
+ out=output[:nt],
328
+ )
vllm/models/inkling/amd/layernorm.py
new
+26
@@ -0,0 +1,26 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inkling RMSNorm (no bias, weight-scaled), backed by the vendored Triton kernel."""
4
+
5
+from __future__ import annotations
6
+
7
+import torch
8
+from torch import nn
9
+
10
+from .ops import rmsnorm
11
+
12
+
13
+class InklingRMSNorm(nn.Module):
14
+ def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
15
+ super().__init__()
16
+ self.weight = nn.Parameter(torch.ones(hidden_size))
17
+ self.variance_epsilon = eps
18
+ self.hidden_size = hidden_size
19
+
20
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
21
+ if x.numel() == 0:
22
+ return x
23
+ original_shape = x.shape
24
+ x_2d = x.contiguous().view(-1, self.hidden_size)
25
+ y = rmsnorm(x_2d, self.weight, self.variance_epsilon)
26
+ return y.view(original_shape)
vllm/models/inkling/amd/logits_processor.py
new
+129
@@ -0,0 +1,129 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inkling logits processor (muP + LoRA aware).
4
+
5
+Inkling divides the final logits by a muP width multiplier
6
+(``logits_mup_width_multiplier``). This applies it two ways, depending on
7
+whether an lm_head LoRA is attached:
8
+
9
+* No LoRA: fold ``1/mup`` into the lm_head GEMM alpha (fp32 epilogue) -- no
10
+ separate elementwise kernel, no extra rounding, no weight mutation.
11
+* LoRA attached: the LoRA manager wraps this layer in
12
+ ``LogitsProcessorWithLoRA``, whose ``forward`` calls
13
+ ``type(base_layer).forward(self=wrapper)`` -- so this ``forward`` runs with
14
+ ``self`` bound to the wrapper. We detect that via ``base_layer`` and take the
15
+ LoRA path: run the wrapper's ``_get_logits`` (base logits + the lm_head LoRA
16
+ delta), then divide the full logits by the multiplier so the delta is scaled
17
+ too. muP thus composes with the LoRA delta, with the dispatch as the only
18
+ model-side branching.
19
+"""
20
+
21
+from __future__ import annotations
22
+
23
+import torch
24
+
25
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
26
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
27
+
28
+
29
+class InklingLogitsProcessor(LogitsProcessor):
30
+ """``LogitsProcessor`` that applies Inkling's muP logits width multiplier.
31
+
32
+ Args:
33
+ vocab_size: Padded vocabulary size.
34
+ org_vocab_size: Unpadded vocabulary size (defaults to ``vocab_size``).
35
+ scale: Base logits scale (kept ``1.0`` for the served checkpoint).
36
+ logits_as_input: Whether the input is already logits.
37
+ soft_cap: Optional logit soft cap (``None`` for the served checkpoint).
38
+ logits_mup_width_multiplier: muP width divisor for the final logits;
39
+ ``None`` or ``0`` disables it.
40
+ """
41
+
42
+ def __init__(
43
+ self,
44
+ vocab_size: int,
45
+ org_vocab_size: int | None = None,
46
+ scale: float = 1.0,
47
+ logits_as_input: bool = False,
48
+ soft_cap: float | None = None,
49
+ logits_mup_width_multiplier: float | None = None,
50
+ ) -> None:
51
+ super().__init__(
52
+ vocab_size=vocab_size,
53
+ org_vocab_size=org_vocab_size,
54
+ scale=scale,
55
+ logits_as_input=logits_as_input,
56
+ soft_cap=soft_cap,
57
+ )
58
+ self.logits_mup_width_multiplier = logits_mup_width_multiplier
59
+ self._logits_zero: torch.Tensor | None = None
60
+
61
+ def forward(
62
+ self,
63
+ lm_head: VocabParallelEmbedding,
64
+ hidden_states: torch.Tensor,
65
+ embedding_bias: torch.Tensor | None = None,
66
+ ) -> torch.Tensor | None:
67
+ # ``base_layer`` exists only on the LogitsProcessorWithLoRA wrapper,
68
+ # which calls this forward with ``self`` bound to the wrapper. The
69
+ # wrapper is not an ``InklingLogitsProcessor`` instance, so dispatch
70
+ # ``_lora_forward`` explicitly through the base_layer's class (it
71
+ # provides ``_get_logits``/``logits_as_input``; only ``_lora_forward``
72
+ # lives on this class).
73
+ if hasattr(self, "base_layer"):
74
+ return type(self.base_layer)._lora_forward(
75
+ self, lm_head, hidden_states, embedding_bias
76
+ )
77
+ return self._base_forward(lm_head, hidden_states, embedding_bias)
78
+
79
+ def _lora_forward(
80
+ self,
81
+ lm_head: VocabParallelEmbedding,
82
+ hidden_states: torch.Tensor,
83
+ embedding_bias: torch.Tensor | None = None,
84
+ ) -> torch.Tensor | None:
85
+ # ``self`` is the LogitsProcessorWithLoRA wrapper here: ``_get_logits``
86
+ # returns the base logits plus the lm_head LoRA delta. Apply the muP
87
+ # divisor on the full logits so the LoRA delta is scaled too.
88
+ mup_multiplier = self.base_layer.logits_mup_width_multiplier
89
+ mup = 1.0 / mup_multiplier if mup_multiplier else None
90
+ if self.logits_as_input:
91
+ logits = hidden_states
92
+ else:
93
+ logits = self._get_logits(hidden_states, lm_head, embedding_bias)
94
+ # TODO: fuse this multiplication
95
+ if logits is not None and mup:
96
+ assert self.base_layer.soft_cap is None
97
+ assert self.base_layer.scale == 1.0
98
+ logits = logits * mup
99
+ return logits
100
+
101
+ def _base_forward(
102
+ self,
103
+ lm_head: VocabParallelEmbedding,
104
+ hidden_states: torch.Tensor,
105
+ embedding_bias: torch.Tensor | None = None,
106
+ ) -> torch.Tensor | None:
107
+ mup = self.logits_mup_width_multiplier
108
+ if not mup:
109
+ return super().forward(lm_head, hidden_states, embedding_bias)
110
+ # Fold the muP width divisor into the lm_head GEMM alpha (fp32 epilogue):
111
+ # no separate elementwise kernel, no bf16 rounding of scaled logits, and
112
+ # no weight mutation. Overfit to the served checkpoint: bf16 lm_head, no
113
+ # soft cap, unit logits scale.
114
+ assert self.soft_cap is None
115
+ assert self.scale == 1.0
116
+ w = lm_head.weight
117
+ if self._logits_zero is None:
118
+ self._logits_zero = w.new_zeros(1)
119
+ logits = torch.addmm(
120
+ self._logits_zero,
121
+ hidden_states,
122
+ w.t(),
123
+ beta=0.0,
124
+ alpha=1.0 / mup,
125
+ )
126
+ logits = self._gather_logits(logits)
127
+ if logits is not None:
128
+ logits = logits[..., : self.org_vocab_size]
129
+ return logits
vllm/models/inkling/amd/mlp.py
new
+63
@@ -0,0 +1,63 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inkling dense SwiGLU MLP (also used as the MoE shared expert).
4
+
5
+The checkpoint stores the gate/up projection as a single fused, *interleaved*
6
+weight (``[gate0, up0, gate1, up1, ...]``), so we use a plain
7
+``ColumnParallelLinear`` whose contiguous row-sharding keeps each gate/up pair
8
+together, and an interleaved SwiGLU activation.
9
+"""
10
+
11
+from __future__ import annotations
12
+
13
+import torch
14
+from torch import nn
15
+
16
+from vllm.model_executor.layers.linear import (
17
+ ColumnParallelLinear,
18
+ RowParallelLinear,
19
+)
20
+from vllm.model_executor.layers.quantization import QuantizationConfig
21
+
22
+
23
+class InklingDenseMLP(nn.Module):
24
+ def __init__(
25
+ self,
26
+ hidden_size: int,
27
+ intermediate_size: int,
28
+ *,
29
+ use_global_scale: bool = False,
30
+ quant_config: QuantizationConfig | None = None,
31
+ prefix: str = "",
32
+ ) -> None:
33
+ super().__init__()
34
+ self.gate_up_proj = ColumnParallelLinear(
35
+ hidden_size,
36
+ 2 * intermediate_size,
37
+ bias=False,
38
+ quant_config=quant_config,
39
+ prefix=f"{prefix}.gate_up_proj",
40
+ )
41
+ self.down_proj = RowParallelLinear(
42
+ intermediate_size,
43
+ hidden_size,
44
+ bias=False,
45
+ quant_config=quant_config,
46
+ reduce_results=False,
47
+ prefix=f"{prefix}.down_proj",
48
+ )
49
+ if use_global_scale:
50
+ self.global_scale = nn.Parameter(torch.empty(1), requires_grad=False)
51
+ else:
52
+ self.global_scale = None
53
+
54
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
55
+ from .ops import silu_and_mul_triton
56
+
57
+ gate_up, _ = self.gate_up_proj(x)
58
+ x = silu_and_mul_triton(gate_up)
59
+ x, _ = self.down_proj(x)
60
+ if self.global_scale is not None:
61
+ x.mul_(self.global_scale)
62
+ # TP-partial output: the layer's reduce-scatter fallback consumes it.
63
+ return x
vllm/models/inkling/amd/model.py
new
+669
@@ -0,0 +1,669 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inkling model implementation for AMD GPUs."""
4
+
5
+from __future__ import annotations
6
+
7
+from collections.abc import Iterable
8
+from typing import Any
9
+
10
+import regex as re
11
+import torch
12
+from torch import nn
13
+
14
+from vllm.config import VllmConfig
15
+from vllm.distributed import (
16
+ get_pp_group,
17
+ get_tensor_model_parallel_rank,
18
+ get_tensor_model_parallel_world_size,
19
+ tensor_model_parallel_all_gather,
20
+ tensor_model_parallel_reduce_scatter,
21
+)
22
+from vllm.model_executor.layers.quantization import QuantizationConfig
23
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
24
+from vllm.model_executor.models.interfaces import (
25
+ MultiModalEmbeddings,
26
+ SupportsLoRA,
27
+ SupportsMultiModal,
28
+ SupportsPP,
29
+)
30
+from vllm.model_executor.models.utils import (
31
+ AutoWeightsLoader,
32
+ WeightsMapper,
33
+ make_empty_intermediate_tensors_factory,
34
+ make_layers,
35
+ maybe_prefix,
36
+)
37
+from vllm.models.inkling.common.mm_preprocess import (
38
+ InklingDummyInputsBuilder,
39
+ InklingMultiModalProcessor,
40
+ InklingProcessingInfo,
41
+ inkling_audio_enabled,
42
+ inkling_vision_enabled,
43
+)
44
+from vllm.models.inkling.common.towers import InklingAudio, InklingVision
45
+from vllm.multimodal import MULTIMODAL_REGISTRY
46
+from vllm.sequence import IntermediateTensors
47
+
48
+from ..configs import InklingMMConfig, InklingModelConfig
49
+from .attention import InklingAttention, compute_log_scaling_tau
50
+from .layernorm import InklingRMSNorm
51
+from .logits_processor import InklingLogitsProcessor
52
+from .mlp import InklingDenseMLP
53
+from .moe import InklingMoE
54
+from .ops.norm import add_rmsnorm, embed_rmsnorm
55
+from .sconv_swa_attn import _ATTN, _MLP, InklingConvState
56
+from .short_conv import InklingShortConv
57
+
58
+
59
+def _layer_id(name: str) -> int | None:
60
+ m = re.search(r"\.layers\.(\d+)\.", name)
61
+ return int(m.group(1)) if m else None
62
+
63
+
64
+def _sconv_add_norm(
65
+ delta: torch.Tensor,
66
+ hidden: torch.Tensor,
67
+ sconv: InklingShortConv,
68
+ norm: InklingRMSNorm | None,
69
+ positions: torch.Tensor,
70
+) -> tuple[torch.Tensor | None, torch.Tensor]:
71
+ """``h = hidden + sconv(TP-sum(delta)); y = rmsnorm(h)``.
72
+
73
+ ROCm uses the portable collective path. The Lamport P2P implementation in
74
+ the NVIDIA model relies on CUDA GDC and CUDA-specific symmetric-memory
75
+ publication semantics which cannot be linked into a gfx950 HSACO."""
76
+ norm_w = norm.weight if norm is not None else None
77
+ eps = norm.variance_epsilon if norm is not None else 0.0
78
+
79
+ # RCCL RS -> shard sconv -> AG -> fused add(+rmsnorm).
80
+ shard = tensor_model_parallel_reduce_scatter(delta, dim=-1)
81
+ shard = sconv(shard.contiguous(), positions)
82
+ full = tensor_model_parallel_all_gather(shard, dim=-1)
83
+ if norm is None:
84
+ return None, hidden + full
85
+ return add_rmsnorm(hidden, full, norm_w, eps)
86
+
87
+
88
+class InklingDecoderLayer(nn.Module):
89
+ def __init__(
90
+ self,
91
+ config: InklingModelConfig,
92
+ layer_id: int,
93
+ is_local: bool,
94
+ quant_config: QuantizationConfig | None,
95
+ prefix: str,
96
+ force_dense_mlp: bool = False,
97
+ ) -> None:
98
+ super().__init__()
99
+ # Per-layer owner of the conv state as a paged SWA cache. The 4 sconv
100
+ # streams (K/V/attn/mlp) are packed head-major into one block and share
101
+ # it. Built first so the attention layer can wire its K/V sconv to it.
102
+ self.conv_state = InklingConvState(
103
+ num_kv_heads=(
104
+ config.swa_num_key_value_heads
105
+ if is_local
106
+ else config.num_key_value_heads
107
+ ),
108
+ head_dim=config.swa_head_dim if is_local else config.head_dim,
109
+ hidden_size=config.hidden_size,
110
+ kernel_size=config.sconv_kernel_size,
111
+ prefix=f"{prefix}.conv_state",
112
+ )
113
+ self.attn_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
114
+ self.attn = InklingAttention(
115
+ config,
116
+ num_heads=(
117
+ config.swa_num_attention_heads
118
+ if is_local
119
+ else config.num_attention_heads
120
+ ),
121
+ num_kv_heads=(
122
+ config.swa_num_key_value_heads
123
+ if is_local
124
+ else config.num_key_value_heads
125
+ ),
126
+ head_dim=config.swa_head_dim if is_local else config.head_dim,
127
+ rel_extent=config.rel_extent,
128
+ local_extent=config.sliding_window_size,
129
+ is_local=is_local,
130
+ prefix=f"{prefix}.attn",
131
+ quant_config=quant_config,
132
+ conv_owner=self.conv_state,
133
+ )
134
+ self.mlp_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
135
+ if force_dense_mlp or layer_id < config.dense_mlp_idx:
136
+ self.mlp: nn.Module = InklingDenseMLP(
137
+ hidden_size=config.hidden_size,
138
+ intermediate_size=config.dense_intermediate_size,
139
+ use_global_scale=config.use_global_scale,
140
+ quant_config=quant_config,
141
+ prefix=f"{prefix}.mlp",
142
+ )
143
+ else:
144
+ # InklingMoE decides per layer (from the checkpoint exclude list)
145
+ # whether the routed experts are NVFP4 or bf16; the shared sink
146
+ # experts are always bf16.
147
+ self.mlp = InklingMoE(
148
+ config,
149
+ prefix=f"{prefix}.mlp",
150
+ quant_config=quant_config,
151
+ )
152
+
153
+ # Short convolution on the attention-output and MLP-output residual
154
+ # streams, hidden-sharded: the sublayer outputs are reduce-scattered
155
+ # to [T, H/tp], the sconv runs on the shard, and an all-gather
156
+ # restores the full residual — all fused with the residual add + next
157
+ # rmsnorm via the Lamport P2P kernels for decode-sized batches.
158
+ tp_size = get_tensor_model_parallel_world_size()
159
+ sconv_dim = config.hidden_size // tp_size
160
+ self.attn_sconv = InklingShortConv(
161
+ sconv_dim, config.sconv_kernel_size, owner=self.conv_state, stream_idx=_ATTN
162
+ )
163
+ self.mlp_sconv = InklingShortConv(
164
+ sconv_dim, config.sconv_kernel_size, owner=self.conv_state, stream_idx=_MLP
165
+ )
166
+
167
+ def forward(
168
+ self,
169
+ positions: torch.Tensor,
170
+ hidden_states: torch.Tensor,
171
+ pending: tuple[torch.Tensor | None, InklingShortConv] | None = None,
172
+ defer_mlp_add: bool = False,
173
+ attn_in: torch.Tensor | None = None,
174
+ log_scaling: torch.Tensor | None = None,
175
+ ) -> (
176
+ torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor | None, InklingShortConv]]
177
+ ):
178
+ # The previous sublayer's (pre-reduce, pre-sconv) delta is folded in
179
+ # fused with its RS/sconv/AG and this layer's pre-attention rmsnorm.
180
+ # A None delta means the partials sit in the NVLS symm buffer.
181
+ if pending is None:
182
+ if attn_in is None:
183
+ # First layer; on the text path attn_norm comes fused with
184
+ # the embedding gather (chain_weight in embed_rmsnorm).
185
+ attn_in = self.attn_norm(hidden_states)
186
+ else:
187
+ attn_in, hidden_states = _sconv_add_norm(
188
+ pending[0], hidden_states, pending[1], self.attn_norm, positions
189
+ )
190
+ attn_output = self.attn(positions, attn_in, log_scaling)
191
+ mlp_in, hidden_states = _sconv_add_norm(
192
+ attn_output, hidden_states, self.attn_sconv, self.mlp_norm, positions
193
+ )
194
+ mlp_output = self.mlp(mlp_in)
195
+ if defer_mlp_add:
196
+ # Caller folds mlp_output (pre-reduce, pre-sconv) into the next
197
+ # fused sconv+add+rmsnorm.
198
+ return hidden_states, (mlp_output, self.mlp_sconv)
199
+ # Standalone (MTP) tail: finish the sublayer without a norm.
200
+ return _sconv_add_norm(
201
+ mlp_output, hidden_states, self.mlp_sconv, None, positions
202
+ )[1]
203
+
204
+
205
+class InklingReplicatedEmbedding(nn.Module):
206
+ """Full-vocab embedding table replicated on every TP rank.
207
+
208
+ Trades the full table per rank (~2.3 GiB at V=201k / H=6144 bf16, vs a
209
+ 1/tp shard) for no masked lookup and no per-lookup TP all-reduce — one
210
+ all-reduce per MTP draft step plus one per verify pass — and keeps the
211
+ full table on-rank for the fused gather+norm kernels (``embed_rmsnorm``,
212
+ ``embed_dual_rmsnorm_cat``). Bit-exact vs vocab-parallel: the all-reduce
213
+ there only ever summed one real row against exact zeros. The LM head
214
+ stays vocab-sharded.
215
+ """
216
+
217
+ def __init__(self, num_embeddings: int, embedding_dim: int) -> None:
218
+ super().__init__()
219
+ self.weight = nn.Parameter(
220
+ torch.empty(num_embeddings, embedding_dim, dtype=torch.get_default_dtype()),
221
+ requires_grad=False,
222
+ )
223
+
224
+ def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
225
+ return embed_rmsnorm(input_ids, self.weight, None, 0.0)
226
+
227
+
228
+class InklingModel(nn.Module):
229
+ def __init__(
230
+ self,
231
+ *,
232
+ config: InklingModelConfig,
233
+ quant_config: QuantizationConfig | None,
234
+ prefix: str,
235
+ ) -> None:
236
+ super().__init__()
237
+ self.config = config
238
+ self.embed_tokens = InklingReplicatedEmbedding(
239
+ config.padded_vocab_size, config.hidden_size
240
+ )
241
+ self.embed_norm = (
242
+ InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
243
+ if config.use_embed_norm
244
+ else None
245
+ )
246
+ local_ids = set(config.local_layer_ids)
247
+
248
+ def get_layer(prefix: str) -> InklingDecoderLayer:
249
+ idx = _layer_id(prefix + ".") or int(prefix.split(".")[-1])
250
+ return InklingDecoderLayer(
251
+ config, idx, idx in local_ids, quant_config, prefix
252
+ )
253
+
254
+ self.start_layer, self.end_layer, self.layers = make_layers(
255
+ config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
256
+ )
257
+ self.norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
258
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
259
+ ["hidden_states"], config.hidden_size
260
+ )
261
+
262
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
263
+ # Row gather + embed_norm in one launch.
264
+ norm = self.embed_norm
265
+ return embed_rmsnorm(
266
+ input_ids,
267
+ self.embed_tokens.weight,
268
+ norm.weight if norm is not None else None,
269
+ norm.variance_epsilon if norm is not None else 0.0,
270
+ )
271
+
272
+ def forward(
273
+ self,
274
+ input_ids: torch.Tensor | None,
275
+ positions: torch.Tensor,
276
+ intermediate_tensors: IntermediateTensors | None,
277
+ inputs_embeds: torch.Tensor | None = None,
278
+ ) -> torch.Tensor | IntermediateTensors:
279
+ attn_in0: torch.Tensor | None = None
280
+ if get_pp_group().is_first_rank:
281
+ if inputs_embeds is not None:
282
+ # embed_norm was already applied when producing inputs_embeds.
283
+ hidden_states = inputs_embeds
284
+ else:
285
+ # Gather + embed_norm + the first layer's attn_norm, one launch.
286
+ norm = self.embed_norm
287
+ hidden_states, attn_in0 = embed_rmsnorm(
288
+ input_ids,
289
+ self.embed_tokens.weight,
290
+ norm.weight if norm is not None else None,
291
+ self.config.rms_norm_eps,
292
+ chain_weight=self.layers[self.start_layer].attn_norm.weight,
293
+ )
294
+ else:
295
+ assert intermediate_tensors is not None
296
+ hidden_states = intermediate_tensors["hidden_states"]
297
+ hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
298
+ log_scaling = None
299
+ if self.config.log_scaling_n_floor is not None:
300
+ log_scaling = compute_log_scaling_tau(
301
+ positions,
302
+ self.config.log_scaling_n_floor,
303
+ self.config.log_scaling_alpha,
304
+ )
305
+
306
+ pending: tuple[torch.Tensor | None, InklingShortConv] | None = None
307
+ for layer in self.layers[self.start_layer : self.end_layer]:
308
+ hidden_states, pending = layer(
309
+ positions,
310
+ hidden_states,
311
+ pending=pending,
312
+ defer_mlp_add=True,
313
+ attn_in=attn_in0,
314
+ log_scaling=log_scaling,
315
+ )
316
+ attn_in0 = None
317
+
318
+ if not get_pp_group().is_last_rank:
319
+ if pending is not None:
320
+ hidden_states = _sconv_add_norm(
321
+ pending[0], hidden_states, pending[1], None, positions
322
+ )[1]
323
+ return IntermediateTensors({"hidden_states": hidden_states})
324
+ if pending is not None:
325
+ # Final RS/sconv/AG + residual add fused with the final rmsnorm.
326
+ norm_out = _sconv_add_norm(
327
+ pending[0], hidden_states, pending[1], self.norm, positions
328
+ )[0]
329
+ assert norm_out is not None
330
+ return norm_out
331
+ return self.norm(hidden_states)
332
+
333
+
334
+class _TmlForCausalLMBase(nn.Module, SupportsPP, SupportsLoRA):
335
+ """Shared text-backbone causal-LM scaffolding for both entry classes."""
336
+
337
+ hf_to_vllm_mapper = WeightsMapper(
338
+ orig_to_new_substr={
339
+ ".w13_dn": ".gate_up_proj",
340
+ ".w2_md": ".down_proj",
341
+ },
342
+ orig_to_new_stacked={
343
+ ".attn.wq_du.": (".attn.qkvr.", 0),
344
+ ".attn.wk_dv.": (".attn.qkvr.", 1),
345
+ ".attn.wv_dv.": (".attn.qkvr.", 2),
346
+ ".attn.wr_du.": (".attn.qkvr.", 3),
347
+ },
348
+ orig_to_new_prefix={
349
+ "model.llm.layers.": "model.layers.",
350
+ "model.llm.embed_norm": "model.embed_norm",
351
+ "model.llm.embed": "model.embed_tokens",
352
+ "model.llm.norm": "model.norm",
353
+ "model.llm.unembed": "lm_head",
354
+ "language_model.layers.": "model.layers.",
355
+ "language_model.lm_head.": "lm_head.",
356
+ },
357
+ orig_to_new_suffix={
358
+ # NVFP4 scale
359
+ ".w13_weight.scale": ".w13_weight_scale",
360
+ ".w13_weight.scale2": ".w13_weight_scale_2",
361
+ ".w2_weight.scale": ".w2_weight_scale",
362
+ ".w2_weight.scale2": ".w2_weight_scale_2",
363
+ },
364
+ )
365
+ # Quark uses this mapping when resolving quantization exclusions for the
366
+ # four checkpoint attention projections fused into qkvr.
367
+ packed_modules_mapping = {
368
+ "qkvr": ["wq_du", "wk_dv", "wv_dv", "wr_du"],
369
+ "w13": ["w1", "w3"],
370
+ }
371
+ embedding_modules = {
372
+ "lm_head": "output_embeddings",
373
+ }
374
+
375
+ def _build(
376
+ self,
377
+ vllm_config: VllmConfig,
378
+ text_config: InklingModelConfig,
379
+ prefix: str,
380
+ ) -> None:
381
+ quant_config = vllm_config.quant_config
382
+ self.config = text_config
383
+ # ROCm checkpoints use Quark OCP MXFP4. The global Quark config is
384
+ # passed into each routed MoE and its exclusion list keeps the rest of
385
+ # Inkling in bf16.
386
+ # Read by the MRV2 runner to publish per-request short-conv metadata.
387
+ # Short convolution is intrinsic to Inkling, so this is always set.
388
+ self.uses_sconv = True
389
+ self.model = InklingModel(
390
+ config=text_config,
391
+ quant_config=quant_config,
392
+ prefix=maybe_prefix(prefix, "model"),
393
+ )
394
+ self.lm_head = ParallelLMHead(
395
+ text_config.padded_vocab_size,
396
+ text_config.hidden_size,
397
+ org_num_embeddings=text_config.padded_vocab_size,
398
+ quant_config=quant_config,
399
+ prefix=maybe_prefix(prefix, "lm_head"),
400
+ )
401
+ self.logits_processor = InklingLogitsProcessor(
402
+ text_config.padded_vocab_size,
403
+ org_vocab_size=text_config.vocab_size,
404
+ soft_cap=text_config.final_logit_softcapping,
405
+ logits_mup_width_multiplier=text_config.logits_mup_width_multiplier,
406
+ )
407
+ self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
408
+ self.model.make_empty_intermediate_tensors
409
+ )
410
+
411
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
412
+ return self.model.embed_input_ids(input_ids)
413
+
414
+ def forward(
415
+ self,
416
+ input_ids: torch.Tensor | None,
417
+ positions: torch.Tensor,
418
+ intermediate_tensors: IntermediateTensors | None = None,
419
+ inputs_embeds: torch.Tensor | None = None,
420
+ **kwargs: object,
421
+ ) -> torch.Tensor | IntermediateTensors:
422
+ return self.model(
423
+ input_ids,
424
+ positions,
425
+ intermediate_tensors,
426
+ inputs_embeds,
427
+ )
428
+
429
+ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
430
+ return self.logits_processor(self.lm_head, hidden_states)
431
+
432
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
433
+ return _load_inkling_weights(self, weights, self.config)
434
+
435
+
436
+class InklingForCausalLM(_TmlForCausalLMBase):
437
+ """Text-only entry point (``inkling_model`` checkpoints)."""
438
+
439
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
440
+ super().__init__()
441
+ self._build(vllm_config, vllm_config.model_config.hf_config, prefix)
442
+
443
+
444
+@MULTIMODAL_REGISTRY.register_processor(
445
+ InklingMultiModalProcessor,
446
+ info=InklingProcessingInfo,
447
+ dummy_inputs=InklingDummyInputsBuilder,
448
+)
449
+class InklingForConditionalGeneration(_TmlForCausalLMBase, SupportsMultiModal):
450
+ """Top-level (multimodal) entry point.
451
+
452
+ Builds the vision + audio towers on top of the shared text backbone. Inkling has
453
+ NO cross-modal fusion (the vision tower emits one token per patch, the audio
454
+ tower one token per frame), so generation reuses the inherited backbone
455
+ ``forward`` / ``compute_logits`` (the latter already applies muP) and this
456
+ class only adds multimodal embedding + merge.
457
+ """
458
+
459
+ hf_to_vllm_mapper = _TmlForCausalLMBase.hf_to_vllm_mapper | WeightsMapper(
460
+ orig_to_new_prefix={
461
+ "model.audio.": "audio.",
462
+ "model.visual.": "visual.vision_encoder.",
463
+ },
464
+ )
465
+
466
+ @classmethod
467
+ def get_placeholder_str(cls, modality: str, i: int) -> str | None:
468
+ if modality.startswith("image"):
469
+ return "<|content_image|>"
470
+ if modality.startswith("audio"):
471
+ return "<|content_audio_input|>"
472
+ raise ValueError("Only image or audio modality is supported")
473
+
474
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
475
+ super().__init__()
476
+ config: InklingMMConfig = vllm_config.model_config.hf_config
477
+
478
+ self.visual = (
479
+ InklingVision(config.vision_config, prefix=maybe_prefix(prefix, "visual"))
480
+ if inkling_vision_enabled(config)
481
+ else None
482
+ )
483
+ self.audio = (
484
+ InklingAudio(config.audio_config, prefix=maybe_prefix(prefix, "audio"))
485
+ if inkling_audio_enabled(config)
486
+ else None
487
+ )
488
+
489
+ self._build(vllm_config, config.text_config, prefix)
490
+
491
+ # -- multimodal embedding -------------------------------------------
492
+
493
+ def _process_image_input(
494
+ self, pixel_values: Any, num_patches: Any
495
+ ) -> tuple[torch.Tensor, ...]:
496
+ assert self.visual is not None
497
+ # pixel_values is a list (per item) of [P_i, 2, P, P, 3] tensors,
498
+ # or a single concatenated tensor. Normalize to a flat batch, run the
499
+ # tower once, then split back per item.
500
+ if isinstance(pixel_values, (list, tuple)):
501
+ if not pixel_values:
502
+ return ()
503
+ sizes = [int(p.shape[0]) for p in pixel_values]
504
+ patches = torch.cat(list(pixel_values), dim=0)
505
+ else:
506
+ patches = pixel_values
507
+ sizes = self._sizes_from(num_patches, patches.shape[0])
508
+
509
+ patches = patches.to(device=self.visual.device, dtype=self.visual.dtype)
510
+ embeds = self.visual(patches) # [total_patches, D]
511
+ return tuple(embeds.split(sizes))
512
+
513
+ def _process_audio_input(
514
+ self, input_audio_features: Any, num_audio_tokens: Any
515
+ ) -> tuple[torch.Tensor, ...]:
516
+ assert self.audio is not None
517
+ if isinstance(input_audio_features, (list, tuple)):
518
+ if not input_audio_features:
519
+ return ()
520
+ sizes = [int(d.shape[0]) for d in input_audio_features]
521
+ dmel = torch.cat(list(input_audio_features), dim=0)
522
+ else:
523
+ dmel = input_audio_features
524
+ sizes = self._sizes_from(num_audio_tokens, dmel.shape[0])
525
+
526
+ dmel = dmel.to(device=self.audio.device)
527
+ embeds = self.audio(dmel) # [total_frames, D]
528
+ return tuple(embeds.split(sizes))
529
+
530
+ @staticmethod
531
+ def _sizes_from(counts: Any, total: int) -> list[int]:
532
+ if counts is None:
533
+ return [total]
534
+ if isinstance(counts, torch.Tensor):
535
+ return [int(c) for c in counts.flatten().tolist()]
536
+ if isinstance(counts, (list, tuple)):
537
+ flat: list[int] = []
538
+ for c in counts:
539
+ flat.append(int(c.item()) if isinstance(c, torch.Tensor) else int(c))
540
+ return flat
541
+ return [int(counts)]
542
+
543
+ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
544
+ # Iterate modalities in a stable order so the returned per-item tensors
545
+ # line up with their appearance order; the positional merge in
546
+ # embed_input_ids handles actual placement.
547
+ pixel_values = kwargs.get("pixel_values")
548
+ num_patches = kwargs.get("num_patches")
549
+ input_audio_features = kwargs.get("input_audio_features")
550
+ num_audio_tokens = kwargs.get("num_audio_tokens")
551
+
552
+ embeddings: tuple[torch.Tensor, ...] = ()
553
+ if pixel_values is not None and self.visual is not None:
554
+ embeddings += self._process_image_input(pixel_values, num_patches)
555
+ if input_audio_features is not None and self.audio is not None:
556
+ embeddings += self._process_audio_input(
557
+ input_audio_features, num_audio_tokens
558
+ )
559
+ return embeddings
560
+
561
+ def embed_input_ids(
562
+ self,
563
+ input_ids: torch.Tensor,
564
+ multimodal_embeddings: MultiModalEmbeddings | None = None,
565
+ *,
566
+ is_multimodal: torch.Tensor | None = None,
567
+ ) -> torch.Tensor:
568
+ # Override the base's 1-arg embed_input_ids: the runner calls this 3-arg
569
+ # signature for multimodal models. Text embeddings come from the shared
570
+ # backbone (which applies embed_norm); MM embeddings are scattered in.
571
+ from vllm.model_executor.models.utils import _merge_multimodal_embeddings
572
+
573
+ # Placeholder ids use unused vocabulary slots and these positions are
574
+ # overwritten by MM embeds below.
575
+ inputs_embeds = self.model.embed_input_ids(input_ids)
576
+ if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
577
+ return inputs_embeds
578
+ assert is_multimodal is not None
579
+ return _merge_multimodal_embeddings(
580
+ inputs_embeds=inputs_embeds,
581
+ multimodal_embeddings=multimodal_embeddings,
582
+ is_multimodal=is_multimodal,
583
+ )
584
+
585
+ def get_language_model(self) -> nn.Module:
586
+ # This class IS the causal LM (the towers are side branches), so the
587
+ # language model is self — callers expect a module exposing ``.model``
588
+ # / ``.lm_head`` (e.g. the MTP/eagle loader shares embeddings via
589
+ # ``get_language_model().model.embed_tokens``).
590
+ return self
591
+
592
+
593
+# ===========================================================================
594
+# Weight loading
595
+# ===========================================================================
596
+
597
+
598
+_MOE_EXPERT_WEIGHT_RE = re.compile(
599
+ r"^(?P<mlp>.*\.mlp)\.(?P<rest>(?:shared_)?experts\..+)$"
600
+)
601
+
602
+
603
+def _is_peft_adapter_weight(name: str) -> bool:
604
+ return ".lora_A." in name or ".lora_B." in name
605
+
606
+
607
+def _load_inkling_weights(
608
+ module: nn.Module,
609
+ weights: Iterable[tuple[str, torch.Tensor]],
610
+ config: InklingModelConfig,
611
+) -> set[str]:
612
+ moe_modules = {
613
+ name: mod for name, mod in module.named_modules() if isinstance(mod, InklingMoE)
614
+ }
615
+ loaded: set[str] = set()
616
+ tp_size = get_tensor_model_parallel_world_size()
617
+ tp_rank = get_tensor_model_parallel_rank()
618
+ local_ids = set(config.local_layer_ids)
619
+
620
+ def _iter_loadable_weights() -> Iterable[tuple[str, torch.Tensor]]:
621
+ for name, weight in module.hf_to_vllm_mapper.apply(weights):
622
+ # LightSeek's MXFP4 conversion bundles a quantized copy of the
623
+ # standalone PEFT adapter in the base safetensors index. These are
624
+ # not base-model parameters; adapters remain opt-in at serving.
625
+ if _is_peft_adapter_weight(name):
626
+ continue
627
+ shard_id = getattr(weight, "shard_id", None)
628
+ # Replicate K/V conv-free GQA heads when tp_size > num_kv_heads.
629
+ if (
630
+ shard_id in (1, 2)
631
+ and name.endswith(".attn.qkvr.weight")
632
+ and weight.shape[0] > 0
633
+ ):
634
+ lid = _layer_id(name)
635
+ if lid is not None:
636
+ is_local = lid in local_ids
637
+ n_kv = (
638
+ config.swa_num_key_value_heads
639
+ if is_local
640
+ else config.num_key_value_heads
641
+ )
642
+ head_dim = config.swa_head_dim if is_local else config.head_dim
643
+ if tp_size > n_kv and weight.shape[0] == n_kv * head_dim:
644
+ kv_idx = (tp_rank * n_kv) // tp_size
645
+ weight = weight.narrow(0, kv_idx * head_dim, head_dim)
646
+ weight.shard_id = shard_id
647
+
648
+ # MoE expert tensors (fused stacked, routed + shared sink): translate
649
+ # the checkpoint layout to per-expert FusedMoE loads.
650
+ moe_match = _MOE_EXPERT_WEIGHT_RE.match(name)
651
+ if moe_match is not None and moe_match.group("mlp") in moe_modules:
652
+ moe = moe_modules[moe_match.group("mlp")]
653
+ for rel in moe.load_expert_weight(moe_match.group("rest"), weight):
654
+ loaded.add(f"{moe_match.group('mlp')}.{rel}")
655
+ continue
656
+
657
+ yield name, weight
658
+
659
+ loader = AutoWeightsLoader(module, skip_prefixes=["model.mtp."])
660
+ loaded |= loader.load_weights(_iter_loadable_weights())
661
+
662
+ # Post-load MoE fixups (default input scales, zeroed EP-padding experts).
663
+ for moe_name, moe in moe_modules.items():
664
+ for rel in moe.finalize_load():
665
+ loaded.add(f"{moe_name}.{rel}")
666
+ return loaded
667
+
668
+
669
+EntryClass = [InklingForCausalLM, InklingForConditionalGeneration]
vllm/models/inkling/amd/moe.py
new
+693
@@ -0,0 +1,693 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inkling mixture-of-experts on vLLM's FusedMoE abstraction.
4
+
5
+Overfit to the served checkpoint: sigmoid gate (+ selection bias) top-k over
6
+the routed experts, log-sigmoid renormalization over the k routed + S shared
7
+"sink" logits, scaled by route_scale * global_scale. The routed top-k goes
8
+through vLLM's FusedMoE (which handles TP/EP); the sink experts run in
9
+:class:`InklingSinkExperts` -- replicated across EP ranks (every token
10
+activates every sink) and always bf16 (the checkpoint excludes every
11
+``shared_experts`` from quantization).
12
+
13
+MXFP4 routed experts reuse vLLM's Quark OCP MXFP4 fused-MoE method; excluded
14
+(bf16) layers fall back to the unquantized method. The checkpoint's fused
15
+stacked tensors (interleaved gate/up rows, ``.scale`` / ``.scale2`` /
16
+``.input_amax`` aux tensors) are translated to the standard per-expert loads
17
+in :meth:`InklingMoE.load_expert_weight`.
18
+"""
19
+
20
+from __future__ import annotations
21
+
22
+import math
23
+from typing import TYPE_CHECKING
24
+
25
+import torch
26
+from torch import nn
27
+from torch.nn.parameter import Parameter
28
+
29
+import vllm.envs as envs
30
+from vllm.config import get_current_vllm_config
31
+from vllm.distributed import (
32
+ get_dp_group,
33
+ get_pcp_group,
34
+ get_tensor_model_parallel_rank,
35
+ get_tensor_model_parallel_world_size,
36
+)
37
+from vllm.model_executor.kernels.linear.cute_dsl import ll_bf16
38
+from vllm.model_executor.layers.fused_moe import FusedMoE
39
+from vllm.model_executor.layers.quantization import QuantizationConfig
40
+from vllm.model_executor.utils import set_weight_attrs
41
+from vllm.platforms import current_platform
42
+from vllm.triton_utils import tl, tldevice, triton
43
+from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
44
+from vllm.utils.torch_utils import aux_stream
45
+
46
+from ..configs import InklingModelConfig
47
+
48
+if TYPE_CHECKING:
49
+ from vllm.model_executor.layers.fused_moe.routed_experts import (
50
+ RoutedExperts,
51
+ )
52
+
53
+# ---------------------------------------------------------------------------
54
+# Gate / expert selection
55
+# ---------------------------------------------------------------------------
56
+
57
+_INKLING_LL_BF16_MAX_TOKENS = 64
58
+_MXFP4_INPUT_SCALE_DENOMINATOR = torch.finfo(torch.float8_e4m3fn).max * 6.0
59
+
60
+
61
+def _linear_with_fp32_out(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
62
+ leading = list(x.shape[:-1])
63
+ flat = x.flatten(0, -2)
64
+ if (
65
+ flat.shape[0] <= _INKLING_LL_BF16_MAX_TOKENS
66
+ and flat.dtype == torch.bfloat16
67
+ and weight.dtype == torch.bfloat16
68
+ and flat.is_cuda
69
+ and flat.is_contiguous()
70
+ and weight.is_contiguous()
71
+ and flat.shape[1] % 8 == 0
72
+ and current_platform.has_device_capability(90)
73
+ and ll_bf16.is_available()
74
+ ):
75
+ out = ll_bf16.ll_bf16_gemm(flat, weight)
76
+ else:
77
+ out = torch.mm(flat, weight.T, out_dtype=torch.float32)
78
+ return out.view(*leading, weight.shape[0])
79
+
80
+
81
+@triton.jit(do_not_specialize=["T", "route_scale"])
82
+def _inkling_gate_select_kernel(
83
+ logits_ptr, # [T, G] fp32 gate logits (stride_logits_0 may include pad)
84
+ bias_ptr, # [R] fp32 selection bias (or 0 ptr if HAS_BIAS=False)
85
+ global_scale_ptr, # [1] fp32 (or unused if HAS_GSCALE=False)
86
+ ids_ptr, # [T, K + S] int32 out: selected expert ids
87
+ weights_ptr, # [T, K + S] fp32 out: renormalized weights
88
+ route_scale,
89
+ T,
90
+ G: tl.constexpr, # total gate experts (routed + shared)
91
+ stride_logits_0,
92
+ R: tl.constexpr, # routed experts
93
+ K: tl.constexpr, # top-k routed
94
+ S: tl.constexpr, # shared (sink) experts
95
+ HAS_BIAS: tl.constexpr,
96
+ HAS_GSCALE: tl.constexpr,
97
+ BLOCK_G: tl.constexpr,
98
+):
99
+ pid = tl.program_id(0).to(tl.int64)
100
+ if pid >= T:
101
+ return
102
+ offs = tl.arange(0, BLOCK_G)
103
+ mask_r = offs < R
104
+ logits = tl.load(
105
+ logits_ptr + pid * stride_logits_0 + offs,
106
+ mask=offs < G,
107
+ other=float("-inf"),
108
+ ).to(tl.float32)
109
+
110
+ # Selection scores: sigmoid(routed logits) (+ bias), non-routed lanes -inf.
111
+ sel = tl.where(mask_r, tl.sigmoid(logits), float("-inf"))
112
+ if HAS_BIAS:
113
+ bias = tl.load(bias_ptr + offs, mask=mask_r, other=0.0).to(tl.float32)
114
+ sel = tl.where(mask_r, sel + bias, float("-inf"))
115
+
116
+ scale = route_scale
117
+ if HAS_GSCALE:
118
+ scale = scale * tl.load(global_scale_ptr).to(tl.float32)
119
+
120
+ # Iterative top-K (K is small); argmax tie-breaks to the lowest index
121
+ # (stable ordering).
122
+ A: tl.constexpr = K + S
123
+ offs_a = tl.arange(0, A)
124
+ top_ids = tl.zeros([A], dtype=tl.int32)
125
+ active = tl.zeros([A], dtype=tl.float32)
126
+ for kk in tl.static_range(K):
127
+ idx = tl.argmax(sel, axis=0).to(tl.int32)
128
+ raw = tl.max(tl.where(offs == idx, logits, float("-inf")), axis=0)
129
+ top_ids = tl.where(offs_a == kk, idx, top_ids)
130
+ active = tl.where(offs_a == kk, raw, active)
131
+ sel = tl.where(offs == idx, float("-inf"), sel)
132
+ if S > 0:
133
+ # Shared sink logits sit at the tail of the gate output; their expert
134
+ # ids continue after the routed range (R + j).
135
+ for jj in tl.static_range(S):
136
+ raw = tl.max(tl.where(offs == R + jj, logits, float("-inf")), axis=0)
137
+ top_ids = tl.where(offs_a == K + jj, tl.full([], R + jj, tl.int32), top_ids)
138
+ active = tl.where(offs_a == K + jj, raw, active)
139
+
140
+ # Log-sigmoid renormalization over the K + S active logits.
141
+ abs_l = tl.abs(active)
142
+ min_l = tl.minimum(active, 0.0)
143
+ log_probs = min_l - tldevice.log1p(tldevice.exp(-abs_l))
144
+ max_lp = tl.max(log_probs, axis=0)
145
+ exp_shifted = tldevice.exp(log_probs - max_lp)
146
+ sum_exp = tl.sum(exp_shifted, axis=0)
147
+ weights = exp_shifted / sum_exp * scale
148
+
149
+ tl.store(ids_ptr + pid * A + offs_a, top_ids)
150
+ tl.store(weights_ptr + pid * A + offs_a, weights)
151
+
152
+
153
+def inkling_gate_select(
154
+ logits: torch.Tensor, # [T, >=G] fp32 (rows may carry GEMM padding)
155
+ n_gate_experts: int,
156
+ n_routed_experts: int,
157
+ topk: int,
158
+ n_shared_experts: int,
159
+ bias: torch.Tensor | None,
160
+ route_scale: float,
161
+ global_scale: torch.Tensor | None,
162
+) -> tuple[torch.Tensor, torch.Tensor]:
163
+ """Sigmoid + bias + top-k + log-sigmoid renorm; returns (weights, ids)."""
164
+ assert logits.dtype == torch.float32
165
+ tokens = logits.shape[0]
166
+ active = topk + n_shared_experts
167
+ topk_ids = torch.empty((tokens, active), dtype=torch.int32, device=logits.device)
168
+ topk_weights = torch.empty(
169
+ (tokens, active), dtype=torch.float32, device=logits.device
170
+ )
171
+ if tokens == 0:
172
+ return topk_weights, topk_ids
173
+ _inkling_gate_select_kernel[(tokens,)](
174
+ logits,
175
+ bias if bias is not None else logits,
176
+ global_scale if global_scale is not None else logits,
177
+ topk_ids,
178
+ topk_weights,
179
+ route_scale,
180
+ tokens,
181
+ n_gate_experts,
182
+ logits.stride(0),
183
+ n_routed_experts,
184
+ topk,
185
+ n_shared_experts,
186
+ HAS_BIAS=bias is not None,
187
+ HAS_GSCALE=global_scale is not None,
188
+ BLOCK_G=triton.next_power_of_2(n_gate_experts),
189
+ )
190
+ return topk_weights, topk_ids
191
+
192
+
193
+class InklingGate(nn.Module):
194
+ """Sigmoid gate with selection bias, log-sigmoid renorm after top-k, and
195
+ global scale (the served checkpoint's only configuration)."""
196
+
197
+ def __init__(
198
+ self,
199
+ d_model: int,
200
+ n_routed_experts: int,
201
+ n_shared_experts: int,
202
+ experts_per_token: int,
203
+ route_scale: float,
204
+ *,
205
+ use_global_scale: bool = False,
206
+ use_gate_bias: bool = False,
207
+ ) -> None:
208
+ super().__init__()
209
+ self.n_routed_experts = n_routed_experts
210
+ self.n_shared_experts = n_shared_experts
211
+ self.n_total_experts = n_routed_experts + n_shared_experts
212
+ self.topk = experts_per_token
213
+ self.route_scale = route_scale
214
+
215
+ padded_experts = self.n_total_experts + (-self.n_total_experts) % 8
216
+ self.weight = Parameter(
217
+ torch.empty(padded_experts, d_model), requires_grad=False
218
+ )
219
+ set_weight_attrs(self.weight, {"weight_loader": self._load_weight})
220
+ self.global_scale: Parameter | None
221
+ if use_global_scale:
222
+ self.global_scale = Parameter(
223
+ torch.empty(1, dtype=torch.float32), requires_grad=False
224
+ )
225
+ else:
226
+ self.global_scale = None
227
+ self.bias: Parameter | None
228
+ if use_gate_bias:
229
+ self.bias = Parameter(
230
+ torch.empty(n_routed_experts, dtype=torch.float32),
231
+ requires_grad=False,
232
+ )
233
+ else:
234
+ self.bias = None
235
+
236
+ @staticmethod
237
+ def _load_weight(param: Parameter, loaded_weight: torch.Tensor) -> None:
238
+ param.data.zero_()
239
+ param.data[: loaded_weight.shape[0]].copy_(loaded_weight)
240
+
241
+ def compute_logits(self, x: torch.Tensor) -> torch.Tensor:
242
+ """fp32 gate logits [T, n_total_experts + pad] (pad columns are junk)."""
243
+ return _linear_with_fp32_out(x, self.weight)
244
+
245
+ def select_experts(
246
+ self, gating_output: torch.Tensor
247
+ ) -> tuple[torch.Tensor, torch.Tensor]:
248
+ """Full selection: (weights, ids) of [T, K + S]. The first K entries
249
+ are the routed top-k; the S trailing entries are the sink gammas."""
250
+ return inkling_gate_select(
251
+ gating_output,
252
+ self.n_total_experts,
253
+ self.n_routed_experts,
254
+ self.topk,
255
+ self.n_shared_experts,
256
+ self.bias,
257
+ self.route_scale,
258
+ self.global_scale,
259
+ )
260
+
261
+
262
+# ---------------------------------------------------------------------------
263
+# MoE layer
264
+# ---------------------------------------------------------------------------
265
+
266
+
267
+def _inkling_moe_ep_size() -> int:
268
+ """EP size the FusedMoE layer will run with (mirrors
269
+ FusedMoEParallelConfig.make: experts shard over tp * dp * pcp when
270
+ expert parallelism is enabled)."""
271
+ parallel_config = get_current_vllm_config().parallel_config
272
+ if not parallel_config.enable_expert_parallel:
273
+ return 1
274
+ world = (
275
+ get_tensor_model_parallel_world_size()
276
+ * get_dp_group().world_size
277
+ * get_pcp_group().world_size
278
+ )
279
+ return world if world > 1 else 1
280
+
281
+
282
+class InklingSinkExperts(nn.Module):
283
+ """Shared "sink" experts with per-token gammas, in bf16.
284
+
285
+ Replicated across EP ranks (every token activates every sink, so
286
+ EP-sharding them would hotspot the owning rank) and TP-sharded on the
287
+ intermediate dim so the output remains a TP-partial sum like the routed
288
+ output. The sinks are always bf16 (the checkpoint excludes every
289
+ ``shared_experts`` from quantization): the experts concatenate into two
290
+ plain dense GEMMs with the fused sink epilogue between them.
291
+ """
292
+
293
+ def __init__(
294
+ self, n_experts: int, d_model: int, d_mlp: int, *, prefix: str = ""
295
+ ) -> None:
296
+ super().__init__()
297
+ self.n_experts = n_experts
298
+ tp_size = get_tensor_model_parallel_world_size()
299
+ self.tp_rank = get_tensor_model_parallel_rank()
300
+ intermediate_pp = d_mlp // tp_size
301
+ self.w13_weight = Parameter(
302
+ torch.empty(n_experts, 2 * intermediate_pp, d_model),
303
+ requires_grad=False,
304
+ )
305
+ self.w2_weight = Parameter(
306
+ torch.empty(d_model, n_experts * intermediate_pp),
307
+ requires_grad=False,
308
+ )
309
+ self._unit: torch.Tensor | None = None
310
+
311
+ def load_weight(self, key: str, weight: torch.Tensor) -> list[str]:
312
+ """Load one checkpoint sink tensor (stacked over the S experts)."""
313
+ if key == "w13_weight":
314
+ if weight.shape != self.w13_weight.shape:
315
+ shard = self.w13_weight.shape[1]
316
+ weight = weight.narrow(1, self.tp_rank * shard, shard)
317
+ self.w13_weight.data.copy_(weight)
318
+ return [key]
319
+
320
+ assert key == "w2_weight"
321
+ shard = self.w2_weight.shape[1] // self.n_experts
322
+ shard_start = 0 if weight.shape[2] == shard else self.tp_rank * shard
323
+ for expert_idx, expert_weight in enumerate(weight):
324
+ local_weight = expert_weight.narrow(1, shard_start, shard)
325
+ start = expert_idx * shard
326
+ self.w2_weight.data[:, start : start + shard].copy_(local_weight)
327
+ return [key]
328
+
329
+ def forward(self, x: torch.Tensor, gammas: torch.Tensor) -> torch.Tensor:
330
+ """``sum_e gammas[:, e] * MLP_e(x)`` (TP-partial along d_mlp)."""
331
+ from .ops import sink_silu_mul_epilogue
332
+
333
+ # One GEMM over the experts' stacked w13 (a view), fused epilogue,
334
+ # then one GEMM whose K-reduction over the K-concatenated w2 performs
335
+ # the expert sum.
336
+ if self._unit is None or self._unit.device != x.device:
337
+ self._unit = torch.ones(
338
+ self.n_experts, dtype=torch.float32, device=x.device
339
+ )
340
+ raw = x @ self.w13_weight.view(-1, x.shape[-1]).T # (T, S*2F)
341
+ h = sink_silu_mul_epilogue(
342
+ raw, self._unit, gammas, self._unit, self.n_experts, x.dtype
343
+ )
344
+ return h @ self.w2_weight.T # (T, D)
345
+
346
+
347
+class InklingSinkExpertsLinear(nn.Module):
348
+ """LoRA-capable implementation of the Inkling sink experts."""
349
+
350
+ def __init__(
351
+ self,
352
+ n_experts: int,
353
+ d_model: int,
354
+ d_mlp: int,
355
+ *,
356
+ prefix: str = "",
357
+ ) -> None:
358
+ super().__init__()
359
+ from vllm.model_executor.layers.linear import (
360
+ MergedColumnParallelLinear,
361
+ RowParallelLinear,
362
+ )
363
+
364
+ self.n_experts = n_experts
365
+ self.d_mlp = d_mlp
366
+ total = n_experts * d_mlp
367
+ self.w13 = MergedColumnParallelLinear(
368
+ input_size=d_model,
369
+ output_sizes=[total, total],
370
+ bias=False,
371
+ prefix=f"{prefix}.w13",
372
+ )
373
+ self.w2 = RowParallelLinear(
374
+ input_size=total,
375
+ output_size=d_model,
376
+ bias=False,
377
+ reduce_results=False,
378
+ prefix=f"{prefix}.w2",
379
+ )
380
+ self._w2_input_pp = self.w2.input_size_per_partition
381
+ self._col_expert: torch.Tensor | None = None
382
+
383
+ def _gamma_expand(self, gammas: torch.Tensor) -> torch.Tensor:
384
+ if self._col_expert is None or self._col_expert.device != gammas.device:
385
+ local = self._w2_input_pp
386
+ start = get_tensor_model_parallel_rank() * local
387
+ cols = torch.arange(start, start + local, device=gammas.device)
388
+ self._col_expert = (cols // self.d_mlp).long()
389
+ return gammas[:, self._col_expert]
390
+
391
+ def load_weight(self, key: str, weight: torch.Tensor) -> list[str]:
392
+ if key == "w13_weight":
393
+ d_model = weight.shape[-1]
394
+ gate = weight[:, 0::2, :].reshape(-1, d_model).contiguous()
395
+ up = weight[:, 1::2, :].reshape(-1, d_model).contiguous()
396
+ self.w13.weight_loader(self.w13.weight, gate, 0)
397
+ self.w13.weight_loader(self.w13.weight, up, 1)
398
+ return ["w13.weight"]
399
+ w = weight.permute(1, 0, 2).reshape(weight.shape[1], -1).contiguous()
400
+ self.w2.weight_loader(self.w2.weight, w)
401
+ return ["w2.weight"]
402
+
403
+ def forward(self, x: torch.Tensor, gammas: torch.Tensor) -> torch.Tensor:
404
+ gate_up, _ = self.w13(x)
405
+ gate, up = gate_up.chunk(2, dim=-1)
406
+ hidden_states = torch.nn.functional.silu(gate) * up
407
+ hidden_states = (hidden_states * self._gamma_expand(gammas)).to(x.dtype)
408
+ output, _ = self.w2(hidden_states)
409
+ return output
410
+
411
+
412
+class InklingMoE(nn.Module):
413
+ def __init__(
414
+ self,
415
+ config: InklingModelConfig,
416
+ *,
417
+ prefix: str = "",
418
+ quant_config: QuantizationConfig | None = None,
419
+ ) -> None:
420
+ super().__init__()
421
+ # Overfit to the served checkpoint: sigmoid gate renormalized after
422
+ # top-k, shared sink experts, interleaved gate/up checkpoint rows.
423
+ assert config.gate_activation == "sigmoid" and config.norm_after_topk
424
+ assert config.n_shared_experts > 0 and config.shared_expert_sink
425
+ assert config.inference_moe_w13_interleaved
426
+ n_routed = config.n_routed_experts
427
+ n_shared = config.n_shared_experts
428
+ self.n_routed_experts = n_routed
429
+ self.gate = InklingGate(
430
+ d_model=config.hidden_size,
431
+ n_routed_experts=n_routed,
432
+ n_shared_experts=n_shared,
433
+ experts_per_token=config.num_experts_per_tok,
434
+ route_scale=config.route_scale,
435
+ use_global_scale=config.use_global_scale,
436
+ use_gate_bias=config.use_gate_bias,
437
+ )
438
+
439
+ # TRTLLM MoE kernels assume equal, contiguous per-rank expert slabs
440
+ # (local_expert_offset = ep_rank * local_num_experts), so pad the
441
+ # expert count to a multiple of the EP size. A no-op for the usual
442
+ # power-of-two EP sizes (n_routed is a power of two).
443
+ num_experts = n_routed + (-n_routed) % _inkling_moe_ep_size()
444
+
445
+ # The released MXFP4 checkpoint keeps the first routed-expert layer
446
+ # in bf16 and lists its two expert weights explicitly in Quark's
447
+ # exclusion list. FusedMoE asks for a quant method at the module
448
+ # prefix, so an exact weight exclusion would otherwise be missed and
449
+ # the bf16 tensors would be loaded into MXFP4 parameters.
450
+ routed_quant_config = quant_config
451
+ if quant_config is not None:
452
+ excluded = set(getattr(quant_config, "quant_config", {}).get("exclude", ()))
453
+ expert_prefix = f"{prefix}.experts"
454
+ routed_weights = {
455
+ f"{expert_prefix}.w13_weight",
456
+ f"{expert_prefix}.w2_weight",
457
+ }
458
+ if routed_weights <= excluded:
459
+ routed_quant_config = None
460
+
461
+ self.experts = FusedMoE(
462
+ num_experts=num_experts,
463
+ top_k=config.num_experts_per_tok,
464
+ hidden_size=config.hidden_size,
465
+ intermediate_size=config.intermediate_size,
466
+ renormalize=False,
467
+ quant_config=routed_quant_config,
468
+ prefix=f"{prefix}.experts",
469
+ custom_routing_function=self._select_routed,
470
+ router_logits_dtype=torch.float32,
471
+ activation="silu",
472
+ )
473
+ # The decoder layer reduce-scatters the MoE delta into the sconv
474
+ # stream itself (RS -> shard sconv -> AG); the runner must return the
475
+ # per-rank partial sum instead of all-reducing.
476
+ self.experts.moe_config.skip_final_all_reduce = True
477
+
478
+ self._routed_sel: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None
479
+ sink_experts_cls = (
480
+ InklingSinkExpertsLinear
481
+ if get_current_vllm_config().lora_config is not None
482
+ else InklingSinkExperts
483
+ )
484
+ self.sink_experts = sink_experts_cls(
485
+ n_experts=n_shared,
486
+ d_model=config.hidden_size,
487
+ d_mlp=config.intermediate_size,
488
+ prefix=f"{prefix}.shared_experts",
489
+ )
490
+
491
+ # Sink chain overlaps the routed MoE call on the aux stream for
492
+ # decode-sized batches (same pattern as the runner's SharedExperts
493
+ # multi-stream overlap). The routed GEMM runs on the default stream and
494
+ # the sink chain on the aux stream, joined via these two events by
495
+ # ``maybe_execute_in_parallel``.
496
+ self._sink_stream: torch.cuda.Stream | None = aux_stream()
497
+ self._sink_events = (torch.cuda.Event(), torch.cuda.Event())
498
+
499
+ def _select_routed(
500
+ self,
501
+ hidden_states: torch.Tensor,
502
+ gating_output: torch.Tensor,
503
+ topk: int,
504
+ renormalize: bool,
505
+ ) -> tuple[torch.Tensor, torch.Tensor]:
506
+ """FusedMoE ``custom_routing_function``: the routed top-k slice of the
507
+ full (routed + sink) selection.
508
+
509
+ forward() stashes its selection (keyed by logits identity) so the
510
+ gate select runs once per layer; the fallback covers paths where the
511
+ runner re-derives the logits (e.g. naive DP dispatch).
512
+ """
513
+ del hidden_states, renormalize
514
+ assert topk == self.gate.topk
515
+ cached = self._routed_sel
516
+ self._routed_sel = None
517
+ if cached is not None and cached[0] is gating_output:
518
+ return cached[1], cached[2]
519
+ weights, ids = self.gate.select_experts(gating_output)
520
+ return weights[:, :topk].contiguous(), ids[:, :topk].contiguous()
521
+
522
+ def forward(self, x: torch.Tensor) -> torch.Tensor | None:
523
+ router_logits = self.gate.compute_logits(x)
524
+ num_tokens = x.shape[0]
525
+ # One gate select per layer: the routed slice is stashed for the
526
+ # routing function inside the FusedMoE op; the sink gammas are the
527
+ # trailing columns.
528
+ k = self.gate.topk
529
+ weights, ids = self.gate.select_experts(router_logits)
530
+ self._routed_sel = (
531
+ router_logits,
532
+ weights[:, :k].contiguous(),
533
+ ids[:, :k].contiguous(),
534
+ )
535
+ gammas = weights[:, k:]
536
+
537
+ out, sink_out = maybe_execute_in_parallel(
538
+ lambda: self.experts(hidden_states=x, router_logits=router_logits),
539
+ lambda: self.sink_experts(x, gammas),
540
+ self._sink_events[0],
541
+ self._sink_events[1],
542
+ self._sink_stream
543
+ if num_tokens <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
544
+ else None,
545
+ )
546
+ self._routed_sel = None
547
+
548
+ return out.add_(sink_out)
549
+
550
+ # -- weight loading ----------------------------------------------------
551
+
552
+ def _local_expert_slots(self) -> dict[int, int]:
553
+ """Global expert id -> local slot for this rank's expert partition."""
554
+ manager = self.experts.routed_experts.expert_map_manager
555
+ if manager.expert_map is None:
556
+ return {g: g for g in range(manager.global_num_experts)}
557
+ emap = manager.expert_map.tolist()
558
+ return {g: slot for g, slot in enumerate(emap) if slot >= 0}
559
+
560
+ def load_expert_weight(self, name: str, weight: torch.Tensor) -> list[str]:
561
+ """Load one checkpoint expert tensor.
562
+
563
+ ``name`` is relative to the mlp module: ``experts.<t>`` (routed
564
+ stack) or ``shared_experts.shared_<t>`` (sink experts). Returns the
565
+ loaded param names (relative to this module).
566
+ """
567
+ if name.startswith("shared_experts."):
568
+ key = name.split(".", 1)[1].replace("shared_", "", 1)
569
+ return [
570
+ f"sink_experts.{p}" for p in self.sink_experts.load_weight(key, weight)
571
+ ]
572
+
573
+ experts: RoutedExperts = self.experts.routed_experts
574
+ key = name.split(".", 1)[1]
575
+
576
+ # original_shape is unused by the vLLM serving layout.
577
+ if key.endswith(".original_shape"):
578
+ return []
579
+ if key.endswith(".input_amax"):
580
+ projection = "w13" if key.startswith("w13") else "w2"
581
+ amax = float(weight.max())
582
+ assert math.isfinite(amax) and amax > 0, (
583
+ f"bad {projection} input_amax: {amax}"
584
+ )
585
+ input_scale = getattr(experts, f"{projection}_input_scale")
586
+ input_scale.data.fill_(amax / _MXFP4_INPUT_SCALE_DENOMINATOR)
587
+ return [f"experts.routed_experts.{projection}_input_scale"]
588
+
589
+ # Quark's OCP MXFP4 converter stores block scales as a two-dimensional
590
+ # tensor with the expert and projection-row dimensions flattened:
591
+ # w13: [E * 2I, H / 32], w2: [E * H, I / 32].
592
+ # Restore the expert dimension before applying the EP/TP slicing used
593
+ # for both packed weights and scales.
594
+ if key.endswith("_weight_scale") and weight.ndim == 2:
595
+ if weight.shape[0] % self.n_routed_experts != 0:
596
+ raise ValueError(
597
+ f"cannot unflatten {name} with shape {tuple(weight.shape)} "
598
+ f"over {self.n_routed_experts} experts"
599
+ )
600
+ weight = weight.view(
601
+ self.n_routed_experts,
602
+ weight.shape[0] // self.n_routed_experts,
603
+ weight.shape[1],
604
+ )
605
+
606
+ param = getattr(experts, key)
607
+ slots = self._local_expert_slots()
608
+ gids = sorted(slots)
609
+ lids = [slots[g] for g in gids]
610
+ tp_rank = experts.moe_config.moe_parallel_config.tp_rank
611
+ tp_size = experts.moe_config.moe_parallel_config.tp_size
612
+
613
+ if key.endswith("_scale_2"):
614
+ # Per-expert scalars, vectorized over the local experts. The
615
+ # fused w13 param carries one slot per gate/up half.
616
+ vals = weight[gids].float().to(param.device)
617
+ param.data[lids] = vals[:, None] if param.data.ndim == 2 else vals
618
+ elif key.startswith("w13"):
619
+ # Checkpoint w13 rows are interleaved [g0, u0, g1, u1, ...]; the
620
+ # fused param layout is [w1(gate); w3(up)]. The TP-local rows form
621
+ # one contiguous slab of the interleaved tensor, so upload just
622
+ # that slab (a single bounded synchronous H2D; pre-uploading whole
623
+ # untrimmed tensors pins the mmap pages of the entire checkpoint
624
+ # and OOMs the host) and de-interleave on device.
625
+ dst_half = param.shape[1] // 2
626
+ if weight.shape[1] % (2 * tp_size) != 0:
627
+ raise ValueError(
628
+ f"cannot TP-shard {name} with shape {tuple(weight.shape)} "
629
+ f"over {tp_size} ranks"
630
+ )
631
+ logical_half = weight.shape[1] // (2 * tp_size)
632
+ if logical_half > dst_half:
633
+ raise ValueError(
634
+ f"checkpoint shard for {name} has {logical_half} rows per "
635
+ f"gate/up half, but destination only has {dst_half}"
636
+ )
637
+ for gid, lid in slots.items():
638
+ slab = weight[gid].narrow(
639
+ 0,
640
+ tp_rank * 2 * logical_half,
641
+ 2 * logical_half,
642
+ )
643
+ slab = slab.to(param.device)
644
+ param.data[lid, :logical_half].copy_(slab[0::2])
645
+ param.data[
646
+ lid,
647
+ dst_half : dst_half + logical_half,
648
+ ].copy_(slab[1::2])
649
+ else:
650
+ # w2: shard the packed intermediate (last) dim. AITER rounds the
651
+ # destination intermediate width to 256, so derive the logical TP
652
+ # slice from the checkpoint and leave the destination tail at its
653
+ # initialized padding value (zero for weights, one for scales).
654
+ if weight.shape[2] % tp_size != 0:
655
+ raise ValueError(
656
+ f"cannot TP-shard {name} with shape {tuple(weight.shape)} "
657
+ f"over {tp_size} ranks"
658
+ )
659
+ shard = weight.shape[2] // tp_size
660
+ if shard > param.shape[2]:
661
+ raise ValueError(
662
+ f"checkpoint shard for {name} has width {shard}, but "
663
+ f"destination only has {param.shape[2]}"
664
+ )
665
+ for gid, lid in slots.items():
666
+ param.data[lid, :, :shard].copy_(
667
+ weight[gid].narrow(1, tp_rank * shard, shard)
668
+ )
669
+ return [f"experts.routed_experts.{key}"]
670
+
671
+ def finalize_load(self) -> list[str]:
672
+ """Post-load fixups for zeroed padding experts."""
673
+ experts = self.experts.routed_experts
674
+ out: list[str] = []
675
+ # Zero the EP-alignment padding experts (if any) so their
676
+ # (never-routed) slots hold defined values.
677
+ slots = self._local_expert_slots()
678
+ for gid in range(self.n_routed_experts, experts.global_num_experts):
679
+ lid = slots.get(gid)
680
+ if lid is None:
681
+ continue
682
+ for pname in (
683
+ "w13_weight",
684
+ "w2_weight",
685
+ "w13_weight_scale",
686
+ "w2_weight_scale",
687
+ "w13_weight_scale_2",
688
+ "w2_weight_scale_2",
689
+ ):
690
+ p = getattr(experts, pname, None)
691
+ if p is not None:
692
+ p.data[lid].zero_()
693
+ return out
vllm/models/inkling/amd/mtp.py
new
+410
@@ -0,0 +1,410 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inkling MTP (Multi-Token Prediction) draft model (NVIDIA).
4
+
5
+Implements the first MTP depth from the reference ``mtp_model.py`` shipped with
6
+the checkpoint. It owns ``hidden_norm`` / ``embed_norm`` RMSNorms, a ``2H -> H``
7
+input projection, and a full Inkling transformer block with a dense bf16 MLP.
8
+
9
+The draft shares the target's token embedding table and LM head
10
+(``load_eagle_model`` wires those references) and applies the backbone
11
+``embed_norm`` on top: the depth layers were trained on the same normed
12
+embeddings the backbone consumes (their own ``embed_norm`` weights are
13
+near-identity trims, unlike the backbone's whitening ``embed_norm``).
14
+"""
15
+
16
+from __future__ import annotations
17
+
18
+from collections.abc import Iterable
19
+
20
+import regex as re
21
+import torch
22
+from torch import nn
23
+
24
+from vllm.config import VllmConfig
25
+from vllm.model_executor.layers.linear import ReplicatedLinear
26
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
27
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
28
+from vllm.model_executor.model_loader.mtp_validation import (
29
+ is_mtp_completeness_check_enabled,
30
+)
31
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
32
+from vllm.model_executor.models.utils import maybe_prefix
33
+from vllm.sequence import IntermediateTensors
34
+
35
+from ..configs import InklingModelConfig
36
+from .layernorm import InklingRMSNorm
37
+from .model import InklingDecoderLayer, InklingReplicatedEmbedding
38
+from .ops.norm import embed_dual_rmsnorm_cat, embed_rmsnorm
39
+
40
+# Checkpoint attention projections (wq_du/wk_dv/wv_dv/wr_du) -> fused qkvr.
41
+# Mirrors the backbone's hf_to_vllm_mapper.orig_to_new_stacked; kept as a
42
+# local (pname, wname, shard) list since the MTP loader remaps by hand.
43
+_ATTENTION_PARAMS_MAPPING = [
44
+ ("qkvr", "wq_du", 0),
45
+ ("qkvr", "wk_dv", 1),
46
+ ("qkvr", "wv_dv", 2),
47
+ ("qkvr", "wr_du", 3),
48
+]
49
+
50
+
51
+def _mtp_depth_from_name(name: str) -> int | None:
52
+ m = re.search(r"\.mtp\.layers\.(\d+)\.", name)
53
+ return int(m.group(1)) if m else None
54
+
55
+
56
+class InklingMTPDepthLayer(nn.Module):
57
+ """One MTP depth: norm both inputs, fuse (2H->H), run a Inkling block."""
58
+
59
+ def __init__(self, config: InklingModelConfig, prefix: str, is_local: bool) -> None:
60
+ super().__init__()
61
+ self.hidden_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
62
+ self.embed_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
63
+ self.input_proj = ReplicatedLinear(
64
+ config.hidden_size * 2,
65
+ config.hidden_size,
66
+ bias=False,
67
+ return_bias=False,
68
+ prefix=f"{prefix}.input_proj",
69
+ )
70
+ # A force-dense-MLP bf16 block; ``is_local`` selects sliding-window vs
71
+ # full attention (the swa_* head config and sliding_window window) to
72
+ # match this depth's checkpoint transformer_block weights.
73
+ self.transformer_block = InklingDecoderLayer(
74
+ config,
75
+ layer_id=0,
76
+ is_local=is_local,
77
+ quant_config=None,
78
+ prefix=f"{prefix}.transformer_block",
79
+ force_dense_mlp=True,
80
+ )
81
+
82
+ def forward(self, combined: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
83
+ # ``combined`` is the fused-normed [rmsnorm(hidden) | embed_norm(emb)]
84
+ # input, built by InklingMultiTokenPredictor.fused_input_cat in one launch.
85
+ hidden = self.input_proj(combined)
86
+ # The short conv self-fetches its paged SWA-cache metadata from the
87
+ # forward context (via its conv_owner prefix); no conv_meta to thread.
88
+ return self.transformer_block(positions, hidden)
89
+
90
+
91
+class InklingMultiTokenPredictor(nn.Module):
92
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
93
+ super().__init__()
94
+ assert vllm_config.speculative_config is not None
95
+ config: InklingModelConfig = (
96
+ vllm_config.speculative_config.draft_model_config.hf_config
97
+ )
98
+ self.config = config
99
+ if vllm_config.speculative_config.num_speculative_tokens != 1:
100
+ raise ValueError(
101
+ "Inkling MTP currently supports exactly one speculative token"
102
+ )
103
+ self.chain_hidden_post_norm = config.chain_hidden_post_norm
104
+ local_ids = set(config.local_layer_ids)
105
+ self.layers = nn.ModuleDict(
106
+ {"0": InklingMTPDepthLayer(config, f"{prefix}.layers.0", 0 in local_ids)}
107
+ )
108
+ self.chain_norm = (
109
+ InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
110
+ if self.chain_hidden_post_norm
111
+ else None
112
+ )
113
+ # The target's raw token embedding (pre embed_norm), attached by
114
+ # load_eagle_model. Never materialized here: building our own
115
+ # replicated copy would transiently double the 2.3 GiB table.
116
+ self.embed_tokens: InklingReplicatedEmbedding = None # type: ignore[assignment]
117
+ # The depth layers consume the *backbone-normed* embedding
118
+ # (embed_norm(embed(ids))), not the raw one: mtp embed_norm weights
119
+ # are near-identity (trained on already-normalized inputs), and
120
+ # feeding raw embeddings drops MTP1 acceptance from ~0.85 to ~0.70.
121
+ # Weight loaded from the target's embed_norm.weight; gated like the
122
+ # target's InklingModel.embed_norm.
123
+ self.backbone_embed_norm = (
124
+ InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
125
+ if config.use_embed_norm
126
+ else None
127
+ )
128
+
129
+ def embed_input_ids(
130
+ self,
131
+ input_ids: torch.Tensor,
132
+ multimodal_embeddings: object | None = None,
133
+ *,
134
+ is_multimodal: torch.Tensor | None = None,
135
+ ) -> torch.Tensor:
136
+ """Draft-prefill embedding: fused gather + backbone embed_norm, then
137
+ the target's tower embeddings scattered in unnormed (the backbone
138
+ convention — MM embeds are merged after embed_norm)."""
139
+ norm = self.backbone_embed_norm
140
+ embeds = embed_rmsnorm(
141
+ input_ids,
142
+ self.embed_tokens.weight,
143
+ norm.weight if norm is not None else None,
144
+ norm.variance_epsilon if norm is not None else 0.0,
145
+ )
146
+ if multimodal_embeddings is None or len(multimodal_embeddings) == 0: # type: ignore[arg-type]
147
+ return embeds
148
+ from vllm.model_executor.models.utils import _merge_multimodal_embeddings
149
+
150
+ assert is_multimodal is not None
151
+ return _merge_multimodal_embeddings(
152
+ inputs_embeds=embeds,
153
+ multimodal_embeddings=multimodal_embeddings,
154
+ is_multimodal=is_multimodal,
155
+ )
156
+
157
+ def fused_input_cat(
158
+ self,
159
+ layer: InklingMTPDepthLayer,
160
+ previous_hidden: torch.Tensor,
161
+ input_ids: torch.Tensor,
162
+ inputs_embeds: torch.Tensor | None,
163
+ ) -> torch.Tensor:
164
+ """The depth layer's [rmsnorm(hidden) | embed_norm(embed)] input in one
165
+ launch: embedding row gather + the backbone embed_norm + the depth
166
+ embed_norm chain on one side, hidden_norm on the other, written
167
+ straight into the cat buffer."""
168
+ hidden_w = layer.hidden_norm.weight
169
+ embed_w = layer.embed_norm.weight
170
+ eps = layer.hidden_norm.variance_epsilon
171
+ if inputs_embeds is not None:
172
+ # Draft prefill with target-merged MM embeddings (already
173
+ # backbone-normed via embed_input_ids); only the depth embed_norm
174
+ # remains.
175
+ return embed_dual_rmsnorm_cat(
176
+ previous_hidden, hidden_w, embed_w, eps, embeds=inputs_embeds
177
+ )
178
+ return embed_dual_rmsnorm_cat(
179
+ previous_hidden,
180
+ hidden_w,
181
+ embed_w,
182
+ eps,
183
+ input_ids=input_ids,
184
+ embed_table=self.embed_tokens.weight,
185
+ pre_norm_weight=(
186
+ self.backbone_embed_norm.weight
187
+ if self.backbone_embed_norm is not None
188
+ else None
189
+ ),
190
+ )
191
+
192
+ def forward(
193
+ self,
194
+ input_ids: torch.Tensor,
195
+ positions: torch.Tensor,
196
+ previous_hidden_states: torch.Tensor,
197
+ inputs_embeds: torch.Tensor | None = None,
198
+ spec_step_idx: int = 0,
199
+ ) -> torch.Tensor:
200
+ # The draft's short conv is a paged SWA-cache layer (its conv_owner is
201
+ # auto-enumerated as a draft attention layer); its per-token metadata is
202
+ # built by the speculator's build_attn_metadata and read from the
203
+ # forward context, so nothing extra is threaded here.
204
+ if spec_step_idx != 0:
205
+ raise ValueError("Inkling MTP only supports spec_step_idx=0")
206
+ layer = self.layers["0"]
207
+ combined = self.fused_input_cat(
208
+ layer, previous_hidden_states, input_ids, inputs_embeds
209
+ )
210
+ hidden = layer(combined, positions)
211
+ if self.chain_norm is not None:
212
+ hidden = self.chain_norm(hidden)
213
+ return hidden
214
+
215
+
216
+class InklingMTP(nn.Module):
217
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
218
+ super().__init__()
219
+ assert vllm_config.speculative_config is not None
220
+ config: InklingModelConfig = (
221
+ vllm_config.speculative_config.draft_model_config.hf_config
222
+ )
223
+ self.config = config
224
+ self.model = InklingMultiTokenPredictor(
225
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
226
+ )
227
+ # The target's (vocab-sharded) LM head, attached by load_eagle_model;
228
+ # never materialized here (same reasoning as model.embed_tokens).
229
+ self.lm_head: ParallelLMHead = None # type: ignore[assignment]
230
+ self.logits_processor = LogitsProcessor(
231
+ config.padded_vocab_size,
232
+ org_vocab_size=config.vocab_size,
233
+ soft_cap=config.final_logit_softcapping,
234
+ )
235
+ self._logits_zero: torch.Tensor | None = None
236
+
237
+ def embed_input_ids(
238
+ self,
239
+ input_ids: torch.Tensor,
240
+ multimodal_embeddings: object | None = None,
241
+ *,
242
+ is_multimodal: torch.Tensor | None = None,
243
+ ) -> torch.Tensor:
244
+ return self.model.embed_input_ids(
245
+ input_ids, multimodal_embeddings, is_multimodal=is_multimodal
246
+ )
247
+
248
+ def forward(
249
+ self,
250
+ input_ids: torch.Tensor,
251
+ positions: torch.Tensor,
252
+ hidden_states: torch.Tensor,
253
+ intermediate_tensors: IntermediateTensors | None = None,
254
+ inputs_embeds: torch.Tensor | None = None,
255
+ spec_step_idx: int = 0,
256
+ ) -> torch.Tensor:
257
+ return self.model(
258
+ input_ids,
259
+ positions,
260
+ hidden_states,
261
+ inputs_embeds,
262
+ spec_step_idx,
263
+ )
264
+
265
+ def compute_logits(
266
+ self,
267
+ hidden_states: torch.Tensor,
268
+ spec_step_idx: int = 0,
269
+ ) -> torch.Tensor | None:
270
+ # The MTP shares the base model's LM head, which is trained on
271
+ # ``hidden / mup``-scaled inputs, so apply the same mup scaling here
272
+ # for a matching logit scale — folded into the lm_head GEMM alpha
273
+ # (fp32 epilogue) like the target's compute_logits. (Argmax-invariant
274
+ # for greedy draft sampling, but it matters for the gumbel sampling
275
+ # distribution at temperature > 0.)
276
+ mup = self.config.logits_mup_width_multiplier
277
+ if not mup:
278
+ return self.logits_processor(self.lm_head, hidden_states)
279
+ assert self.logits_processor.soft_cap is None
280
+ assert self.logits_processor.scale == 1.0
281
+ w = self.lm_head.weight
282
+ if self._logits_zero is None:
283
+ self._logits_zero = w.new_zeros(1)
284
+ logits = torch.addmm(
285
+ self._logits_zero,
286
+ hidden_states,
287
+ w.t(),
288
+ beta=0.0,
289
+ alpha=1.0 / mup,
290
+ )
291
+ logits = self.logits_processor._gather_logits(logits)
292
+ if logits is not None:
293
+ logits = logits[..., : self.logits_processor.org_vocab_size]
294
+ return logits
295
+
296
+ def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor:
297
+ """Greedy draft tokens via rank-local argmax + tiny (value, index)
298
+ reduction — no full-vocab logits all-gather. The muP divisor is a
299
+ positive scalar, so the argmax is invariant and the scaling is
300
+ skipped entirely."""
301
+ return self.logits_processor.get_top_tokens(self.lm_head, hidden_states)
302
+
303
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
304
+ return _load_inkling_mtp_weights(self, weights)
305
+
306
+
307
+def _load_inkling_mtp_weights(
308
+ module: InklingMTP,
309
+ weights: Iterable[tuple[str, torch.Tensor]],
310
+) -> set[str]:
311
+ """Load ``model.mtp.*`` weights into the MTP module.
312
+
313
+ Checkpoint keys look like ``model.mtp.chain_norm.weight`` and
314
+ ``model.mtp.layers.{i}.{...}``. The transformer block reuses the backbone
315
+ layer's fused-projection layout, so we apply the same qkvr / gate_up / down
316
+ remapping as ``_load_inkling_weights``. Token embedding and LM head are shared
317
+ (provided by ``load_eagle_model``) and are not present in mtp.safetensors.
318
+ """
319
+ # Per-depth attention is full or sliding-window (config.local_layer_ids);
320
+ # each depth's qkvr MergedColumnParallelLinear is built with the matching
321
+ # (swa_)num_key_value_heads, and its weight_loader handles the TP sharding.
322
+ # The sconv SWA cache pins tp_size <= num_key_value_heads, so tp never
323
+ # exceeds a layer's kv-head count and no GQA K/V replication is needed here.
324
+ params = dict(module.named_parameters())
325
+ loaded: set[str] = set()
326
+
327
+ def _load(name: str, weight: torch.Tensor, shard_id: object = None) -> bool:
328
+ param = params.get(name)
329
+ if param is None:
330
+ return False
331
+ loader = getattr(param, "weight_loader", default_weight_loader)
332
+ if shard_id is None:
333
+ if loader is default_weight_loader or param.shape == weight.shape:
334
+ default_weight_loader(param, weight)
335
+ else:
336
+ loader(param, weight)
337
+ else:
338
+ loader(param, weight, shard_id) # type: ignore[call-arg]
339
+ loaded.add(name)
340
+ return True
341
+
342
+ for name, weight in weights:
343
+ depth = _mtp_depth_from_name(name)
344
+ # Token embedding and LM head are never materialized on the draft
345
+ # (no params to load into); load_eagle_model attaches the target's.
346
+ if name in ("model.llm.embed.weight", "model.llm.unembed.weight"):
347
+ continue
348
+ # The backbone embed_norm, applied to the shared embedding before the
349
+ # depth layers (see InklingMultiTokenPredictor.embed_input_ids). The
350
+ # per-depth mtp.layers.{i}.embed_norm keys carry ".mtp." and are loaded
351
+ # below. Only the shared backbone key routes here.
352
+ if name == "model.llm.embed_norm.weight":
353
+ _load("model.backbone_embed_norm.weight", weight)
354
+ continue
355
+ # Only consume the MTP weights; everything else belongs to the target.
356
+ if ".mtp." not in name:
357
+ continue
358
+ # Only the first checkpoint depth is used for MTP=1.
359
+ if depth is not None and depth != 0:
360
+ continue
361
+ # model.mtp.chain_norm.weight -> model.chain_norm.weight
362
+ # model.mtp.layers.{i}.X -> model.layers.{i}.X
363
+ original_name = name
364
+ name = name.replace(".mtp.layers.", ".layers.").replace(
365
+ ".mtp.chain_norm.", ".chain_norm."
366
+ )
367
+
368
+ if ".chain_norm." in name and module.model.chain_norm is None:
369
+ raise ValueError(
370
+ "Inkling checkpoint contains chain_norm weights but "
371
+ "chain_hidden_post_norm is disabled."
372
+ )
373
+
374
+ # Fused attention qkvr (wq_du/wk_dv/wv_dv/wr_du -> qkvr).
375
+ matched = False
376
+ for pname, wname, shard in _ATTENTION_PARAMS_MAPPING:
377
+ if f".attn.{wname}." in name:
378
+ mapped_name = name.replace(f".{wname}.", f".{pname}.")
379
+ if not _load(mapped_name, weight, shard):
380
+ raise ValueError(f"Unexpected Inkling MTP weight: {original_name}")
381
+ matched = True
382
+ break
383
+ if matched:
384
+ continue
385
+
386
+ # Dense MLP fused gate/up + down.
387
+ if ".mlp.w13_dn.weight" in name:
388
+ loaded_weight = _load(name.replace(".w13_dn.", ".gate_up_proj."), weight)
389
+ elif ".mlp.w2_md.weight" in name:
390
+ loaded_weight = _load(name.replace(".w2_md.", ".down_proj."), weight)
391
+ else:
392
+ if name.endswith(".bias") and name not in params:
393
+ continue
394
+ loaded_weight = _load(name, weight)
395
+ if not loaded_weight:
396
+ raise ValueError(f"Unexpected Inkling MTP weight: {original_name}")
397
+ required = {
398
+ name
399
+ for name in params
400
+ if name.startswith("model.layers.") or name.startswith("model.chain_norm.")
401
+ }
402
+ if (missing := sorted(required - loaded)) and is_mtp_completeness_check_enabled():
403
+ raise ValueError(
404
+ "Inkling MTP checkpoint is missing required parameters: "
405
+ + ", ".join(missing)
406
+ )
407
+ return loaded
408
+
409
+
410
+EntryClass = [InklingMTP]
vllm/models/inkling/amd/ops/__init__.py
new
+44
@@ -0,0 +1,44 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inkling kernels (NVIDIA).
4
+
5
+``rmsnorm`` / ``sconv`` import eagerly. The SwiGLU kernels and the FA4
6
+relative-attention wrapper are exposed lazily to keep this package's
7
+import path lightweight.
8
+"""
9
+
10
+from typing import TYPE_CHECKING
11
+
12
+from .norm import add_rmsnorm, rmsnorm
13
+from .sconv import fused_sconv
14
+
15
+_LAZY_EXPORTS = {
16
+ "silu_and_mul_triton": "silu_and_mul",
17
+ "sink_silu_mul_epilogue": "silu_and_mul",
18
+ "inkling_fa4_rel_attention": "fa4_rel_attention",
19
+ "inkling_rel_attention_split_kv_decode": "rel_attention_decode",
20
+}
21
+
22
+if TYPE_CHECKING:
23
+ from .fa4_rel_attention import inkling_fa4_rel_attention # noqa: F401
24
+ from .silu_and_mul import ( # noqa: F401
25
+ silu_and_mul_triton,
26
+ sink_silu_mul_epilogue,
27
+ )
28
+
29
+__all__ = [
30
+ "add_rmsnorm",
31
+ "rmsnorm",
32
+ "fused_sconv",
33
+ *sorted(_LAZY_EXPORTS),
34
+]
35
+
36
+
37
+def __getattr__(name: str):
38
+ module = _LAZY_EXPORTS.get(name)
39
+ if module is not None:
40
+ import importlib
41
+
42
+ mod = importlib.import_module(f".{module}", __name__)
43
+ return getattr(mod, name)
44
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
vllm/models/inkling/amd/ops/fa4_rel_attention.py
new
+438
@@ -0,0 +1,438 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ROCm paged attention with Inkling's query-dependent relative bias.
4
+
5
+The NVIDIA implementation uses the score-mod hook in tml-fa4. ROCm Flash
6
+Attention and AITER do not expose an equivalent hook, so this module implements
7
+the same operation directly in Triton. Query heads belonging to one KV head
8
+are processed together and KV pages are gathered through vLLM's block table.
9
+"""
10
+
11
+from __future__ import annotations
12
+
13
+import os
14
+from typing import cast
15
+
16
+import torch
17
+
18
+from vllm.models.inkling.amd.ops.rel_attention_decode import (
19
+ inkling_rel_attention_split_kv_decode,
20
+ use_split_kv_decode,
21
+)
22
+from vllm.platforms.rocm import on_gfx950
23
+from vllm.triton_utils import tl, triton
24
+
25
+
26
+def bucket_max_seqlen_q(max_seqlen_q: int) -> int:
27
+ """Round the scheduling bound up to a power of two."""
28
+ return 1 << max(0, max_seqlen_q - 1).bit_length()
29
+
30
+
31
+def use_gfx950_gluon_decode(
32
+ *, max_query_len: int, page_size: int, head_dim: int
33
+) -> bool:
34
+ """Use the vendored TokenSpeed CDNA4 decode kernel where it is supported."""
35
+ return (
36
+ os.getenv("INKLING_GFX950_GLUON", "1") == "1"
37
+ and on_gfx950()
38
+ and max_query_len == 1
39
+ and page_size in (64, 128, 256)
40
+ and head_dim in (64, 128)
41
+ )
42
+
43
+
44
+def use_gfx950_gluon_extend(
45
+ *,
46
+ max_query_len: int,
47
+ max_kv_len: int,
48
+ page_size: int,
49
+ head_dim: int,
50
+ window_left: int,
51
+) -> bool:
52
+ """Use Gluon only for the long full-attention extend regime it wins."""
53
+ return (
54
+ os.getenv("INKLING_GFX950_GLUON", "1") == "1"
55
+ and on_gfx950()
56
+ and max_query_len > 1
57
+ and max_kv_len >= 8192
58
+ and window_left < 0
59
+ and page_size in (64, 128, 256)
60
+ and head_dim in (64, 128)
61
+ )
62
+
63
+
64
+def inkling_fa4_num_splits(
65
+ *,
66
+ is_local: bool,
67
+ batch_size: int,
68
+ max_query_len: int,
69
+ num_heads: int,
70
+ num_kv_heads: int,
71
+ max_kv_len: int,
72
+) -> int:
73
+ """Keep the NVIDIA-facing split heuristic as API-compatible metadata.
74
+
75
+ The ROCm Triton implementation performs online softmax in one program and
76
+ does not consume the result. Keeping this function unchanged avoids
77
+ platform-specific scheduling branches in :mod:`inkling.amd.attention`.
78
+ """
79
+ if is_local:
80
+ return 1
81
+
82
+ q_rows = max_query_len * (num_heads // num_kv_heads)
83
+ q_tiles = (q_rows + 255) // 256
84
+ base_ctas = batch_size * num_kv_heads * q_tiles
85
+ target_ctas = (
86
+ 256 if q_tiles == 1 and batch_size == 1 else (128 if q_tiles == 1 else 64)
87
+ )
88
+ max_splits = 128
89
+ if q_tiles == 1 and batch_size == 1:
90
+ if num_kv_heads == 8:
91
+ max_splits = 16
92
+ elif num_kv_heads == 4 or max_kv_len <= 8192:
93
+ max_splits = 32
94
+ elif max_kv_len <= 65536:
95
+ max_splits = 64
96
+ return max(
97
+ 1,
98
+ min(target_ctas // base_ctas, max_splits, (max_kv_len + 127) // 128),
99
+ )
100
+
101
+
102
+@triton.heuristics(
103
+ {
104
+ "BLOCK_D": lambda args: triton.next_power_of_2(args["head_dim"]),
105
+ "BLOCK_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]),
106
+ "BLOCK_QH": lambda args: (
107
+ args["BLOCK_Q"] * triton.next_power_of_2(args["gqa_group_size"])
108
+ ),
109
+ }
110
+)
111
+@triton.jit(do_not_specialize_on_alignment=["cache_seqlens", "cu_seqlens_q"])
112
+def _inkling_rel_attention_kernel(
113
+ q_ptr, # [total_q, Hq, D]
114
+ k_ptr, # [blocks, page, Hkv, D]
115
+ v_ptr, # [blocks, page, Hkv, D]
116
+ rel_ptr, # [total_q, Hq, rel_extent]
117
+ out_ptr, # [total_q, Hq, D]
118
+ block_table_ptr, # [batch, max_pages]
119
+ cache_seqlens,
120
+ cu_seqlens_q,
121
+ gqa_group_size,
122
+ head_dim,
123
+ softmax_scale,
124
+ stride_q_t,
125
+ stride_q_h,
126
+ stride_q_d,
127
+ stride_k_b,
128
+ stride_k_p,
129
+ stride_k_h,
130
+ stride_k_d,
131
+ stride_v_b,
132
+ stride_v_p,
133
+ stride_v_h,
134
+ stride_v_d,
135
+ stride_r_t,
136
+ stride_r_h,
137
+ stride_r_e,
138
+ stride_o_t,
139
+ stride_o_h,
140
+ stride_o_d,
141
+ stride_bt_b,
142
+ page_size: tl.constexpr,
143
+ rel_extent: tl.constexpr,
144
+ window_left: tl.constexpr,
145
+ BLOCK_Q: tl.constexpr,
146
+ BLOCK_K: tl.constexpr,
147
+ BLOCK_D: tl.constexpr,
148
+ BLOCK_H: tl.constexpr,
149
+ BLOCK_QH: tl.constexpr,
150
+):
151
+ # A program owns BLOCK_Q query positions and every query head sharing one
152
+ # KV head. This makes the QK/PV operations MFMA-friendly on CDNA.
153
+ pid_q = tl.program_id(0)
154
+ pid_kh = tl.program_id(1)
155
+ pid_b = tl.program_id(2)
156
+
157
+ q_start = tl.load(cu_seqlens_q + pid_b)
158
+ q_len = tl.load(cu_seqlens_q + pid_b + 1) - q_start
159
+ q_block = pid_q * BLOCK_Q
160
+ if q_block >= q_len:
161
+ return
162
+
163
+ kv_len = tl.load(cache_seqlens + pid_b)
164
+ prefix_len = kv_len - q_len
165
+ q_head_start = pid_kh * gqa_group_size
166
+ bt_row = block_table_ptr + pid_b * stride_bt_b
167
+
168
+ q_block_ptr = tl.make_block_ptr(
169
+ base=q_ptr + q_start * stride_q_t + q_head_start * stride_q_h,
170
+ shape=(q_len, gqa_group_size, head_dim),
171
+ strides=(stride_q_t, stride_q_h, stride_q_d),
172
+ offsets=(q_block, 0, 0),
173
+ block_shape=(BLOCK_Q, BLOCK_H, BLOCK_D),
174
+ order=(2, 1, 0),
175
+ )
176
+ q = tl.load(q_block_ptr, boundary_check=(0, 1, 2), padding_option="zero")
177
+ q = tl.reshape(q, (BLOCK_QH, BLOCK_D))
178
+
179
+ q_rows = q_block + tl.arange(0, BLOCK_Q)
180
+ q_abs = prefix_len + q_rows
181
+ q_valid = q_rows < q_len
182
+ off_d = tl.arange(0, BLOCK_D)
183
+ d_valid = off_d < head_dim
184
+
185
+ m_i = tl.full((BLOCK_QH,), float("-inf"), dtype=tl.float32)
186
+ l_i = tl.zeros((BLOCK_QH,), dtype=tl.float32)
187
+ acc = tl.zeros((BLOCK_QH, BLOCK_D), dtype=tl.float32)
188
+ log2e: tl.constexpr = 1.4426950408889634
189
+
190
+ for k_start in range(0, kv_len, BLOCK_K):
191
+ k_pos = k_start + tl.arange(0, BLOCK_K)
192
+ k_valid = k_pos < kv_len
193
+ # Masked lanes still need an in-range page-table load.
194
+ safe_pos = tl.minimum(k_pos, tl.maximum(kv_len - 1, 0))
195
+ page = tl.load(bt_row + safe_pos // page_size).to(tl.int64)
196
+ page_off = safe_pos % page_size
197
+
198
+ k = tl.load(
199
+ k_ptr
200
+ + page[None, :] * stride_k_b
201
+ + page_off[None, :] * stride_k_p
202
+ + pid_kh * stride_k_h
203
+ + off_d[:, None] * stride_k_d,
204
+ mask=d_valid[:, None] & k_valid[None, :],
205
+ other=0.0,
206
+ )
207
+ scores = tl.dot(q, k) * (softmax_scale * log2e)
208
+
209
+ dist = q_abs[:, None] - k_pos[None, :]
210
+ score_valid = q_valid[:, None] & k_valid[None, :] & (dist >= 0)
211
+ if window_left >= 0:
212
+ score_valid &= dist <= window_left
213
+
214
+ # rel_logits is query- and head-dependent. Expand [Q, K] distance
215
+ # indices across the GQA heads, then flatten to match QK's row order.
216
+ off_h = tl.arange(0, BLOCK_H)
217
+ rel_dist = dist[:, None, :]
218
+ rel_valid = (
219
+ q_valid[:, None, None]
220
+ & (off_h[None, :, None] < gqa_group_size)
221
+ & (rel_dist >= 0)
222
+ & (rel_dist < rel_extent)
223
+ )
224
+ safe_dist = tl.maximum(0, tl.minimum(rel_dist, rel_extent - 1))
225
+ bias = tl.load(
226
+ rel_ptr
227
+ + (q_start + q_rows[:, None, None]) * stride_r_t
228
+ + (q_head_start + off_h[None, :, None]) * stride_r_h
229
+ + safe_dist * stride_r_e,
230
+ mask=rel_valid,
231
+ other=0.0,
232
+ )
233
+ bias = tl.reshape(bias, (BLOCK_QH, BLOCK_K))
234
+ scores += bias.to(tl.float32) * log2e
235
+ score_valid = tl.reshape(
236
+ score_valid[:, None, :] & (off_h[None, :, None] < gqa_group_size),
237
+ (BLOCK_QH, BLOCK_K),
238
+ )
239
+ scores = tl.where(score_valid, scores, float("-inf"))
240
+
241
+ m_ij = tl.maximum(m_i, tl.max(scores, axis=1))
242
+ has_scores = tl.sum(score_valid.to(tl.int32), axis=1) > 0
243
+ # A sliding-window query can have entire early KV tiles masked. Avoid
244
+ # the -inf - -inf NaN in online softmax until its first live tile.
245
+ alpha = tl.where(has_scores, tl.exp2(m_i - m_ij), 1.0)
246
+ p = tl.where(
247
+ score_valid,
248
+ tl.exp2(scores - m_ij[:, None]),
249
+ 0.0,
250
+ )
251
+ l_i = l_i * alpha + tl.sum(p, axis=1)
252
+ acc *= alpha[:, None]
253
+
254
+ v = tl.load(
255
+ v_ptr
256
+ + page[:, None] * stride_v_b
257
+ + page_off[:, None] * stride_v_p
258
+ + pid_kh * stride_v_h
259
+ + off_d[None, :] * stride_v_d,
260
+ mask=k_valid[:, None] & d_valid[None, :],
261
+ other=0.0,
262
+ )
263
+ acc += tl.dot(p.to(v.dtype), v)
264
+ m_i = m_ij
265
+
266
+ acc /= l_i[:, None]
267
+ acc = tl.reshape(acc, (BLOCK_Q, BLOCK_H, BLOCK_D))
268
+ out_block_ptr = tl.make_block_ptr(
269
+ base=out_ptr + q_start * stride_o_t + q_head_start * stride_o_h,
270
+ shape=(q_len, gqa_group_size, head_dim),
271
+ strides=(stride_o_t, stride_o_h, stride_o_d),
272
+ offsets=(q_block, 0, 0),
273
+ block_shape=(BLOCK_Q, BLOCK_H, BLOCK_D),
274
+ order=(2, 1, 0),
275
+ )
276
+ tl.store(
277
+ out_block_ptr,
278
+ acc.to(out_ptr.dtype.element_ty),
279
+ boundary_check=(0, 1, 2),
280
+ )
281
+
282
+
283
+@torch.no_grad()
284
+def inkling_fa4_rel_attention(
285
+ q: torch.Tensor,
286
+ key_cache: torch.Tensor,
287
+ value_cache: torch.Tensor,
288
+ *,
289
+ block_table: torch.Tensor,
290
+ cache_seqlens: torch.Tensor,
291
+ cu_seqlens_q: torch.Tensor,
292
+ max_seqlen_q: int,
293
+ softmax_scale: float,
294
+ causal: bool,
295
+ window_size: tuple[int, int],
296
+ rel_extent: int,
297
+ rel_logits: torch.Tensor,
298
+ num_splits: int = 32,
299
+ max_kv_len: int | None = None,
300
+ out: torch.Tensor | None = None,
301
+) -> torch.Tensor:
302
+ """Paged varlen attention with Inkling's relative score modification."""
303
+ del num_splits
304
+ if not causal or window_size[1] not in (0, -1):
305
+ raise NotImplementedError("Inkling ROCm attention requires causal masking")
306
+ if q.ndim != 3 or key_cache.ndim != 4 or value_cache.ndim != 4:
307
+ raise ValueError("expected q [T,H,D] and paged K/V [B,P,Hkv,D]")
308
+ if rel_logits.shape != (q.shape[0], q.shape[1], rel_extent):
309
+ raise ValueError(
310
+ f"relative logits have shape {tuple(rel_logits.shape)}, expected "
311
+ f"{(q.shape[0], q.shape[1], rel_extent)}"
312
+ )
313
+
314
+ num_kv_heads = key_cache.shape[2]
315
+ if q.shape[1] % num_kv_heads:
316
+ raise ValueError("query heads must be divisible by KV heads")
317
+ if out is None:
318
+ out = torch.empty_like(q)
319
+ gqa_group_size = q.shape[1] // num_kv_heads
320
+ batch = cu_seqlens_q.shape[0] - 1
321
+ if max_kv_len is None:
322
+ max_kv_len = key_cache.shape[0] * key_cache.shape[1]
323
+ if use_gfx950_gluon_decode(
324
+ max_query_len=max_seqlen_q,
325
+ page_size=key_cache.shape[1],
326
+ head_dim=q.shape[2],
327
+ ):
328
+ # Import only after the architecture guard. The implementation uses
329
+ # gfx950-only CDNA4 Gluon layouts and async-copy operations.
330
+ from vllm.models.inkling.amd.ops.gluon.rel_mha_decode_gfx950 import (
331
+ gluon_rel_mha_decode_gfx950,
332
+ )
333
+
334
+ return gluon_rel_mha_decode_gfx950(
335
+ q,
336
+ key_cache,
337
+ value_cache,
338
+ block_table,
339
+ cache_seqlens,
340
+ max_kv_len,
341
+ rel_logits,
342
+ cu_seqlens_q,
343
+ max_seqlen_q=max_seqlen_q,
344
+ window_left=window_size[0],
345
+ softmax_scale=softmax_scale,
346
+ out=out,
347
+ )
348
+ if use_gfx950_gluon_extend(
349
+ max_query_len=max_seqlen_q,
350
+ max_kv_len=max_kv_len,
351
+ page_size=key_cache.shape[1],
352
+ head_dim=q.shape[2],
353
+ window_left=window_size[0],
354
+ ):
355
+ from vllm.models.inkling.amd.ops.gluon.rel_mha_extend_gfx950 import (
356
+ gluon_rel_mha_extend_gfx950,
357
+ )
358
+
359
+ # The copied kernel retains TokenSpeed's cu_seqlens_kv parameter for
360
+ # API compatibility but does not consume it.
361
+ return cast(
362
+ torch.Tensor,
363
+ gluon_rel_mha_extend_gfx950(
364
+ q,
365
+ cu_seqlens_q,
366
+ cu_seqlens_q,
367
+ key_cache,
368
+ value_cache,
369
+ block_table,
370
+ cache_seqlens,
371
+ window_left=window_size[0],
372
+ max_seqlen_q=max_seqlen_q,
373
+ max_seqlen_k=max_kv_len,
374
+ rel_logits=rel_logits,
375
+ softmax_scale=softmax_scale,
376
+ out=out,
377
+ ),
378
+ )
379
+ if use_split_kv_decode(
380
+ max_query_len=max_seqlen_q,
381
+ max_kv_len=max_kv_len,
382
+ page_size=key_cache.shape[1],
383
+ window_left=window_size[0],
384
+ ):
385
+ return inkling_rel_attention_split_kv_decode(
386
+ q,
387
+ key_cache,
388
+ value_cache,
389
+ block_table=block_table,
390
+ cache_seqlens=cache_seqlens,
391
+ softmax_scale=softmax_scale,
392
+ window_left=window_size[0],
393
+ rel_extent=rel_extent,
394
+ rel_logits=rel_logits,
395
+ max_kv_len=max_kv_len,
396
+ out=out,
397
+ )
398
+ block_q = 1 if max_seqlen_q == 1 else 4
399
+ grid = (triton.cdiv(max_seqlen_q, block_q), num_kv_heads, batch)
400
+ _inkling_rel_attention_kernel[grid](
401
+ q,
402
+ key_cache,
403
+ value_cache,
404
+ rel_logits,
405
+ out,
406
+ block_table,
407
+ cache_seqlens,
408
+ cu_seqlens_q,
409
+ gqa_group_size,
410
+ q.shape[2],
411
+ softmax_scale,
412
+ q.stride(0),
413
+ q.stride(1),
414
+ q.stride(2),
415
+ key_cache.stride(0),
416
+ key_cache.stride(1),
417
+ key_cache.stride(2),
418
+ key_cache.stride(3),
419
+ value_cache.stride(0),
420
+ value_cache.stride(1),
421
+ value_cache.stride(2),
422
+ value_cache.stride(3),
423
+ rel_logits.stride(0),
424
+ rel_logits.stride(1),
425
+ rel_logits.stride(2),
426
+ out.stride(0),
427
+ out.stride(1),
428
+ out.stride(2),
429
+ block_table.stride(0),
430
+ page_size=key_cache.shape[1],
431
+ rel_extent=rel_extent,
432
+ window_left=window_size[0],
433
+ BLOCK_Q=block_q,
434
+ BLOCK_K=64,
435
+ num_warps=4,
436
+ num_stages=1,
437
+ )
438
+ return out
vllm/models/inkling/amd/ops/fa4_warmup.py
new
+35
@@ -0,0 +1,35 @@
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
vllm/models/inkling/amd/ops/gluon/__init__.py
new
+4
@@ -0,0 +1,4 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""GFX950-only Gluon relative-attention kernels for Inkling."""
vllm/models/inkling/amd/ops/gluon/rel_mha_decode_gfx950.py
new
+1037
@@ -0,0 +1,1037 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# Copyright (c) 2026 LightSeek Foundation
5
+#
6
+# Permission is hereby granted, free of charge, to any person obtaining a copy
7
+# of this software and associated documentation files (the "Software"), to deal
8
+# in the Software without restriction, including without limitation the rights
9
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
10
+# copies of the Software, and to permit persons to whom the Software is
11
+# furnished to do so, subject to the following conditions:
12
+#
13
+# The above copyright notice and this permission notice shall be included in
14
+# all copies or substantial portions of the Software.
15
+#
16
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22
+# SOFTWARE.
23
+
24
+"""rel_mha decode Gluon kernel for AMD GFX950."""
25
+
26
+import math
27
+from typing import NamedTuple
28
+
29
+import torch
30
+
31
+from vllm.models.inkling.amd.ops.gluon.utils import (
32
+ _INV_LN2_VALUE,
33
+ InputStrides,
34
+ PagedKVStrides,
35
+ max,
36
+ maximum,
37
+)
38
+from vllm.triton_utils import aggregate, gl, gluon
39
+
40
+cdna4 = gl.amd.cdna4
41
+async_copy = cdna4.async_copy
42
+cdiv = gl.cdiv
43
+_GFX950_SM_COUNT = 256
44
+
45
+
46
+# ===-----------------------------------------------------------------------===#
47
+# Kernel Config
48
+# ===-----------------------------------------------------------------------===#
49
+
50
+
51
+@aggregate
52
+class AttentionConfig:
53
+ SM_SCALE: gl.constexpr
54
+ PAGE_TABLE_STRIDE: gl.constexpr
55
+ PAGE_SIZE: gl.constexpr
56
+ NUM_KV_SPLITS: gl.constexpr
57
+ MAX_SEQLEN_Q: gl.constexpr
58
+ NUM_Q_HEADS: gl.constexpr
59
+ NUM_KV_HEADS: gl.constexpr
60
+ HEAD_DIM: gl.constexpr
61
+ BLOCK_M: gl.constexpr
62
+ BLOCK_N: gl.constexpr
63
+ IS_SLIDING: gl.constexpr
64
+ WINDOW_LEFT: gl.constexpr
65
+ REL_EXTENT: gl.constexpr
66
+ REL_BIAS_QK_SCALE: gl.constexpr
67
+ IS_FP8: gl.constexpr
68
+ GROUP_SIZE: gl.constexpr
69
+ NUM_GROUPS: gl.constexpr
70
+ q_strides: InputStrides
71
+ rel_strides: InputStrides
72
+ k_strides: PagedKVStrides
73
+ v_strides: PagedKVStrides
74
+ qk_layout: gl.constexpr
75
+ pv_layout: gl.constexpr
76
+ q_layout: gl.constexpr
77
+ k_layout: gl.constexpr
78
+ p_layout: gl.constexpr
79
+ v_layout: gl.constexpr
80
+ load_layout: gl.constexpr
81
+ store_layout: gl.constexpr
82
+ reduce_layout: gl.constexpr
83
+ k_smem_layout: gl.constexpr
84
+ v_smem_layout: gl.constexpr
85
+
86
+ @gluon.constexpr_function
87
+ def __init__(
88
+ self,
89
+ SM_SCALE,
90
+ PAGE_TABLE_STRIDE,
91
+ PAGE_SIZE,
92
+ NUM_KV_SPLITS,
93
+ MAX_SEQLEN_Q,
94
+ NUM_Q_HEADS,
95
+ NUM_KV_HEADS,
96
+ HEAD_DIM,
97
+ BLOCK_M,
98
+ BLOCK_N,
99
+ IS_SLIDING,
100
+ WINDOW_LEFT,
101
+ REL_EXTENT,
102
+ REL_BIAS_QK_SCALE,
103
+ IS_FP8,
104
+ q_strides,
105
+ rel_strides,
106
+ k_strides,
107
+ v_strides,
108
+ ):
109
+ assert NUM_Q_HEADS % NUM_KV_HEADS == 0
110
+ assert HEAD_DIM in (64, 128)
111
+ assert BLOCK_N == PAGE_SIZE
112
+ if IS_SLIDING:
113
+ assert WINDOW_LEFT >= 0
114
+ else:
115
+ assert WINDOW_LEFT == -1
116
+
117
+ mfma_layout = gl.amd.AMDMFMALayout(
118
+ version=4,
119
+ instr_shape=[16, 16, 32],
120
+ transposed=True,
121
+ warps_per_cta=[1, 1],
122
+ )
123
+ qk_layout = mfma_layout
124
+ pv_layout = mfma_layout
125
+ # qk_kw is derived from a 128-bit load / dtype bitwidth.
126
+ # pv_kw is empirically tuned.
127
+ qk_kw = 16 if IS_FP8 else 8
128
+ pv_kw = 8 if IS_FP8 else 4
129
+ q_layout = gl.DotOperandLayout(0, qk_layout, k_width=qk_kw)
130
+ k_layout = gl.DotOperandLayout(1, qk_layout, k_width=qk_kw)
131
+ p_layout = gl.DotOperandLayout(0, pv_layout, k_width=pv_kw)
132
+ v_layout = gl.DotOperandLayout(1, pv_layout, k_width=pv_kw)
133
+ # Elements loaded per lane depend on the input dtype, matching qk_kw.
134
+ # load_threads is how many lanes span HEAD_DIM.
135
+ load_vec = 16 if IS_FP8 else 8
136
+ load_threads = HEAD_DIM // load_vec
137
+ load_layout = gl.BlockedLayout(
138
+ [1, load_vec], [64 // load_threads, load_threads], [1, 1], [1, 0]
139
+ )
140
+ # Output is always 16-bit, so a 128-bit store has 8 elements.
141
+ # store_threads is how many lanes span HEAD_DIM.
142
+ store_vec = 8
143
+ store_threads = HEAD_DIM // store_vec
144
+ store_layout = gl.BlockedLayout(
145
+ [1, store_vec], [64 // store_threads, store_threads], [1, 1], [1, 0]
146
+ )
147
+ reduce_layout = gl.BlockedLayout([1, HEAD_DIM // 64], [1, 64], [1, 1], [1, 0])
148
+ # Padding interval is 64 lanes * load_vec elems.
149
+ pad_interval = 64 * load_vec
150
+ # Empirically tuned.
151
+ pad_k = 16 if IS_FP8 else 8
152
+ pad_v = 32
153
+ k_smem_layout = gl.PaddedSharedLayout.with_identity_for(
154
+ [[pad_interval, pad_k]], [BLOCK_N, HEAD_DIM], [1, 0]
155
+ )
156
+ v_smem_layout = gl.PaddedSharedLayout.with_identity_for(
157
+ [[pad_interval, pad_v]], [BLOCK_N, HEAD_DIM], [1, 0]
158
+ )
159
+
160
+ self.SM_SCALE = gl.constexpr(SM_SCALE)
161
+ self.PAGE_TABLE_STRIDE = gl.constexpr(PAGE_TABLE_STRIDE)
162
+ self.PAGE_SIZE = gl.constexpr(PAGE_SIZE)
163
+ self.NUM_KV_SPLITS = gl.constexpr(NUM_KV_SPLITS)
164
+ self.MAX_SEQLEN_Q = gl.constexpr(MAX_SEQLEN_Q)
165
+ self.NUM_Q_HEADS = gl.constexpr(NUM_Q_HEADS)
166
+ self.NUM_KV_HEADS = gl.constexpr(NUM_KV_HEADS)
167
+ self.HEAD_DIM = gl.constexpr(HEAD_DIM)
168
+ self.BLOCK_M = gl.constexpr(BLOCK_M)
169
+ self.BLOCK_N = gl.constexpr(BLOCK_N)
170
+ self.IS_SLIDING = gl.constexpr(IS_SLIDING)
171
+ self.WINDOW_LEFT = gl.constexpr(WINDOW_LEFT)
172
+ self.REL_EXTENT = gl.constexpr(REL_EXTENT)
173
+ self.REL_BIAS_QK_SCALE = gl.constexpr(REL_BIAS_QK_SCALE)
174
+ self.IS_FP8 = gl.constexpr(IS_FP8)
175
+ self.GROUP_SIZE = gl.constexpr(NUM_Q_HEADS // NUM_KV_HEADS)
176
+ self.NUM_GROUPS = gl.constexpr((self.GROUP_SIZE + BLOCK_M - 1) // BLOCK_M)
177
+ self.q_strides = q_strides
178
+ self.rel_strides = rel_strides
179
+ self.k_strides = k_strides
180
+ self.v_strides = v_strides
181
+ self.qk_layout = gl.constexpr(qk_layout)
182
+ self.pv_layout = gl.constexpr(pv_layout)
183
+ self.q_layout = gl.constexpr(q_layout)
184
+ self.k_layout = gl.constexpr(k_layout)
185
+ self.p_layout = gl.constexpr(p_layout)
186
+ self.v_layout = gl.constexpr(v_layout)
187
+ self.load_layout = gl.constexpr(load_layout)
188
+ self.store_layout = gl.constexpr(store_layout)
189
+ self.reduce_layout = gl.constexpr(reduce_layout)
190
+ self.k_smem_layout = gl.constexpr(k_smem_layout)
191
+ self.v_smem_layout = gl.constexpr(v_smem_layout)
192
+
193
+
194
+# ===-----------------------------------------------------------------------===#
195
+# Kernel Program
196
+# ===-----------------------------------------------------------------------===#
197
+
198
+
199
+@aggregate
200
+class AttentionProgram:
201
+ cfg: gl.constexpr
202
+ q_ptr: gl.tensor
203
+ rel_logits_ptr: gl.tensor
204
+ k_cache_ptr: gl.tensor
205
+ v_cache_ptr: gl.tensor
206
+ page_table_ptr: gl.tensor
207
+ cache_seqlens_ptr: gl.tensor
208
+ mid_o_ptr: gl.tensor
209
+ mid_lse_ptr: gl.tensor
210
+ q_index: gl.tensor
211
+ batch: gl.tensor
212
+ kv_head: gl.tensor
213
+ group_start: gl.tensor
214
+ split_id: gl.tensor
215
+ cache_len: gl.tensor
216
+ kv_start: gl.tensor
217
+ split_start: gl.tensor
218
+ split_end: gl.tensor
219
+
220
+ @gluon.constexpr_function
221
+ def __init__(
222
+ self,
223
+ cfg,
224
+ q_ptr,
225
+ rel_logits_ptr,
226
+ k_cache_ptr,
227
+ v_cache_ptr,
228
+ page_table_ptr,
229
+ cache_seqlens_ptr,
230
+ mid_o_ptr,
231
+ mid_lse_ptr,
232
+ q_index,
233
+ batch,
234
+ kv_head,
235
+ group_start,
236
+ split_id,
237
+ cache_len,
238
+ kv_start,
239
+ split_start,
240
+ split_end,
241
+ ):
242
+ self.cfg = gl.constexpr(cfg)
243
+ self.q_ptr = q_ptr
244
+ self.rel_logits_ptr = rel_logits_ptr
245
+ self.k_cache_ptr = k_cache_ptr
246
+ self.v_cache_ptr = v_cache_ptr
247
+ self.page_table_ptr = page_table_ptr
248
+ self.cache_seqlens_ptr = cache_seqlens_ptr
249
+ self.mid_o_ptr = mid_o_ptr
250
+ self.mid_lse_ptr = mid_lse_ptr
251
+ self.q_index = q_index
252
+ self.batch = batch
253
+ self.kv_head = kv_head
254
+ self.group_start = group_start
255
+ self.split_id = split_id
256
+ self.cache_len = cache_len
257
+ self.kv_start = kv_start
258
+ self.split_start = split_start
259
+ self.split_end = split_end
260
+
261
+ @gluon.jit
262
+ def create(
263
+ cfg,
264
+ q_ptr,
265
+ rel_logits_ptr,
266
+ k_cache_ptr,
267
+ v_cache_ptr,
268
+ page_table_ptr,
269
+ cache_seqlens_ptr,
270
+ mid_o_ptr,
271
+ mid_lse_ptr,
272
+ ):
273
+ # Gluon treats ``cfg`` as the constexpr factory argument, while mypy
274
+ # models the first parameter of this aggregate method as ``self``.
275
+ q_index = gl.program_id(0)
276
+ batch = q_index // cfg.MAX_SEQLEN_Q # type: ignore[attr-defined]
277
+ q_pos = q_index - batch * cfg.MAX_SEQLEN_Q # type: ignore[attr-defined]
278
+ head_block = gl.program_id(1)
279
+ kv_head = head_block // cfg.NUM_GROUPS # type: ignore[attr-defined]
280
+ group_block = head_block - kv_head * cfg.NUM_GROUPS # type: ignore[attr-defined]
281
+ group_start = group_block * cfg.BLOCK_M # type: ignore[attr-defined]
282
+ split_id = gl.program_id(2)
283
+ cache_len = gl.load(cache_seqlens_ptr + batch)
284
+ cache_len = cache_len - (
285
+ cfg.MAX_SEQLEN_Q - 1 - q_pos # type: ignore[attr-defined]
286
+ )
287
+ cache_len = maximum(cache_len, 0)
288
+ if cfg.IS_SLIDING: # type: ignore[attr-defined]
289
+ # WINDOW_LEFT is defined as exclusive (keys strictly to the left, not
290
+ # counting the current token), so the window is WINDOW_LEFT + 1 keys
291
+ # once the current token is included. E.g. with
292
+ # WINDOW_LEFT = 127 and cache_len = 500 (current token at index 499),
293
+ # kv_start = 500 - (127 + 1) = 372 keeps keys [372, 499] = 128 keys.
294
+ # Without the + 1, kv_start = 373 would drop the leftmost key.
295
+ kv_start = cache_len - min(
296
+ cache_len,
297
+ cfg.WINDOW_LEFT + 1, # type: ignore[attr-defined]
298
+ )
299
+ else:
300
+ kv_start = cache_len - cache_len
301
+ first_page = kv_start // cfg.PAGE_SIZE # type: ignore[attr-defined]
302
+ end_page = cdiv(cache_len, cfg.PAGE_SIZE) # type: ignore[attr-defined]
303
+ num_pages = end_page - first_page
304
+ pages_per_split = cdiv(
305
+ num_pages,
306
+ cfg.NUM_KV_SPLITS, # type: ignore[attr-defined]
307
+ )
308
+ split_start_page = first_page + split_id * pages_per_split
309
+ split_end_page = min(split_start_page + pages_per_split, end_page)
310
+ split_start = (
311
+ split_start_page * cfg.PAGE_SIZE # type: ignore[attr-defined]
312
+ )
313
+ split_end = min(
314
+ split_end_page * cfg.PAGE_SIZE, # type: ignore[attr-defined]
315
+ cache_len,
316
+ )
317
+ return AttentionProgram(
318
+ gl.constexpr(cfg),
319
+ q_ptr,
320
+ rel_logits_ptr,
321
+ k_cache_ptr,
322
+ v_cache_ptr,
323
+ page_table_ptr,
324
+ cache_seqlens_ptr,
325
+ mid_o_ptr,
326
+ mid_lse_ptr,
327
+ q_index,
328
+ batch,
329
+ kv_head,
330
+ group_start,
331
+ split_id,
332
+ cache_len,
333
+ kv_start,
334
+ split_start,
335
+ split_end,
336
+ )
337
+
338
+ @gluon.jit
339
+ def load_q(self):
340
+ cfg = self.cfg
341
+ offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.q_layout))
342
+ offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.q_layout))
343
+ q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m
344
+ valid = (self.group_start + offs_m) < cfg.GROUP_SIZE
345
+ offsets = cfg.q_strides.offsets(self.q_index, q_heads[:, None], offs_d[None, :])
346
+ return cdna4.buffer_load(self.q_ptr, offsets, mask=valid[:, None], other=0.0)
347
+
348
+ @gluon.jit
349
+ def init_state(self):
350
+ cfg = self.cfg
351
+ m_i = gl.full(
352
+ [cfg.BLOCK_M],
353
+ value=-float("inf"),
354
+ dtype=gl.float32,
355
+ layout=gl.SliceLayout(1, cfg.pv_layout),
356
+ )
357
+ l_i = gl.full(
358
+ [cfg.BLOCK_M],
359
+ value=0.0,
360
+ dtype=gl.float32,
361
+ layout=gl.SliceLayout(1, cfg.pv_layout),
362
+ )
363
+ acc = gl.zeros(
364
+ [cfg.BLOCK_M, cfg.HEAD_DIM], dtype=gl.float32, layout=cfg.pv_layout
365
+ )
366
+ return m_i, l_i, acc
367
+
368
+ @gluon.jit
369
+ def load_page(self, start_n):
370
+ cfg = self.cfg
371
+ page_index = start_n // cfg.PAGE_SIZE
372
+ valid = start_n < self.split_end
373
+ return gl.load(
374
+ self.page_table_ptr + self.batch * cfg.PAGE_TABLE_STRIDE + page_index,
375
+ mask=valid,
376
+ other=0,
377
+ )
378
+
379
+ @gluon.jit
380
+ def issue_load_k(self, physical_page, k_smem):
381
+ cfg = self.cfg
382
+ offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout))
383
+ offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout))
384
+ offsets = cfg.k_strides.offsets(
385
+ physical_page,
386
+ offs_n[:, None],
387
+ self.kv_head,
388
+ offs_d[None, :],
389
+ )
390
+ # can't use buffer_load: paged KV offsets may exceed its 32-bit range.
391
+ async_copy.global_load_to_shared(k_smem, self.k_cache_ptr + offsets)
392
+ async_copy.commit_group()
393
+
394
+ @gluon.jit
395
+ def issue_load_v(self, physical_page, v_smem):
396
+ cfg = self.cfg
397
+ offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout))
398
+ offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout))
399
+ offsets = cfg.v_strides.offsets(
400
+ physical_page,
401
+ offs_n[:, None],
402
+ self.kv_head,
403
+ offs_d[None, :],
404
+ )
405
+ # can't use buffer_load: paged KV offsets may exceed its 32-bit range.
406
+ async_copy.global_load_to_shared(v_smem, self.v_cache_ptr + offsets)
407
+ async_copy.commit_group()
408
+
409
+ @gluon.jit
410
+ def shared_load_k(self, k_smem):
411
+ return k_smem.permute([1, 0]).load(self.cfg.k_layout)
412
+
413
+ @gluon.jit
414
+ def shared_load_v(self, v_smem):
415
+ return v_smem.load(self.cfg.v_layout)
416
+
417
+ @gluon.jit
418
+ def compute_qk(self, q, k):
419
+ cfg = self.cfg
420
+ qk = gl.zeros(
421
+ [cfg.BLOCK_M, cfg.BLOCK_N], dtype=gl.float32, layout=cfg.qk_layout
422
+ )
423
+ return cdna4.mfma(q, k, qk)
424
+
425
+ @gluon.jit
426
+ def apply_rel_bias(self, qk, start_n):
427
+ cfg = self.cfg
428
+ offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.qk_layout))
429
+ offs_n = start_n + gl.arange(
430
+ 0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout)
431
+ )
432
+ q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m
433
+ rel_dist = (self.cache_len - 1) - offs_n
434
+ rel_valid = (rel_dist >= 0) & (rel_dist < cfg.REL_EXTENT)
435
+ rel_idx = gl.where(rel_dist >= 0, rel_dist, 0)
436
+ rel_idx = gl.where(rel_idx < cfg.REL_EXTENT, rel_idx, cfg.REL_EXTENT - 1)
437
+ offsets = cfg.rel_strides.offsets(
438
+ self.q_index, q_heads[:, None], rel_idx[None, :]
439
+ )
440
+ head_valid = (self.group_start + offs_m) < cfg.GROUP_SIZE
441
+ rel_bias = cdna4.buffer_load(
442
+ self.rel_logits_ptr,
443
+ offsets,
444
+ mask=head_valid[:, None] & rel_valid[None, :],
445
+ other=0.0,
446
+ ).to(gl.float32)
447
+ return qk + rel_bias * cfg.REL_BIAS_QK_SCALE
448
+
449
+ @gluon.jit
450
+ def apply_kv_mask(self, qk, start_n):
451
+ cfg = self.cfg
452
+ offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout))
453
+ tokens = start_n + offs_n[None, :]
454
+ mask = (tokens >= self.kv_start) & (tokens < self.split_end)
455
+ return gl.where(mask, qk, -float("inf"))
456
+
457
+ @gluon.jit
458
+ def softmax(self, qk, m_i, l_i, acc):
459
+ cfg = self.cfg
460
+ row_max = max(qk, axis=1)
461
+ row_max = gl.convert_layout(row_max, gl.SliceLayout(1, cfg.pv_layout))
462
+ m_new = maximum(m_i, row_max)
463
+ m_new_scaled = m_new * cfg.SM_SCALE
464
+ qk_shifted = qk * cfg.SM_SCALE - m_new_scaled[:, None]
465
+ p = gl.exp2(qk_shifted)
466
+ m_diff = m_i * cfg.SM_SCALE - m_new_scaled
467
+ alpha = gl.exp2(m_diff)
468
+ l_ij = gl.sum(p, axis=1)
469
+ l_i = l_i * alpha + l_ij
470
+ acc = acc * alpha[:, None]
471
+ p = p.to(self.q_ptr.dtype.element_ty)
472
+ p = gl.convert_layout(p, cfg.p_layout)
473
+ return p, m_new, l_i, acc
474
+
475
+ @gluon.jit
476
+ def compute_pv(self, p, v, acc):
477
+ return cdna4.mfma(p, v, acc)
478
+
479
+ @gluon.jit
480
+ def store_split(self, acc, l_i, m_i):
481
+ cfg = self.cfg
482
+ offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.store_layout))
483
+ offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.store_layout))
484
+ q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m
485
+ valid = ((self.group_start + offs_m) < cfg.GROUP_SIZE) & (
486
+ self.split_start < self.split_end
487
+ )
488
+ acc = gl.convert_layout(acc, cfg.store_layout)
489
+ l_i = gl.convert_layout(l_i, gl.SliceLayout(1, cfg.store_layout))
490
+ m_i = gl.convert_layout(m_i, gl.SliceLayout(1, cfg.store_layout))
491
+ recip_l_i = 1.0 / l_i
492
+ part_o = acc * recip_l_i[:, None]
493
+ part_lse = m_i * cfg.SM_SCALE + gl.log2(l_i)
494
+ mid_o_offsets = (
495
+ (self.q_index * cfg.NUM_Q_HEADS + q_heads[:, None]) * cfg.NUM_KV_SPLITS
496
+ + self.split_id
497
+ ) * cfg.HEAD_DIM + offs_d[None, :]
498
+ mid_lse_offsets = (
499
+ self.q_index * cfg.NUM_Q_HEADS + q_heads
500
+ ) * cfg.NUM_KV_SPLITS + self.split_id
501
+ cdna4.buffer_store(part_o, self.mid_o_ptr, mid_o_offsets, mask=valid[:, None])
502
+ cdna4.buffer_store(part_lse, self.mid_lse_ptr, mid_lse_offsets, mask=valid)
503
+
504
+ @gluon.jit
505
+ def store_output(self, acc, l_i):
506
+ cfg = self.cfg
507
+ offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.store_layout))
508
+ offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.store_layout))
509
+ q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m
510
+ valid = (self.group_start + offs_m) < cfg.GROUP_SIZE
511
+ acc = gl.convert_layout(acc, cfg.store_layout)
512
+ l_i = gl.convert_layout(l_i, gl.SliceLayout(1, cfg.store_layout))
513
+ output = acc * (1.0 / l_i)[:, None]
514
+ output = output.to(self.mid_o_ptr.dtype.element_ty)
515
+ offsets = (self.q_index * cfg.NUM_Q_HEADS + q_heads[:, None]) * cfg.HEAD_DIM
516
+ offsets += offs_d[None, :]
517
+ cdna4.buffer_store(output, self.mid_o_ptr, offsets, mask=valid[:, None])
518
+
519
+
520
+# ===-----------------------------------------------------------------------===#
521
+# Entry Point
522
+# ===-----------------------------------------------------------------------===#
523
+
524
+
525
+@gluon.jit
526
+def _rel_mha_decode_fp16(
527
+ q_ptr,
528
+ rel_logits_ptr,
529
+ k_cache_ptr,
530
+ v_cache_ptr,
531
+ page_table_ptr,
532
+ cache_seqlens_ptr,
533
+ mid_o_ptr,
534
+ mid_lse_ptr,
535
+ Q_STRIDE_B: gl.constexpr,
536
+ Q_STRIDE_H: gl.constexpr,
537
+ Q_STRIDE_D: gl.constexpr,
538
+ K_STRIDE_B: gl.constexpr,
539
+ K_STRIDE_P: gl.constexpr,
540
+ K_STRIDE_H: gl.constexpr,
541
+ K_STRIDE_D: gl.constexpr,
542
+ V_STRIDE_B: gl.constexpr,
543
+ V_STRIDE_P: gl.constexpr,
544
+ V_STRIDE_H: gl.constexpr,
545
+ V_STRIDE_D: gl.constexpr,
546
+ SM_SCALE: gl.constexpr,
547
+ PAGE_TABLE_STRIDE: gl.constexpr,
548
+ PAGE_SIZE: gl.constexpr,
549
+ NUM_KV_SPLITS: gl.constexpr,
550
+ MAX_SEQLEN_Q: gl.constexpr,
551
+ NUM_Q_HEADS: gl.constexpr,
552
+ NUM_KV_HEADS: gl.constexpr,
553
+ HEAD_DIM: gl.constexpr,
554
+ BLOCK_M: gl.constexpr,
555
+ BLOCK_N: gl.constexpr,
556
+ IS_SLIDING: gl.constexpr,
557
+ WINDOW_LEFT: gl.constexpr,
558
+ REL_STRIDE_T: gl.constexpr,
559
+ REL_STRIDE_H: gl.constexpr,
560
+ REL_STRIDE_E: gl.constexpr,
561
+ REL_EXTENT: gl.constexpr,
562
+ REL_BIAS_QK_SCALE: gl.constexpr,
563
+ IS_FP8: gl.constexpr,
564
+):
565
+ cfg = AttentionConfig(
566
+ SM_SCALE,
567
+ PAGE_TABLE_STRIDE,
568
+ PAGE_SIZE,
569
+ NUM_KV_SPLITS,
570
+ MAX_SEQLEN_Q,
571
+ NUM_Q_HEADS,
572
+ NUM_KV_HEADS,
573
+ HEAD_DIM,
574
+ BLOCK_M,
575
+ BLOCK_N,
576
+ IS_SLIDING,
577
+ WINDOW_LEFT,
578
+ REL_EXTENT,
579
+ REL_BIAS_QK_SCALE,
580
+ IS_FP8,
581
+ InputStrides(Q_STRIDE_B, Q_STRIDE_H, Q_STRIDE_D),
582
+ InputStrides(REL_STRIDE_T, REL_STRIDE_H, REL_STRIDE_E),
583
+ PagedKVStrides(K_STRIDE_B, K_STRIDE_P, K_STRIDE_H, K_STRIDE_D),
584
+ PagedKVStrides(V_STRIDE_B, V_STRIDE_P, V_STRIDE_H, V_STRIDE_D),
585
+ )
586
+ program = AttentionProgram.create(
587
+ cfg,
588
+ q_ptr,
589
+ rel_logits_ptr,
590
+ k_cache_ptr,
591
+ v_cache_ptr,
592
+ page_table_ptr,
593
+ cache_seqlens_ptr,
594
+ mid_o_ptr,
595
+ mid_lse_ptr,
596
+ )
597
+ k_smem = gl.allocate_shared_memory(
598
+ k_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.k_smem_layout
599
+ )
600
+ v_smem = gl.allocate_shared_memory(
601
+ v_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.v_smem_layout
602
+ )
603
+
604
+ q = program.load_q()
605
+ m_i, l_i, acc = program.init_state()
606
+
607
+ physical_page = program.load_page(program.split_start)
608
+
609
+ for start_n in range(program.split_start, program.split_end, cfg.BLOCK_N):
610
+ program.issue_load_k(physical_page, k_smem)
611
+ program.issue_load_v(physical_page, v_smem)
612
+ physical_page = program.load_page(start_n + cfg.BLOCK_N)
613
+
614
+ async_copy.wait_group(1)
615
+ k = program.shared_load_k(k_smem)
616
+ qk = program.compute_qk(q, k)
617
+ qk = program.apply_rel_bias(qk, start_n)
618
+ qk = program.apply_kv_mask(qk, start_n)
619
+ p, m_i, l_i, acc = program.softmax(qk, m_i, l_i, acc)
620
+
621
+ async_copy.wait_group(0)
622
+ v = program.shared_load_v(v_smem)
623
+ acc = program.compute_pv(p, v, acc)
624
+
625
+ program.store_split(acc, l_i, m_i)
626
+
627
+
628
+@gluon.jit
629
+def _rel_mha_decode_sliding_fp16(
630
+ q_ptr,
631
+ rel_logits_ptr,
632
+ k_cache_ptr,
633
+ v_cache_ptr,
634
+ page_table_ptr,
635
+ cache_seqlens_ptr,
636
+ out_ptr,
637
+ Q_STRIDE_B: gl.constexpr,
638
+ Q_STRIDE_H: gl.constexpr,
639
+ Q_STRIDE_D: gl.constexpr,
640
+ K_STRIDE_B: gl.constexpr,
641
+ K_STRIDE_P: gl.constexpr,
642
+ K_STRIDE_H: gl.constexpr,
643
+ K_STRIDE_D: gl.constexpr,
644
+ V_STRIDE_B: gl.constexpr,
645
+ V_STRIDE_P: gl.constexpr,
646
+ V_STRIDE_H: gl.constexpr,
647
+ V_STRIDE_D: gl.constexpr,
648
+ SM_SCALE: gl.constexpr,
649
+ PAGE_TABLE_STRIDE: gl.constexpr,
650
+ PAGE_SIZE: gl.constexpr,
651
+ MAX_SEQLEN_Q: gl.constexpr,
652
+ NUM_Q_HEADS: gl.constexpr,
653
+ NUM_KV_HEADS: gl.constexpr,
654
+ HEAD_DIM: gl.constexpr,
655
+ BLOCK_M: gl.constexpr,
656
+ BLOCK_N: gl.constexpr,
657
+ IS_SLIDING: gl.constexpr,
658
+ WINDOW_LEFT: gl.constexpr,
659
+ REL_STRIDE_T: gl.constexpr,
660
+ REL_STRIDE_H: gl.constexpr,
661
+ REL_STRIDE_E: gl.constexpr,
662
+ REL_EXTENT: gl.constexpr,
663
+ REL_BIAS_QK_SCALE: gl.constexpr,
664
+ IS_FP8: gl.constexpr,
665
+):
666
+ cfg = AttentionConfig(
667
+ SM_SCALE,
668
+ PAGE_TABLE_STRIDE,
669
+ PAGE_SIZE,
670
+ 1,
671
+ MAX_SEQLEN_Q,
672
+ NUM_Q_HEADS,
673
+ NUM_KV_HEADS,
674
+ HEAD_DIM,
675
+ BLOCK_M,
676
+ BLOCK_N,
677
+ IS_SLIDING,
678
+ WINDOW_LEFT,
679
+ REL_EXTENT,
680
+ REL_BIAS_QK_SCALE,
681
+ IS_FP8,
682
+ InputStrides(Q_STRIDE_B, Q_STRIDE_H, Q_STRIDE_D),
683
+ InputStrides(REL_STRIDE_T, REL_STRIDE_H, REL_STRIDE_E),
684
+ PagedKVStrides(K_STRIDE_B, K_STRIDE_P, K_STRIDE_H, K_STRIDE_D),
685
+ PagedKVStrides(V_STRIDE_B, V_STRIDE_P, V_STRIDE_H, V_STRIDE_D),
686
+ )
687
+ program = AttentionProgram.create(
688
+ cfg,
689
+ q_ptr,
690
+ rel_logits_ptr,
691
+ k_cache_ptr,
692
+ v_cache_ptr,
693
+ page_table_ptr,
694
+ cache_seqlens_ptr,
695
+ out_ptr,
696
+ out_ptr,
697
+ )
698
+ k_smem = gl.allocate_shared_memory(
699
+ k_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.k_smem_layout
700
+ )
701
+ v_smem = gl.allocate_shared_memory(
702
+ v_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.v_smem_layout
703
+ )
704
+
705
+ q = program.load_q()
706
+ m_i, l_i, acc = program.init_state()
707
+
708
+ for start_n in range(program.split_start, program.split_end, cfg.BLOCK_N):
709
+ physical_page = program.load_page(start_n)
710
+ program.issue_load_k(physical_page, k_smem)
711
+ program.issue_load_v(physical_page, v_smem)
712
+ async_copy.wait_group(1)
713
+ k = program.shared_load_k(k_smem)
714
+ qk = program.compute_qk(q, k)
715
+ qk = program.apply_rel_bias(qk, start_n)
716
+ qk = program.apply_kv_mask(qk, start_n)
717
+ p, m_i, l_i, acc = program.softmax(qk, m_i, l_i, acc)
718
+
719
+ async_copy.wait_group(0)
720
+ v = program.shared_load_v(v_smem)
721
+ acc = program.compute_pv(p, v, acc)
722
+
723
+ program.store_output(acc, l_i)
724
+
725
+
726
+@gluon.jit
727
+def _rel_mha_decode_reduce_fp16(
728
+ mid_o_ptr,
729
+ mid_lse_ptr,
730
+ out_ptr,
731
+ cache_seqlens_ptr,
732
+ SM_SCALE: gl.constexpr,
733
+ PAGE_TABLE_STRIDE: gl.constexpr,
734
+ NUM_KV_SPLITS: gl.constexpr,
735
+ MAX_SEQLEN_Q: gl.constexpr,
736
+ PAGE_SIZE: gl.constexpr,
737
+ NUM_Q_HEADS: gl.constexpr,
738
+ NUM_KV_HEADS: gl.constexpr,
739
+ HEAD_DIM: gl.constexpr,
740
+ BLOCK_M: gl.constexpr,
741
+ BLOCK_N: gl.constexpr,
742
+ IS_SLIDING: gl.constexpr,
743
+ WINDOW_LEFT: gl.constexpr,
744
+ IS_FP8: gl.constexpr,
745
+):
746
+ cfg = AttentionConfig(
747
+ SM_SCALE,
748
+ PAGE_TABLE_STRIDE,
749
+ PAGE_SIZE,
750
+ NUM_KV_SPLITS,
751
+ MAX_SEQLEN_Q,
752
+ NUM_Q_HEADS,
753
+ NUM_KV_HEADS,
754
+ HEAD_DIM,
755
+ BLOCK_M,
756
+ BLOCK_N,
757
+ IS_SLIDING,
758
+ WINDOW_LEFT,
759
+ 1, # REL_EXTENT
760
+ 1.0, # REL_BIAS_QK_SCALE
761
+ IS_FP8,
762
+ InputStrides(1, 1, 1),
763
+ InputStrides(1, 1, 1),
764
+ PagedKVStrides(1, 1, 1, 1),
765
+ PagedKVStrides(1, 1, 1, 1),
766
+ )
767
+ q_index = gl.program_id(0)
768
+ batch = q_index // MAX_SEQLEN_Q
769
+ q_pos = q_index - batch * MAX_SEQLEN_Q
770
+ q_head = gl.program_id(1)
771
+ cache_len = gl.load(cache_seqlens_ptr + batch)
772
+ cache_len = cache_len - (MAX_SEQLEN_Q - 1 - q_pos)
773
+ cache_len = maximum(cache_len, 0)
774
+ if cfg.IS_SLIDING:
775
+ kv_start = cache_len - min(cache_len, cfg.WINDOW_LEFT + 1)
776
+ else:
777
+ kv_start = cache_len - cache_len
778
+ first_page = kv_start // cfg.PAGE_SIZE
779
+ end_page = cdiv(cache_len, cfg.PAGE_SIZE)
780
+ num_pages = end_page - first_page
781
+ pages_per_split = cdiv(num_pages, cfg.NUM_KV_SPLITS)
782
+
783
+ # SPLIT_TILE pads NUM_KV_SPLITS up to a power of 2.
784
+ SPLIT_TILE: gl.constexpr = 1 << (NUM_KV_SPLITS - 1).bit_length()
785
+ offs_s = gl.arange(0, SPLIT_TILE, layout=gl.SliceLayout(1, cfg.reduce_layout))
786
+ offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.reduce_layout))
787
+ # split_valid masks out empty splits and the power-of-2 padding tail.
788
+ split_start_page = first_page + offs_s * pages_per_split
789
+ split_end_page_raw = split_start_page + pages_per_split
790
+ split_end_page = gl.where(
791
+ split_end_page_raw < end_page, split_end_page_raw, end_page
792
+ )
793
+ split_start_tok = split_start_page * cfg.PAGE_SIZE
794
+ split_end_raw = split_end_page * cfg.PAGE_SIZE
795
+ split_end_tok = gl.where(split_end_raw < cache_len, split_end_raw, cache_len)
796
+ split_valid = (split_start_tok < split_end_tok) & (offs_s < cfg.NUM_KV_SPLITS)
797
+ # Load every split's partial output and lse.
798
+ base = (q_index * cfg.NUM_Q_HEADS + q_head) * cfg.NUM_KV_SPLITS + offs_s
799
+ part_lse = gl.load(mid_lse_ptr + base, mask=split_valid, other=-float("inf"))
800
+ o_off = base[:, None] * cfg.HEAD_DIM + offs_d[None, :]
801
+ part_o = cdna4.buffer_load(mid_o_ptr, o_off, mask=split_valid[:, None], other=0.0)
802
+
803
+ # Global softmax max over all splits.
804
+ m_i = max(part_lse, axis=0)
805
+ # Weighted sum of the split partials, normalized by the total softmax mass.
806
+ beta = gl.exp2(part_lse - m_i)
807
+ l_i = gl.sum(beta, axis=0)
808
+ acc = gl.sum(part_o * beta[:, None], axis=0)
809
+
810
+ out_base = (q_index * cfg.NUM_Q_HEADS + q_head) * cfg.HEAD_DIM
811
+ output = acc * (1.0 / l_i)
812
+ output = output.to(out_ptr.dtype.element_ty)
813
+ cdna4.buffer_store(output, out_ptr, out_base + offs_d)
814
+
815
+
816
+def _select_num_kv_splits(
817
+ *,
818
+ batch: int,
819
+ num_kv_heads: int,
820
+ num_groups: int,
821
+ num_pages: int,
822
+ sm_count: int,
823
+) -> int:
824
+ """Pick num_kv_splits to balance occupancy against reduce overhead.
825
+
826
+ The launch grid is (batch * num_kv_heads * num_groups) * num_kv_splits
827
+ work-groups. Too few splits under-fill the machine at low batch; too many
828
+ leave each split with a handful of pages, so the reduce kernel dominates.
829
+
830
+ Return the smaller of two candidate counts: splits_for_occupancy (enough to
831
+ fill ~wave_target waves of CUs) and splits_for_pages (~min_pages_per_split
832
+ pages per split), with the pages candidate clamped to [min_page_splits,
833
+ max_page_splits] so a short context still splits without launching empty work
834
+ and a long one does not over-split where reduce cost outgrows the decode win.
835
+ """
836
+ wave_target = 2
837
+ min_pages_per_split = 2
838
+ min_page_splits = 8
839
+ max_page_splits = 32
840
+
841
+ base_ctas = batch * num_kv_heads * num_groups
842
+ target_ctas = sm_count * wave_target
843
+ splits_for_occupancy = (target_ctas + base_ctas - 1) // base_ctas
844
+
845
+ splits_for_pages = num_pages // min_pages_per_split
846
+ min_page_splits = min(min_page_splits, num_pages)
847
+ if splits_for_pages < min_page_splits:
848
+ splits_for_pages = min_page_splits
849
+ if splits_for_pages > max_page_splits:
850
+ splits_for_pages = max_page_splits
851
+ return min(splits_for_occupancy, splits_for_pages)
852
+
853
+
854
+class LaunchConfig(NamedTuple):
855
+ num_q_heads: int
856
+ num_kv_heads: int
857
+ num_groups: int
858
+ head_dim: int
859
+ page_size: int
860
+ num_kv_splits: int
861
+ block_m: int
862
+ block_n: int
863
+ sm_scale: float
864
+ rel_bias_qk_scale: float
865
+ is_sliding: bool
866
+ window_left: int
867
+
868
+
869
+def get_config(
870
+ *,
871
+ q: torch.Tensor,
872
+ k_cache: torch.Tensor,
873
+ max_seqlen_k: int,
874
+ window_left: int,
875
+ softmax_scale: float | None,
876
+) -> LaunchConfig:
877
+ head_dim = q.shape[2]
878
+ page_size = k_cache.shape[1]
879
+ block_m = 16
880
+ block_n = page_size
881
+ group_size = q.shape[1] // k_cache.shape[2]
882
+ num_groups = math.ceil(group_size / block_m)
883
+ is_sliding = window_left >= 0
884
+ window_left = window_left if is_sliding else -1
885
+ if softmax_scale is None:
886
+ softmax_scale = 1.0 / math.sqrt(head_dim)
887
+ sm_scale = softmax_scale
888
+ effective_seqlen_k = (
889
+ min(max_seqlen_k, window_left + 1) if is_sliding else max_seqlen_k
890
+ )
891
+ num_pages = (effective_seqlen_k + page_size - 1) // page_size
892
+ if is_sliding:
893
+ num_kv_splits = 4
894
+ else:
895
+ num_kv_splits = _select_num_kv_splits(
896
+ batch=q.shape[0],
897
+ num_kv_heads=k_cache.shape[2],
898
+ num_groups=num_groups,
899
+ num_pages=num_pages,
900
+ sm_count=_GFX950_SM_COUNT,
901
+ )
902
+ return LaunchConfig(
903
+ num_q_heads=q.shape[1],
904
+ num_kv_heads=k_cache.shape[2],
905
+ num_groups=num_groups,
906
+ head_dim=head_dim,
907
+ page_size=page_size,
908
+ num_kv_splits=num_kv_splits,
909
+ block_m=block_m,
910
+ block_n=block_n,
911
+ sm_scale=sm_scale * _INV_LN2_VALUE,
912
+ rel_bias_qk_scale=1.0 / softmax_scale,
913
+ is_sliding=is_sliding,
914
+ window_left=window_left,
915
+ )
916
+
917
+
918
+def gluon_rel_mha_decode_gfx950(
919
+ q: torch.Tensor,
920
+ k_cache: torch.Tensor,
921
+ v_cache: torch.Tensor,
922
+ page_table: torch.Tensor,
923
+ cache_seqlens: torch.Tensor,
924
+ max_seqlen_k: int,
925
+ rel_logits: torch.Tensor,
926
+ cu_seqlens_q: torch.Tensor,
927
+ max_seqlen_q: int = 1,
928
+ window_left: int = -1,
929
+ softmax_scale: float | None = None,
930
+ q_scale: torch.Tensor | None = None,
931
+ k_scale: torch.Tensor | None = None,
932
+ v_scale: torch.Tensor | None = None,
933
+ out: torch.Tensor | None = None,
934
+) -> torch.Tensor:
935
+ total_q = q.shape[0]
936
+
937
+ config = get_config(
938
+ q=q,
939
+ k_cache=k_cache,
940
+ max_seqlen_k=max_seqlen_k,
941
+ window_left=window_left,
942
+ softmax_scale=softmax_scale,
943
+ )
944
+
945
+ is_fp8 = q.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
946
+ out_dtype = torch.bfloat16 if is_fp8 else q.dtype
947
+ if k_cache.shape[1] not in (64, 128, 256):
948
+ raise ValueError(
949
+ "gfx950 Gluon relative-attention decode requires page size "
950
+ f"64, 128, or 256; got {k_cache.shape[1]}"
951
+ )
952
+ output = (
953
+ torch.empty(q.shape, device=q.device, dtype=out_dtype) if out is None else out
954
+ )
955
+
956
+ # Always use split-k for full attention and for the small-batch sliding
957
+ # decode path. Sliding uses a fixed 8 splits, one page per split for the
958
+ # TP-4 local-attention shape with 512-token window and 64-token pages.
959
+ mid_o = torch.empty(
960
+ (total_q, config.num_q_heads, config.num_kv_splits, config.head_dim),
961
+ device=q.device,
962
+ dtype=torch.float32,
963
+ )
964
+ mid_lse = torch.empty(
965
+ (total_q, config.num_q_heads, config.num_kv_splits),
966
+ device=q.device,
967
+ dtype=torch.float32,
968
+ )
969
+
970
+ grid = (
971
+ total_q,
972
+ config.num_kv_heads * config.num_groups,
973
+ config.num_kv_splits,
974
+ )
975
+ _rel_mha_decode_fp16[grid](
976
+ q,
977
+ rel_logits,
978
+ k_cache,
979
+ v_cache,
980
+ page_table,
981
+ cache_seqlens,
982
+ mid_o,
983
+ mid_lse,
984
+ q.stride(0),
985
+ q.stride(1),
986
+ q.stride(2),
987
+ k_cache.stride(0),
988
+ k_cache.stride(1),
989
+ k_cache.stride(2),
990
+ k_cache.stride(3),
991
+ v_cache.stride(0),
992
+ v_cache.stride(1),
993
+ v_cache.stride(2),
994
+ v_cache.stride(3),
995
+ config.sm_scale,
996
+ page_table.stride(0),
997
+ config.page_size,
998
+ config.num_kv_splits,
999
+ max_seqlen_q,
1000
+ config.num_q_heads,
1001
+ config.num_kv_heads,
1002
+ config.head_dim,
1003
+ config.block_m,
1004
+ config.block_n,
1005
+ config.is_sliding,
1006
+ config.window_left,
1007
+ rel_logits.stride(0),
1008
+ rel_logits.stride(1),
1009
+ rel_logits.stride(2),
1010
+ rel_logits.shape[2],
1011
+ config.rel_bias_qk_scale,
1012
+ is_fp8,
1013
+ num_warps=1,
1014
+ )
1015
+
1016
+ reduce_grid = (total_q, config.num_q_heads)
1017
+ _rel_mha_decode_reduce_fp16[reduce_grid](
1018
+ mid_o,
1019
+ mid_lse,
1020
+ output,
1021
+ cache_seqlens,
1022
+ config.sm_scale,
1023
+ page_table.stride(0),
1024
+ config.num_kv_splits,
1025
+ max_seqlen_q,
1026
+ config.page_size,
1027
+ config.num_q_heads,
1028
+ config.num_kv_heads,
1029
+ config.head_dim,
1030
+ config.block_m,
1031
+ config.block_n,
1032
+ config.is_sliding,
1033
+ config.window_left,
1034
+ is_fp8,
1035
+ num_warps=1,
1036
+ )
1037
+ return output
vllm/models/inkling/amd/ops/gluon/rel_mha_extend_gfx950.py
new
+760
@@ -0,0 +1,760 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# Copyright (c) 2026 LightSeek Foundation
5
+#
6
+# Permission is hereby granted, free of charge, to any person obtaining a copy
7
+# of this software and associated documentation files (the "Software"), to deal
8
+# in the Software without restriction, including without limitation the rights
9
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
10
+# copies of the Software, and to permit persons to whom the Software is
11
+# furnished to do so, subject to the following conditions:
12
+#
13
+# The above copyright notice and this permission notice shall be included in
14
+# all copies or substantial portions of the Software.
15
+#
16
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22
+# SOFTWARE.
23
+
24
+"""rel_mha extend Gluon kernel for AMD GFX950.
25
+
26
+This handles ragged, multi-token queries against a paged KV cache. The query
27
+axis is tiled into the MFMA ``M`` dimension (prefill-style): ``BLOCK_M`` query
28
+rows of a single q-head share each paged KV tile, so every KV tile is loaded
29
+once and reused across all rows in the tile. The grid is
30
+``(blocks_per_req, batch, n_heads)`` -- all host-known sizes -- and each program
31
+self-locates its request from ``cu_seqlens_q`` / ``cache_seqlens`` in-kernel, so
32
+the launch stays CUDA-graph static with no device->host sync.
33
+
34
+Visibility per query row depends on ``is_causal`` and the optional sliding
35
+window:
36
+
37
+* ``is_causal=False``: every query token attends the full visible cache, i.e.
38
+ ``visible_kv = cache_seqlens[batch]``.
39
+* ``is_causal=True``: query tokens are a causal suffix, so the ``i``-th query
40
+ token of a request (0-indexed) attends ``prefix + i + 1`` tokens, where
41
+ ``prefix = cache_seqlens[batch] - query_len[batch]``.
42
+
43
+Causal masking only touches the few KV tiles that reach the diagonal; the long
44
+prefix is a mask-free fast path. Sinks and sliding windows (causal or not) are
45
+applied per tile. This is the sole Gluon extend implementation.
46
+"""
47
+
48
+import math
49
+
50
+import torch
51
+
52
+from vllm.models.inkling.amd.ops.gluon.utils import (
53
+ _INV_LN2_VALUE,
54
+ _LN2,
55
+ InputStrides,
56
+ PagedKVStrides,
57
+ max,
58
+ maximum,
59
+)
60
+from vllm.triton_utils import aggregate, gl, gluon
61
+
62
+cdna4 = gl.amd.cdna4
63
+async_copy = cdna4.async_copy
64
+
65
+
66
+# ===-----------------------------------------------------------------------===#
67
+# Query-batched path
68
+#
69
+# This path tiles BLOCK_M query rows (of a single q-head) into the MFMA M
70
+# dimension -- like the prefill kernel -- so each paged KV tile is loaded once
71
+# and reused across all BLOCK_M rows. KV still comes from the paged cache
72
+# (decode-style async page loads). Causal masking only touches the few KV tiles
73
+# that reach the diagonal; the long prefix is a mask-free fast path.
74
+# ===-----------------------------------------------------------------------===#
75
+
76
+_EXTEND_SHORT_Q_BLOCK_M = 64
77
+_EXTEND_SHORT_Q_NUM_WARPS = 2
78
+_EXTEND_LONG_Q_BLOCK_M = 128
79
+_EXTEND_LONG_Q_NUM_WARPS = 4
80
+
81
+
82
+def _select_extend_tile(max_seqlen_q: int, page_size: int) -> tuple[int, int, int]:
83
+ """Return (BLOCK_M, BLOCK_N, NUM_WARPS) for the given max query length.
84
+
85
+ Queries that fit in a single short tile use it (least padding, most
86
+ occupancy); longer ones use the tall tile that covers more rows per
87
+ shared-KV pass.
88
+ """
89
+ if max_seqlen_q <= _EXTEND_SHORT_Q_BLOCK_M:
90
+ return _EXTEND_SHORT_Q_BLOCK_M, page_size, _EXTEND_SHORT_Q_NUM_WARPS
91
+ return _EXTEND_LONG_Q_BLOCK_M, page_size, _EXTEND_LONG_Q_NUM_WARPS
92
+
93
+
94
+@aggregate
95
+class ExtendConfig:
96
+ N_HEADS: gl.constexpr
97
+ N_KV_HEADS: gl.constexpr
98
+ GROUP_SIZE: gl.constexpr
99
+ HEAD_DIM: gl.constexpr
100
+ SM_SCALE: gl.constexpr
101
+ BLOCK_M: gl.constexpr
102
+ BLOCK_N: gl.constexpr
103
+ NUM_WARPS: gl.constexpr
104
+ PAGE_SIZE: gl.constexpr
105
+ PAGE_TABLE_STRIDE: gl.constexpr
106
+ IS_CAUSAL: gl.constexpr
107
+ HAS_LSE: gl.constexpr
108
+ WINDOW_LEFT: gl.constexpr
109
+ REL_EXTENT: gl.constexpr
110
+ REL_BIAS_QK_SCALE: gl.constexpr
111
+ IS_FP8: gl.constexpr
112
+ q_strides: InputStrides
113
+ rel_strides: InputStrides
114
+ k_strides: PagedKVStrides
115
+ v_strides: PagedKVStrides
116
+ qk_layout: gl.constexpr
117
+ pv_layout: gl.constexpr
118
+ q_layout: gl.constexpr
119
+ k_layout: gl.constexpr
120
+ p_layout: gl.constexpr
121
+ v_layout: gl.constexpr
122
+ load_layout: gl.constexpr
123
+ store_layout: gl.constexpr
124
+ k_smem_layout: gl.constexpr
125
+ v_smem_layout: gl.constexpr
126
+
127
+ @gluon.constexpr_function
128
+ def __init__(
129
+ self,
130
+ N_HEADS,
131
+ N_KV_HEADS,
132
+ HEAD_DIM,
133
+ SM_SCALE,
134
+ BLOCK_M,
135
+ BLOCK_N,
136
+ NUM_WARPS,
137
+ PAGE_SIZE,
138
+ PAGE_TABLE_STRIDE,
139
+ IS_CAUSAL,
140
+ HAS_LSE,
141
+ WINDOW_LEFT,
142
+ REL_EXTENT,
143
+ REL_BIAS_QK_SCALE,
144
+ IS_FP8,
145
+ q_strides,
146
+ rel_strides,
147
+ k_strides,
148
+ v_strides,
149
+ ):
150
+ assert HEAD_DIM in (64, 128)
151
+ assert BLOCK_N == PAGE_SIZE
152
+
153
+ instr_shape = [32, 32, 16]
154
+ mfma_layout = gl.amd.AMDMFMALayout(
155
+ version=4,
156
+ instr_shape=instr_shape,
157
+ transposed=True,
158
+ warps_per_cta=[NUM_WARPS, 1],
159
+ )
160
+ qk_layout = mfma_layout
161
+ pv_layout = mfma_layout
162
+ # Elements loaded per lane depend on the input dtype, matching qk_kw.
163
+ # load_threads is how many lanes span HEAD_DIM.
164
+ load_vec = 16 if IS_FP8 else 8
165
+ load_threads = HEAD_DIM // load_vec
166
+ load_layout = gl.BlockedLayout(
167
+ [1, load_vec], [64 // load_threads, load_threads], [NUM_WARPS, 1], [1, 0]
168
+ )
169
+ # Output is always 16-bit, so a 128-bit store has 8 elements.
170
+ # store_threads is how many lanes span HEAD_DIM.
171
+ store_vec = 8
172
+ store_threads = HEAD_DIM // store_vec
173
+ store_layout = gl.BlockedLayout(
174
+ [1, store_vec], [64 // store_threads, store_threads], [NUM_WARPS, 1], [1, 0]
175
+ )
176
+ # Padding interval is 64 lanes * load_vec elems.
177
+ pad_interval = 64 * load_vec
178
+ # Empirically tuned.
179
+ pad_k = 16 if IS_FP8 else 8
180
+ pad_v = 16 if IS_FP8 else 32
181
+ k_smem_layout = gl.PaddedSharedLayout.with_identity_for(
182
+ [[pad_interval, pad_k]], [BLOCK_N, HEAD_DIM], [1, 0]
183
+ )
184
+ v_smem_layout = gl.PaddedSharedLayout.with_identity_for(
185
+ [[pad_interval, pad_v]], [BLOCK_N, HEAD_DIM], [1, 0]
186
+ )
187
+
188
+ self.N_HEADS = gl.constexpr(N_HEADS)
189
+ self.N_KV_HEADS = gl.constexpr(N_KV_HEADS)
190
+ self.GROUP_SIZE = gl.constexpr(N_HEADS // N_KV_HEADS)
191
+ self.HEAD_DIM = gl.constexpr(HEAD_DIM)
192
+ self.SM_SCALE = gl.constexpr(SM_SCALE)
193
+ self.BLOCK_M = gl.constexpr(BLOCK_M)
194
+ self.BLOCK_N = gl.constexpr(BLOCK_N)
195
+ self.NUM_WARPS = gl.constexpr(NUM_WARPS)
196
+ self.PAGE_SIZE = gl.constexpr(PAGE_SIZE)
197
+ self.PAGE_TABLE_STRIDE = gl.constexpr(PAGE_TABLE_STRIDE)
198
+ self.IS_CAUSAL = gl.constexpr(IS_CAUSAL)
199
+ self.HAS_LSE = gl.constexpr(HAS_LSE)
200
+ self.WINDOW_LEFT = gl.constexpr(WINDOW_LEFT)
201
+ self.REL_EXTENT = gl.constexpr(REL_EXTENT)
202
+ self.REL_BIAS_QK_SCALE = gl.constexpr(REL_BIAS_QK_SCALE)
203
+ self.IS_FP8 = gl.constexpr(IS_FP8)
204
+ self.q_strides = q_strides
205
+ self.rel_strides = rel_strides
206
+ self.k_strides = k_strides
207
+ self.v_strides = v_strides
208
+ self.qk_layout = gl.constexpr(qk_layout)
209
+ self.pv_layout = gl.constexpr(pv_layout)
210
+ # qk_kw is derived from a 128-bit load / dtype bitwidth.
211
+ # pv_kw is empirically tuned.
212
+ qk_kw = 16 if IS_FP8 else 8
213
+ pv_kw = 16 if IS_FP8 else 4
214
+ self.q_layout = gl.constexpr(gl.DotOperandLayout(0, qk_layout, k_width=qk_kw))
215
+ self.k_layout = gl.constexpr(gl.DotOperandLayout(1, qk_layout, k_width=qk_kw))
216
+ self.p_layout = gl.constexpr(gl.DotOperandLayout(0, pv_layout, k_width=pv_kw))
217
+ self.v_layout = gl.constexpr(gl.DotOperandLayout(1, pv_layout, k_width=pv_kw))
218
+ self.load_layout = gl.constexpr(load_layout)
219
+ self.store_layout = gl.constexpr(store_layout)
220
+ self.k_smem_layout = gl.constexpr(k_smem_layout)
221
+ self.v_smem_layout = gl.constexpr(v_smem_layout)
222
+
223
+
224
+@aggregate
225
+class ExtendProgram:
226
+ cfg: gl.constexpr
227
+ q_ptr: gl.tensor
228
+ rel_logits_ptr: gl.tensor
229
+ k_cache_ptr: gl.tensor
230
+ v_cache_ptr: gl.tensor
231
+ page_table_ptr: gl.tensor
232
+ output_ptr: gl.tensor
233
+ lse_ptr: gl.tensor
234
+ batch: gl.tensor
235
+ q_head: gl.tensor
236
+ kv_head: gl.tensor
237
+ q_start: gl.tensor
238
+ seq_base: gl.tensor
239
+ seq_len: gl.tensor
240
+ prefix: gl.tensor
241
+ cache_len: gl.tensor
242
+
243
+ @gluon.constexpr_function
244
+ def __init__(
245
+ self,
246
+ cfg,
247
+ q_ptr,
248
+ rel_logits_ptr,
249
+ k_cache_ptr,
250
+ v_cache_ptr,
251
+ page_table_ptr,
252
+ output_ptr,
253
+ lse_ptr,
254
+ batch,
255
+ q_head,
256
+ kv_head,
257
+ q_start,
258
+ seq_base,
259
+ seq_len,
260
+ prefix,
261
+ cache_len,
262
+ ):
263
+ self.cfg = gl.constexpr(cfg)
264
+ self.q_ptr = q_ptr
265
+ self.rel_logits_ptr = rel_logits_ptr
266
+ self.k_cache_ptr = k_cache_ptr
267
+ self.v_cache_ptr = v_cache_ptr
268
+ self.page_table_ptr = page_table_ptr
269
+ self.output_ptr = output_ptr
270
+ self.lse_ptr = lse_ptr
271
+ self.batch = batch
272
+ self.q_head = q_head
273
+ self.kv_head = kv_head
274
+ self.q_start = q_start
275
+ self.seq_base = seq_base
276
+ self.seq_len = seq_len
277
+ self.prefix = prefix
278
+ self.cache_len = cache_len
279
+
280
+ @gluon.jit
281
+ def create(
282
+ cfg,
283
+ q_ptr,
284
+ rel_logits_ptr,
285
+ k_cache_ptr,
286
+ v_cache_ptr,
287
+ page_table_ptr,
288
+ output_ptr,
289
+ lse_ptr,
290
+ cu_seqlens_q_ptr,
291
+ cache_seqlens_ptr,
292
+ ):
293
+ # Gluon treats ``cfg`` as the constexpr factory argument, while mypy
294
+ # models the first parameter of this aggregate method as ``self``.
295
+ block_in_req = gl.program_id(0)
296
+ batch = gl.program_id(1)
297
+ q_head = gl.program_id(2)
298
+ q_start = block_in_req * cfg.BLOCK_M # type: ignore[attr-defined]
299
+ seq_base = gl.load(cu_seqlens_q_ptr + batch)
300
+ seq_end = gl.load(cu_seqlens_q_ptr + batch + 1)
301
+ seq_len = seq_end - seq_base
302
+ cache_len = gl.load(cache_seqlens_ptr + batch)
303
+ prefix = cache_len - seq_len
304
+ kv_head = q_head // cfg.GROUP_SIZE # type: ignore[attr-defined]
305
+ return ExtendProgram(
306
+ gl.constexpr(cfg),
307
+ q_ptr,
308
+ rel_logits_ptr,
309
+ k_cache_ptr,
310
+ v_cache_ptr,
311
+ page_table_ptr,
312
+ output_ptr,
313
+ lse_ptr,
314
+ batch,
315
+ q_head,
316
+ kv_head,
317
+ q_start,
318
+ seq_base,
319
+ seq_len,
320
+ prefix,
321
+ cache_len,
322
+ )
323
+
324
+ @gluon.jit
325
+ def load_q(self):
326
+ cfg = self.cfg
327
+ offs_m = self.q_start + gl.arange(
328
+ 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.q_layout)
329
+ )
330
+ offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.q_layout))
331
+ row = self.seq_base + offs_m
332
+ offsets = cfg.q_strides.offsets(row[:, None], self.q_head, offs_d[None, :])
333
+ mask = offs_m[:, None] < self.seq_len
334
+ return cdna4.buffer_load(self.q_ptr, offsets, mask=mask, other=0.0)
335
+
336
+ @gluon.jit
337
+ def load_page(self, start_n):
338
+ cfg = self.cfg
339
+ page_index = start_n // cfg.PAGE_SIZE
340
+ valid = start_n < self.cache_len
341
+ return gl.load(
342
+ self.page_table_ptr + self.batch * cfg.PAGE_TABLE_STRIDE + page_index,
343
+ mask=valid,
344
+ other=0,
345
+ )
346
+
347
+ @gluon.jit
348
+ def issue_load_k(self, physical_page, k_smem):
349
+ cfg = self.cfg
350
+ offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout))
351
+ offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout))
352
+ offsets = cfg.k_strides.offsets(
353
+ physical_page,
354
+ offs_n[:, None],
355
+ self.kv_head,
356
+ offs_d[None, :],
357
+ )
358
+ async_copy.global_load_to_shared(k_smem, self.k_cache_ptr + offsets)
359
+ async_copy.commit_group()
360
+
361
+ @gluon.jit
362
+ def issue_load_v(self, physical_page, v_smem):
363
+ cfg = self.cfg
364
+ offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout))
365
+ offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout))
366
+ offsets = cfg.v_strides.offsets(
367
+ physical_page,
368
+ offs_n[:, None],
369
+ self.kv_head,
370
+ offs_d[None, :],
371
+ )
372
+ async_copy.global_load_to_shared(v_smem, self.v_cache_ptr + offsets)
373
+ async_copy.commit_group()
374
+
375
+ @gluon.jit
376
+ def shared_load_k(self, k_smem):
377
+ cfg = self.cfg
378
+ k_buffer = k_smem.permute([1, 0])
379
+ return k_buffer.load(cfg.k_layout)
380
+
381
+ @gluon.jit
382
+ def shared_load_v(self, v_smem):
383
+ cfg = self.cfg
384
+ return v_smem.load(cfg.v_layout)
385
+
386
+ @gluon.jit
387
+ def compute_qk(self, q, k):
388
+ cfg = self.cfg
389
+ qk = gl.zeros(
390
+ [cfg.BLOCK_M, cfg.BLOCK_N], dtype=gl.float32, layout=cfg.qk_layout
391
+ )
392
+ return cdna4.mfma(q, k, qk)
393
+
394
+ @gluon.jit
395
+ def apply_rel_bias(self, qk, start_n):
396
+ cfg = self.cfg
397
+ offs_m = self.q_start + gl.arange(
398
+ 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.qk_layout)
399
+ )
400
+ offs_n_abs = start_n + gl.arange(
401
+ 0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout)
402
+ )
403
+ q_pos = self.prefix + offs_m
404
+ rel_dist = q_pos[:, None] - offs_n_abs[None, :]
405
+ rel_valid = (rel_dist >= 0) & (rel_dist < cfg.REL_EXTENT)
406
+ rel_idx = gl.where(rel_dist >= 0, rel_dist, 0)
407
+ rel_idx = gl.where(rel_idx < cfg.REL_EXTENT, rel_idx, cfg.REL_EXTENT - 1)
408
+ offsets = cfg.rel_strides.offsets(
409
+ self.seq_base + offs_m[:, None], self.q_head, rel_idx
410
+ )
411
+ mask = (offs_m[:, None] < self.seq_len) & rel_valid
412
+ rel_bias = cdna4.buffer_load(
413
+ self.rel_logits_ptr, offsets, mask=mask, other=0.0
414
+ ).to(gl.float32)
415
+ return qk + rel_bias * cfg.REL_BIAS_QK_SCALE
416
+
417
+ @gluon.jit
418
+ def compute_pv(self, p, v, acc):
419
+ return cdna4.mfma(p, v, acc)
420
+
421
+ @gluon.jit
422
+ def init_attention_state(self):
423
+ cfg = self.cfg
424
+ m_i = gl.full(
425
+ [cfg.BLOCK_M],
426
+ value=-float("inf"),
427
+ dtype=gl.float32,
428
+ layout=gl.SliceLayout(1, cfg.pv_layout),
429
+ )
430
+ l_i = gl.full(
431
+ [cfg.BLOCK_M],
432
+ value=0,
433
+ dtype=gl.float32,
434
+ layout=gl.SliceLayout(1, cfg.pv_layout),
435
+ )
436
+ acc = gl.zeros(
437
+ [cfg.BLOCK_M, cfg.HEAD_DIM], dtype=gl.float32, layout=cfg.pv_layout
438
+ )
439
+ return m_i, l_i, acc
440
+
441
+ @gluon.jit
442
+ def softmax(self, qk, m_i, l_i, acc):
443
+ cfg = self.cfg
444
+ # In sliding window case, some rows can see fully masked tiles before
445
+ # any valid KV. Guard the online softmax state so `-inf - -inf` does not
446
+ # produce NaNs.
447
+ HAS_INVALID: gl.constexpr = cfg.WINDOW_LEFT >= 0
448
+
449
+ row_max = max(qk, 1)
450
+ m_new = maximum(m_i, row_max)
451
+ m_new_scaled = m_new * cfg.SM_SCALE
452
+ if HAS_INVALID:
453
+ invalid = m_new == -float("inf")
454
+ m_new_scaled = gl.where(invalid, 0.0, m_new_scaled)
455
+
456
+ qk_shifted = qk * cfg.SM_SCALE - m_new_scaled[:, None]
457
+ p = gl.exp2(qk_shifted)
458
+ m_diff = m_i * cfg.SM_SCALE - m_new_scaled
459
+ if HAS_INVALID:
460
+ m_diff = gl.where(invalid, 0.0, m_diff)
461
+
462
+ alpha = gl.exp2(m_diff)
463
+ l_ij = gl.sum(p, axis=1)
464
+ l_i = l_i * alpha + l_ij
465
+ acc = acc * alpha[:, None]
466
+ p = p.to(self.q_ptr.dtype.element_ty)
467
+ p = gl.convert_layout(p, cfg.p_layout)
468
+ return p, m_new, l_i, acc
469
+
470
+ @gluon.jit
471
+ def store_output(self, output):
472
+ cfg = self.cfg
473
+ offs_m = self.q_start + gl.arange(
474
+ 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.store_layout)
475
+ )
476
+ offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.store_layout))
477
+ offsets = (
478
+ ((self.seq_base + offs_m[:, None]) * cfg.N_HEADS + self.q_head)
479
+ * cfg.HEAD_DIM
480
+ + offs_d[None, :]
481
+ ).to(gl.int32)
482
+ mask = offs_m[:, None] < self.seq_len
483
+ output = output.to(self.output_ptr.dtype.element_ty)
484
+ cdna4.buffer_store(output, self.output_ptr, offsets, mask=mask)
485
+
486
+ @gluon.jit
487
+ def store_lse(self, l_i, m_i):
488
+ cfg = self.cfg
489
+ if cfg.HAS_LSE:
490
+ offs_m = self.q_start + gl.arange(
491
+ 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.pv_layout)
492
+ )
493
+ offsets = ((self.seq_base + offs_m) * cfg.N_HEADS + self.q_head).to(
494
+ gl.int32
495
+ )
496
+ mask = offs_m < self.seq_len
497
+ lse_l_i = gl.where(l_i > 0.0, l_i, 1.0)
498
+ # Softmax runs in base-2 (exp2 hardware fast path), so m_i*SM_SCALE +
499
+ # log2(l_i) is the LSE in base-2 units. Convert to natural log (the
500
+ # public op contract / torch.logsumexp convention) by scaling by ln2.
501
+ lse = (m_i * cfg.SM_SCALE + gl.log2(lse_l_i)) * _LN2
502
+ cdna4.buffer_store(lse, self.lse_ptr, offsets, mask=mask)
503
+
504
+
505
+@gluon.jit
506
+def _rel_mha_extend_fp16(
507
+ q_ptr,
508
+ rel_logits_ptr,
509
+ k_cache_ptr,
510
+ v_cache_ptr,
511
+ page_table_ptr,
512
+ output_ptr,
513
+ lse_ptr,
514
+ cu_seqlens_q_ptr,
515
+ cache_seqlens_ptr,
516
+ Q_STRIDE_T: gl.constexpr,
517
+ Q_STRIDE_H: gl.constexpr,
518
+ Q_STRIDE_D: gl.constexpr,
519
+ K_STRIDE_B: gl.constexpr,
520
+ K_STRIDE_P: gl.constexpr,
521
+ K_STRIDE_H: gl.constexpr,
522
+ K_STRIDE_D: gl.constexpr,
523
+ V_STRIDE_B: gl.constexpr,
524
+ V_STRIDE_P: gl.constexpr,
525
+ V_STRIDE_H: gl.constexpr,
526
+ V_STRIDE_D: gl.constexpr,
527
+ SM_SCALE: gl.constexpr,
528
+ PAGE_TABLE_STRIDE: gl.constexpr,
529
+ PAGE_SIZE: gl.constexpr,
530
+ N_HEADS: gl.constexpr,
531
+ N_KV_HEADS: gl.constexpr,
532
+ HEAD_DIM: gl.constexpr,
533
+ BLOCK_M: gl.constexpr,
534
+ BLOCK_N: gl.constexpr,
535
+ NUM_WARPS: gl.constexpr,
536
+ IS_CAUSAL: gl.constexpr,
537
+ HAS_LSE: gl.constexpr,
538
+ WINDOW_LEFT: gl.constexpr,
539
+ REL_STRIDE_T: gl.constexpr,
540
+ REL_STRIDE_H: gl.constexpr,
541
+ REL_STRIDE_E: gl.constexpr,
542
+ REL_EXTENT: gl.constexpr,
543
+ REL_BIAS_QK_SCALE: gl.constexpr,
544
+ IS_FP8: gl.constexpr,
545
+):
546
+ cfg = ExtendConfig(
547
+ N_HEADS,
548
+ N_KV_HEADS,
549
+ HEAD_DIM,
550
+ SM_SCALE,
551
+ BLOCK_M,
552
+ BLOCK_N,
553
+ NUM_WARPS,
554
+ PAGE_SIZE,
555
+ PAGE_TABLE_STRIDE,
556
+ IS_CAUSAL,
557
+ HAS_LSE,
558
+ WINDOW_LEFT,
559
+ REL_EXTENT,
560
+ REL_BIAS_QK_SCALE,
561
+ IS_FP8,
562
+ InputStrides(Q_STRIDE_T, Q_STRIDE_H, Q_STRIDE_D),
563
+ InputStrides(REL_STRIDE_T, REL_STRIDE_H, REL_STRIDE_E),
564
+ PagedKVStrides(K_STRIDE_B, K_STRIDE_P, K_STRIDE_H, K_STRIDE_D),
565
+ PagedKVStrides(V_STRIDE_B, V_STRIDE_P, V_STRIDE_H, V_STRIDE_D),
566
+ )
567
+ program = ExtendProgram.create(
568
+ cfg,
569
+ q_ptr,
570
+ rel_logits_ptr,
571
+ k_cache_ptr,
572
+ v_cache_ptr,
573
+ page_table_ptr,
574
+ output_ptr,
575
+ lse_ptr,
576
+ cu_seqlens_q_ptr,
577
+ cache_seqlens_ptr,
578
+ )
579
+ # Over-provisioned tile past this request's real query rows: nothing to do.
580
+ if program.q_start >= program.seq_len:
581
+ return
582
+ k_smem = gl.allocate_shared_memory(
583
+ k_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.k_smem_layout
584
+ )
585
+ v_smem = gl.allocate_shared_memory(
586
+ v_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.v_smem_layout
587
+ )
588
+
589
+ q = program.load_q()
590
+ m_i, l_i, acc = program.init_attention_state()
591
+
592
+ offs_m_q = program.q_start + gl.arange(
593
+ 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.qk_layout)
594
+ )
595
+ offs_n_q = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout))
596
+ diag_row = program.prefix + offs_m_q
597
+
598
+ if IS_CAUSAL:
599
+ kv_end = min(program.cache_len, program.prefix + program.q_start + cfg.BLOCK_M)
600
+ else:
601
+ kv_end = program.cache_len
602
+
603
+ # Sliding window (inclusive-left, matches flash-attn window_size=(W, 0)):
604
+ # skip KV tiles entirely below the window's lower edge. The tile's top query
605
+ # row sits at absolute position pos_top; its window opens at
606
+ # pos_top - WINDOW_LEFT (that key is visible -> W + 1 keys total). The min()
607
+ # form clamps to 0 without the shadowed builtin max().
608
+ if cfg.WINDOW_LEFT >= 0:
609
+ pos_top = program.prefix + program.q_start
610
+ kv_start = pos_top - min(pos_top, cfg.WINDOW_LEFT)
611
+ kv_start = (kv_start // cfg.BLOCK_N) * cfg.BLOCK_N
612
+ else:
613
+ kv_start = 0
614
+
615
+ for start_n in range(kv_start, kv_end, cfg.BLOCK_N):
616
+ physical_page = program.load_page(start_n)
617
+ program.issue_load_k(physical_page, k_smem)
618
+ program.issue_load_v(physical_page, v_smem)
619
+
620
+ async_copy.wait_group(1)
621
+ k = program.shared_load_k(k_smem)
622
+ qk = program.compute_qk(q, k)
623
+ qk = program.apply_rel_bias(qk, start_n)
624
+
625
+ if cfg.WINDOW_LEFT >= 0:
626
+ # Window lower edge + cache bound always apply; causal upper edge only
627
+ # when IS_CAUSAL (independent layering keeps non-causal + window correct).
628
+ offs_n_abs = start_n + offs_n_q
629
+ mask = (offs_n_abs[None, :] >= diag_row[:, None] - cfg.WINDOW_LEFT) & (
630
+ offs_n_abs[None, :] < program.cache_len
631
+ )
632
+ if IS_CAUSAL:
633
+ mask &= offs_n_abs[None, :] <= diag_row[:, None]
634
+ qk = gl.where(mask, qk, -float("inf"))
635
+ elif IS_CAUSAL:
636
+ if start_n + cfg.BLOCK_N > program.prefix + program.q_start:
637
+ offs_n_abs = start_n + offs_n_q
638
+ mask = (offs_n_abs[None, :] <= diag_row[:, None]) & (
639
+ offs_n_abs[None, :] < program.cache_len
640
+ )
641
+ qk = gl.where(mask, qk, -float("inf"))
642
+ else:
643
+ if start_n + cfg.BLOCK_N > program.cache_len:
644
+ offs_n_abs = start_n + offs_n_q
645
+ qk = gl.where(
646
+ offs_n_abs[None, :] < program.cache_len, qk, -float("inf")
647
+ )
648
+
649
+ p, m_i, l_i, acc = program.softmax(qk, m_i, l_i, acc)
650
+
651
+ async_copy.wait_group(0)
652
+ v = program.shared_load_v(v_smem)
653
+ acc = program.compute_pv(p, v, acc)
654
+
655
+ denom = gl.where(l_i > 0.0, l_i, 1.0)
656
+ output = acc * (1.0 / denom)[:, None]
657
+ output = gl.convert_layout(output, cfg.store_layout)
658
+ program.store_output(output)
659
+ program.store_lse(l_i, m_i)
660
+
661
+
662
+def gluon_rel_mha_extend_gfx950(
663
+ q: torch.Tensor,
664
+ cu_seqlens_q: torch.Tensor,
665
+ cu_seqlens_kv: torch.Tensor,
666
+ k_cache: torch.Tensor,
667
+ v_cache: torch.Tensor,
668
+ page_table: torch.Tensor,
669
+ cache_seqlens: torch.Tensor,
670
+ window_left: int = -1,
671
+ return_lse: bool = False,
672
+ max_seqlen_q: int = 1,
673
+ max_seqlen_k: int = 1,
674
+ rel_logits: torch.Tensor | None = None,
675
+ softmax_scale: float | None = None,
676
+ q_scale: torch.Tensor | None = None,
677
+ k_scale: torch.Tensor | None = None,
678
+ v_scale: torch.Tensor | None = None,
679
+ out: torch.Tensor | None = None,
680
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
681
+ assert rel_logits is not None
682
+ head_dim = q.shape[2]
683
+ n_heads = q.shape[1]
684
+ n_kv_heads = k_cache.shape[2]
685
+ page_size = k_cache.shape[1]
686
+ block_m, block_n, num_warps = _select_extend_tile(max_seqlen_q, page_size)
687
+ if softmax_scale is None:
688
+ softmax_scale = 1.0 / math.sqrt(head_dim)
689
+ sm_scale = softmax_scale * _INV_LN2_VALUE
690
+ rel_bias_qk_scale = 1.0 / softmax_scale
691
+
692
+ # max_seqlen_q must be >= the true max query length; extra tiles early-exit.
693
+ batch = cu_seqlens_q.shape[0] - 1
694
+ safe_max_q = max_seqlen_q if max_seqlen_q > 0 else 1
695
+ blocks_per_req = (safe_max_q + block_m - 1) // block_m
696
+ cu_q_i32 = cu_seqlens_q.to(torch.int32).contiguous()
697
+ cache_i32 = cache_seqlens.to(torch.int32).contiguous()
698
+
699
+ is_fp8 = q.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
700
+ out_dtype = torch.bfloat16 if is_fp8 else q.dtype
701
+ if page_size not in (64, 128, 256):
702
+ raise ValueError(
703
+ "gfx950 Gluon relative-attention extend requires page size "
704
+ f"64, 128, or 256; got {page_size}"
705
+ )
706
+ output = (
707
+ torch.empty(q.shape, device=q.device, dtype=out_dtype) if out is None else out
708
+ )
709
+ if return_lse:
710
+ lse = torch.empty((q.shape[0], n_heads), device=q.device, dtype=torch.float32)
711
+ lse_arg = lse
712
+ else:
713
+ lse = None
714
+ lse_arg = q
715
+ grid = (blocks_per_req, batch, n_heads)
716
+ _rel_mha_extend_fp16[grid](
717
+ q,
718
+ rel_logits,
719
+ k_cache,
720
+ v_cache,
721
+ page_table,
722
+ output,
723
+ lse_arg,
724
+ cu_q_i32,
725
+ cache_i32,
726
+ q.stride(0),
727
+ q.stride(1),
728
+ q.stride(2),
729
+ k_cache.stride(0),
730
+ k_cache.stride(1),
731
+ k_cache.stride(2),
732
+ k_cache.stride(3),
733
+ v_cache.stride(0),
734
+ v_cache.stride(1),
735
+ v_cache.stride(2),
736
+ v_cache.stride(3),
737
+ sm_scale,
738
+ page_table.stride(0),
739
+ page_size,
740
+ n_heads,
741
+ n_kv_heads,
742
+ head_dim,
743
+ block_m,
744
+ block_n,
745
+ num_warps,
746
+ True,
747
+ return_lse,
748
+ window_left,
749
+ rel_logits.stride(0),
750
+ rel_logits.stride(1),
751
+ rel_logits.stride(2),
752
+ rel_logits.shape[2],
753
+ rel_bias_qk_scale,
754
+ is_fp8,
755
+ num_warps=num_warps,
756
+ )
757
+ if return_lse:
758
+ assert lse is not None
759
+ return output, lse
760
+ return output
vllm/models/inkling/amd/ops/gluon/utils.py
new
+154
@@ -0,0 +1,154 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# Copyright (c) 2026 LightSeek Foundation
5
+#
6
+# Permission is hereby granted, free of charge, to any person obtaining a copy
7
+# of this software and associated documentation files (the "Software"), to deal
8
+# in the Software without restriction, including without limitation the rights
9
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
10
+# copies of the Software, and to permit persons to whom the Software is
11
+# furnished to do so, subject to the following conditions:
12
+#
13
+# The above copyright notice and this permission notice shall be included in
14
+# all copies or substantial portions of the Software.
15
+#
16
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22
+# SOFTWARE.
23
+
24
+from vllm.triton_utils import aggregate, gl, gluon, tl
25
+
26
+_INV_LN2_VALUE = 1.4426950408889634
27
+_INV_LN2 = tl.constexpr(_INV_LN2_VALUE)
28
+_LN2_VALUE = 0.6931471805599453
29
+_LN2 = tl.constexpr(_LN2_VALUE)
30
+_PROPAGATE_NAN_ALL = gl.constexpr(tl.PropagateNan.ALL)
31
+
32
+
33
+@gluon.jit
34
+def maximum(a, b, propagate_nan: gl.constexpr = _PROPAGATE_NAN_ALL):
35
+ return gl.maximum(a, b, propagate_nan=propagate_nan)
36
+
37
+
38
+@gluon.jit
39
+def max(input, axis=None, keep_dims=False):
40
+ return gl.reduce(input, axis, maximum, keep_dims=keep_dims)
41
+
42
+
43
+@gluon.constexpr_function
44
+def attention_layouts(head_dim, block_n, is_fp8, dtype, num_warps, instr_shape):
45
+ mfma = gl.amd.AMDMFMALayout(
46
+ version=4,
47
+ instr_shape=instr_shape,
48
+ transposed=True,
49
+ warps_per_cta=[num_warps, 1],
50
+ )
51
+ qk_layout = mfma
52
+ pv_layout = mfma
53
+ # qk_kw is derived from a 128-bit load / dtype bitwidth; pv_kw is tuned.
54
+ qk_kw = 16 if is_fp8 else 8
55
+ pv_kw = 8 if is_fp8 else 4
56
+ q_layout = gl.DotOperandLayout(0, qk_layout, k_width=qk_kw)
57
+ k_layout = gl.DotOperandLayout(1, qk_layout, k_width=qk_kw)
58
+ p_layout = gl.DotOperandLayout(0, pv_layout, k_width=pv_kw)
59
+ v_layout = gl.DotOperandLayout(1, pv_layout, k_width=pv_kw)
60
+ # load_vec = elems/lane (dtype-dependent, == qk_kw); load_threads span HEAD_DIM.
61
+ load_vec = 16 if is_fp8 else 8
62
+ load_threads = head_dim // load_vec
63
+ load_layout = gl.BlockedLayout(
64
+ [1, load_vec], [64 // load_threads, load_threads], [num_warps, 1], [1, 0]
65
+ )
66
+ # store_vec is always 16-bit (128 / 16 = 8) regardless of input dtype.
67
+ store_vec = 8
68
+ store_threads = head_dim // store_vec
69
+ store_layout = gl.BlockedLayout(
70
+ [1, store_vec], [64 // store_threads, store_threads], [num_warps, 1], [1, 0]
71
+ )
72
+ # Take only the built-in's padding, not its swizzle: the swizzle scatters
73
+ # head_dim across banks, making the LDS write stride non-constant, so it can't
74
+ # lower through async_copy's affine [1, 0] load. Padding-only is affine and
75
+ # DMA-legal.
76
+ # TODO(perf): to also use the swizzle, co-design a matched load layout so the
77
+ # DMA stays legal.
78
+ k_api = gl.amd.cdna4.compute_efficient_padded_shared_layout(
79
+ k_layout, [block_n, head_dim], dtype, is_k_contig=True
80
+ )
81
+ v_api = gl.amd.cdna4.compute_efficient_padded_shared_layout(
82
+ v_layout, [block_n, head_dim], dtype, is_k_contig=False
83
+ )
84
+ assert k_api is not None and v_api is not None, (
85
+ "no CDNA4 padded shared layout for this operand/dtype"
86
+ )
87
+ k_pairs = list(k_api.interval_padding_pairs)
88
+ v_pairs = list(v_api.interval_padding_pairs)
89
+ assert len(k_pairs) == 1 and len(v_pairs) == 1, (
90
+ "expected a single interval padding pair from the built-in"
91
+ )
92
+ k_smem_layout = gl.PaddedSharedLayout.with_identity_for(
93
+ [[int(k_pairs[0][0]), int(k_pairs[0][1])]], [block_n, head_dim], [1, 0]
94
+ )
95
+ v_smem_layout = gl.PaddedSharedLayout.with_identity_for(
96
+ [[int(v_pairs[0][0]), int(v_pairs[0][1])]], [block_n, head_dim], [1, 0]
97
+ )
98
+ return (
99
+ qk_layout,
100
+ pv_layout,
101
+ q_layout,
102
+ k_layout,
103
+ p_layout,
104
+ v_layout,
105
+ load_layout,
106
+ store_layout,
107
+ k_smem_layout,
108
+ v_smem_layout,
109
+ )
110
+
111
+
112
+@aggregate
113
+class InputStrides:
114
+ stride_t: gl.constexpr
115
+ stride_h: gl.constexpr
116
+ stride_d: gl.constexpr
117
+
118
+ @gluon.constexpr_function
119
+ def __init__(self, stride_t, stride_h, stride_d):
120
+ self.stride_t = gl.constexpr(stride_t)
121
+ self.stride_h = gl.constexpr(stride_h)
122
+ self.stride_d = gl.constexpr(stride_d)
123
+
124
+ @gluon.jit
125
+ def offsets(self, token, head, dim):
126
+ return (token * self.stride_t + head * self.stride_h + dim * self.stride_d).to(
127
+ gl.int32
128
+ )
129
+
130
+
131
+@aggregate
132
+class PagedKVStrides:
133
+ stride_b: gl.constexpr
134
+ stride_p: gl.constexpr
135
+ stride_h: gl.constexpr
136
+ stride_d: gl.constexpr
137
+
138
+ @gluon.constexpr_function
139
+ def __init__(self, stride_b, stride_p, stride_h, stride_d):
140
+ self.stride_b = gl.constexpr(stride_b)
141
+ self.stride_p = gl.constexpr(stride_p)
142
+ self.stride_h = gl.constexpr(stride_h)
143
+ self.stride_d = gl.constexpr(stride_d)
144
+
145
+ @gluon.jit
146
+ def offsets(self, page, token, head, dim):
147
+ # KV pools can exceed the 32-bit buffer-offset range. Keep the page
148
+ # arithmetic in int64, matching the original kernel's large-cache path.
149
+ return (
150
+ page.to(gl.int64) * self.stride_b
151
+ + token.to(gl.int64) * self.stride_p
152
+ + head * self.stride_h
153
+ + dim * self.stride_d
154
+ )
vllm/models/inkling/amd/ops/lamport.py
new
+766
@@ -0,0 +1,766 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Deadlock-free fused RS + short-conv + AG + residual + RMSNorm.
4
+
5
+The public integration surface is ``LamportRSConv.rs_sconv_ag_add_norm``.
6
+
7
+Liveness
8
+--------
9
+Large grids are deliberately split at the two communication dependencies:
10
+
11
+ 1. ``_publish_input_kernel`` only publishes rank partials.
12
+ 2. ``_reduce_insert_kernel`` waits, reduces, and inserts the local shard.
13
+ 3. ``_sconv_publish_kernel`` only computes and publishes the local result.
14
+ 4. ``_gather_norm_kernel`` waits, gathers, and normalizes.
15
+
16
+Without PDL, CUDA stream order completes a producer before its consumer. With
17
+PDL, every producer CTA posts all peer stores before triggering its dependent,
18
+and the consumer executes ``gdc_wait`` before polling. A consumer therefore
19
+waits only for stores from a producer that is already running (or complete) on
20
+another GPU. It never waits for another block in its own grid. Consequently
21
+spinning consumers cannot occupy resources needed by any producer, and the
22
+wait-for graph has no cycle. The proof is independent of grid size, block
23
+dispatch order, and occupancy.
24
+
25
+For one token, the first three phases use eight independent channel slices to
26
+expose enough CTA parallelism for decode latency. Each slice has exclusive
27
+ownership of its cache and Lamport columns; the gather/RMSNorm phase retains
28
+one CTA per token so no cross-CTA reduction or completion counter is needed.
29
+
30
+Immediate buffer reuse (including replay of a captured CUDA graph) is also
31
+safe. Rank R cannot republish input for call n+1 until its gather for n has
32
+finished; that gather waited for owner O's output, which O publishes only
33
+after consuming R's input for n. Likewise, R cannot republish output for n+1
34
+until its reduction for n+1 has observed destination D's input; D publishes
35
+that input only after its gather consumed R's output for n. Thus every prior
36
+read happens-before a same-slot rewrite. Three generations reduce incidental
37
+coupling for ordinary launches, but correctness does not depend on rotation.
38
+
39
+The payload itself is the Lamport flag. A 32-bit store publishes two bf16s
40
+atomically; 0x80008000 (two negative zeroes) denotes an empty pair. Real
41
+negative zeroes are changed to positive zero before publication. Consumers
42
+use volatile 32-bit loads and restore the sentinel after consuming a slot.
43
+"""
44
+
45
+from __future__ import annotations
46
+
47
+import os
48
+
49
+import torch
50
+
51
+from vllm.distributed import get_tp_group
52
+from vllm.distributed.parallel_state import in_the_same_node_as
53
+from vllm.logger import init_logger
54
+from vllm.triton_utils import tl, triton
55
+
56
+logger = init_logger(__name__)
57
+
58
+_MAX_TOKENS = 16384
59
+_EMPTY_PAIR = tl.constexpr(0x80008000)
60
+
61
+
62
+@triton.jit
63
+def _pack_bf16_pairs(values):
64
+ """Pack bf16 values into atomic u32 pairs and reserve negative zero."""
65
+ lo, hi = tl.split(values.reshape([values.shape[0] // 2, 2]))
66
+ lo = lo.to(tl.uint16, bitcast=True)
67
+ hi = hi.to(tl.uint16, bitcast=True)
68
+ lo = tl.where(lo == 0x8000, 0, lo).to(tl.uint32)
69
+ hi = tl.where(hi == 0x8000, 0, hi).to(tl.uint32)
70
+ return lo | (hi << 16)
71
+
72
+
73
+@triton.jit
74
+def _unpack_bf16_pairs(values):
75
+ lo = (values & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True)
76
+ hi = (values >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True)
77
+ return tl.interleave(lo, hi)
78
+
79
+
80
+@triton.jit
81
+def _wait_pairs(ptr, offsets, mask):
82
+ values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True)
83
+ while tl.max(tl.where(mask & (values == _EMPTY_PAIR), 1, 0)) != 0:
84
+ values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True)
85
+ return values
86
+
87
+
88
+@triton.jit
89
+def _publish_input_kernel(
90
+ stage_ptr,
91
+ peer_ptrs,
92
+ peer_offset_u32,
93
+ stride_stage_t,
94
+ C: tl.constexpr,
95
+ CS: tl.constexpr,
96
+ CS_P2: tl.constexpr,
97
+ SPLITS: tl.constexpr,
98
+ RANK: tl.constexpr,
99
+ WORLD: tl.constexpr,
100
+ USE_PDL: tl.constexpr,
101
+ launch_pdl: tl.constexpr,
102
+):
103
+ """Publish this rank's full partial row into every shard owner's slots."""
104
+ token = tl.program_id(0).to(tl.int64)
105
+ split = tl.program_id(1).to(tl.int64)
106
+ CSS: tl.constexpr = CS // SPLITS
107
+ pair = tl.arange(0, CS_P2 // 2)
108
+ pair_mask = pair < CSS // 2
109
+ elem = tl.arange(0, CS_P2)
110
+ elem_mask = elem < CSS
111
+ ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64))
112
+ # Address generation is independent of the preceding stream kernel. The
113
+ # acquire remains before stage/peer loads, which may consume its writes.
114
+ if USE_PDL:
115
+ tl.extra.cuda.gdc_wait()
116
+
117
+ for owner in tl.static_range(WORLD):
118
+ values = tl.load(
119
+ stage_ptr + token * stride_stage_t + owner * CS + split * CSS + elem,
120
+ mask=elem_mask,
121
+ other=0.0,
122
+ )
123
+ packed = _pack_bf16_pairs(values)
124
+ base = tl.load(ptrs + owner).to(tl.pointer_type(tl.uint32))
125
+ dst = (token * WORLD + RANK) * (CS // 2) + split * (CSS // 2) + pair
126
+ tl.store(base + peer_offset_u32 + dst, packed, mask=pair_mask)
127
+ if USE_PDL:
128
+ tl.extra.cuda.gdc_launch_dependents()
129
+
130
+
131
+@triton.jit
132
+def _reduce_insert_kernel(
133
+ input_peer_ptrs,
134
+ input_peer_offset_u32,
135
+ cache_ptr,
136
+ slot_ptr,
137
+ stride_cache_block,
138
+ stride_cache_head,
139
+ stride_cache_token,
140
+ stride_cache_dim,
141
+ block_size,
142
+ C: tl.constexpr,
143
+ CS: tl.constexpr,
144
+ CS_P2: tl.constexpr,
145
+ SPLITS: tl.constexpr,
146
+ HEAD_SIZE: tl.constexpr,
147
+ CACHE_OFFSET: tl.constexpr,
148
+ RANK: tl.constexpr,
149
+ WORLD: tl.constexpr,
150
+ USE_PDL: tl.constexpr,
151
+ launch_pdl: tl.constexpr,
152
+):
153
+ """Consume all partials for this rank and insert the reduced cache row."""
154
+ token = tl.program_id(0).to(tl.int64)
155
+ split = tl.program_id(1).to(tl.int64)
156
+ CSS: tl.constexpr = CS // SPLITS
157
+ source = tl.arange(0, WORLD)
158
+ pair = tl.arange(0, CS_P2 // 2)
159
+ pair_mask = pair < CSS // 2
160
+ offsets = (
161
+ (token * WORLD + source)[:, None] * (CS // 2)
162
+ + split * (CSS // 2)
163
+ + pair[None, :]
164
+ )
165
+ mask = tl.full([WORLD], True, tl.int1)[:, None] & pair_mask[None, :]
166
+ input_ptrs = input_peer_ptrs.to(tl.pointer_type(tl.uint64))
167
+ input_u32 = tl.load(input_ptrs + RANK).to(tl.pointer_type(tl.uint32))
168
+ input_u32 += input_peer_offset_u32
169
+ # Slot metadata and the cache destination do not depend on input publish.
170
+ slot = tl.load(slot_ptr + token)
171
+ valid = slot >= 0
172
+ channel = tl.arange(0, CS_P2)
173
+ channel_mask = channel < CSS
174
+ global_channel = split * CSS + channel
175
+ head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1)
176
+ dim = CACHE_OFFSET + global_channel % HEAD_SIZE
177
+ safe_slot = tl.maximum(slot, 0).to(tl.int64)
178
+ dst = (
179
+ cache_ptr
180
+ + (safe_slot // block_size) * stride_cache_block
181
+ + head * stride_cache_head
182
+ + (safe_slot % block_size) * stride_cache_token
183
+ + dim * stride_cache_dim
184
+ )
185
+ if USE_PDL:
186
+ tl.extra.cuda.gdc_wait()
187
+ packed = _wait_pairs(input_u32, offsets, mask)
188
+
189
+ lo = (packed & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True)
190
+ hi = (packed >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True)
191
+ reduced = tl.interleave(
192
+ tl.sum(lo.to(tl.float32), axis=0).to(tl.bfloat16),
193
+ tl.sum(hi.to(tl.float32), axis=0).to(tl.bfloat16),
194
+ )
195
+ tl.store(
196
+ input_u32 + offsets,
197
+ tl.full([WORLD, CS_P2 // 2], _EMPTY_PAIR, tl.uint32),
198
+ mask=mask,
199
+ )
200
+
201
+ tl.store(dst, reduced, mask=valid & channel_mask)
202
+ if USE_PDL:
203
+ tl.extra.cuda.gdc_launch_dependents()
204
+
205
+
206
+@triton.jit
207
+def _sconv_publish_kernel(
208
+ peer_ptrs,
209
+ peer_offset_u32,
210
+ residual_ptr,
211
+ weight_ptr,
212
+ cache_ptr,
213
+ position_ptr,
214
+ sequence_ptr,
215
+ slot_ptr,
216
+ block_table_ptr,
217
+ stride_residual_t,
218
+ stride_cache_block,
219
+ stride_cache_head,
220
+ stride_cache_token,
221
+ stride_cache_dim,
222
+ stride_block_table_r,
223
+ max_blocks,
224
+ block_size,
225
+ C: tl.constexpr,
226
+ CS: tl.constexpr,
227
+ CS_P2: tl.constexpr,
228
+ SPLITS: tl.constexpr,
229
+ HEAD_SIZE: tl.constexpr,
230
+ CACHE_OFFSET: tl.constexpr,
231
+ RANK: tl.constexpr,
232
+ WORLD: tl.constexpr,
233
+ WINDOW: tl.constexpr,
234
+ USE_PDL: tl.constexpr,
235
+ launch_pdl: tl.constexpr,
236
+):
237
+ """Compute this rank's shard, then publish it to every rank."""
238
+ tl.static_assert(WINDOW == 4)
239
+ token = tl.program_id(0).to(tl.int64)
240
+ split = tl.program_id(1).to(tl.int64)
241
+ CSS: tl.constexpr = CS // SPLITS
242
+ channel = tl.arange(0, CS_P2)
243
+ channel_mask = channel < CSS
244
+ global_channel = split * CSS + channel
245
+ head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1)
246
+ dim = CACHE_OFFSET + global_channel % HEAD_SIZE
247
+ slot = tl.load(slot_ptr + token)
248
+ valid = slot >= 0
249
+ position = tl.load(position_ptr + token)
250
+ sequence = tl.load(sequence_ptr + token)
251
+ safe_slot = tl.maximum(slot, 0).to(tl.int64)
252
+ own_ptr = (
253
+ cache_ptr
254
+ + (safe_slot // block_size) * stride_cache_block
255
+ + head * stride_cache_head
256
+ + (safe_slot % block_size) * stride_cache_token
257
+ + dim * stride_cache_dim
258
+ )
259
+ # These loads were made visible before reduce/insert could signal us and
260
+ # are independent of its cache store. Hoisting them gives PDL useful work
261
+ # to overlap while retaining the acquire before every cache read.
262
+ residual = tl.load(
263
+ residual_ptr + token * stride_residual_t + RANK * CS + global_channel,
264
+ mask=channel_mask,
265
+ other=0.0,
266
+ )
267
+ weight0 = tl.load(
268
+ weight_ptr + global_channel * WINDOW,
269
+ mask=channel_mask,
270
+ other=0.0,
271
+ ).to(tl.float32)
272
+ weight1 = tl.load(
273
+ weight_ptr + global_channel * WINDOW + 1,
274
+ mask=channel_mask,
275
+ other=0.0,
276
+ ).to(tl.float32)
277
+ weight2 = tl.load(
278
+ weight_ptr + global_channel * WINDOW + 2,
279
+ mask=channel_mask,
280
+ other=0.0,
281
+ ).to(tl.float32)
282
+ weight3 = tl.load(
283
+ weight_ptr + global_channel * WINDOW + 3,
284
+ mask=channel_mask,
285
+ other=0.0,
286
+ ).to(tl.float32)
287
+ ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64))
288
+ if USE_PDL:
289
+ tl.extra.cuda.gdc_wait()
290
+ current = tl.load(own_ptr, mask=valid & channel_mask, other=0.0)
291
+
292
+ conv = tl.zeros([CS_P2], tl.float32)
293
+ for tap_idx in tl.static_range(WINDOW):
294
+ source_position = position - (WINDOW - 1) + tap_idx
295
+ take = valid & (source_position >= 0)
296
+ if tap_idx == WINDOW - 1:
297
+ value = tl.where(take, current.to(tl.float32), 0.0)
298
+ else:
299
+ safe_position = tl.maximum(source_position, 0)
300
+ logical_block = tl.minimum(safe_position // block_size, max_blocks - 1)
301
+ physical_block = tl.load(
302
+ block_table_ptr + sequence * stride_block_table_r + logical_block,
303
+ mask=take,
304
+ other=0,
305
+ ).to(tl.int64)
306
+ source_ptr = (
307
+ cache_ptr
308
+ + physical_block * stride_cache_block
309
+ + head * stride_cache_head
310
+ + (safe_position % block_size) * stride_cache_token
311
+ + dim * stride_cache_dim
312
+ )
313
+ cached = tl.load(source_ptr, mask=take & channel_mask, other=0.0)
314
+ value = tl.where(take, cached.to(tl.float32), 0.0)
315
+ if tap_idx == 0:
316
+ weight = weight0
317
+ elif tap_idx == 1:
318
+ weight = weight1
319
+ elif tap_idx == 2:
320
+ weight = weight2
321
+ else:
322
+ weight = weight3
323
+ conv += value * weight
324
+
325
+ # Preserve both bf16 rounding points of the original sublayer.
326
+ short_conv_with_skip = (conv + current.to(tl.float32)).to(tl.bfloat16)
327
+ output = (residual.to(tl.float32) + short_conv_with_skip.to(tl.float32)).to(
328
+ tl.bfloat16
329
+ )
330
+
331
+ pair = tl.arange(0, CS_P2 // 2)
332
+ pair_mask = pair < CSS // 2
333
+ packed = _pack_bf16_pairs(output)
334
+ row_offset = token * (C // 2) + RANK * (CS // 2) + split * (CSS // 2) + pair
335
+
336
+ for destination in tl.static_range(WORLD):
337
+ base = tl.load(ptrs + destination).to(tl.pointer_type(tl.uint32))
338
+ tl.store(base + peer_offset_u32 + row_offset, packed, mask=pair_mask)
339
+ if USE_PDL:
340
+ tl.extra.cuda.gdc_launch_dependents()
341
+
342
+
343
+@triton.jit
344
+def _gather_norm_kernel(
345
+ output_peer_ptrs,
346
+ output_peer_offset_u32,
347
+ norm_weight_ptr,
348
+ normed_ptr,
349
+ residual_out_ptr,
350
+ eps,
351
+ stride_output_t,
352
+ C: tl.constexpr,
353
+ C_P2: tl.constexpr,
354
+ RANK: tl.constexpr,
355
+ HAS_NORM: tl.constexpr,
356
+ USE_PDL: tl.constexpr,
357
+ launch_pdl: tl.constexpr,
358
+):
359
+ """Consume a complete row; one CTA owns all outputs for one token."""
360
+ token = tl.program_id(0).to(tl.int64)
361
+ pair = tl.arange(0, C_P2 // 2)
362
+ pair_mask = pair < C // 2
363
+ output_ptrs = output_peer_ptrs.to(tl.pointer_type(tl.uint64))
364
+ output_u32 = tl.load(output_ptrs + RANK).to(tl.pointer_type(tl.uint32))
365
+ output_u32 += output_peer_offset_u32
366
+ offsets = token * (C // 2) + pair
367
+ channel = tl.arange(0, C_P2)
368
+ channel_mask = channel < C
369
+ # Gamma is independent of the preceding sconv publication. Keep the
370
+ # acquire immediately before polling the Lamport output slots.
371
+ if HAS_NORM:
372
+ weight = tl.load(norm_weight_ptr + channel, mask=channel_mask, other=0.0)
373
+ if USE_PDL:
374
+ tl.extra.cuda.gdc_wait()
375
+ packed = _wait_pairs(output_u32, offsets, pair_mask)
376
+ row = _unpack_bf16_pairs(packed)
377
+ tl.store(
378
+ residual_out_ptr + token * stride_output_t + channel, row, mask=channel_mask
379
+ )
380
+ if HAS_NORM:
381
+ row_f32 = tl.where(channel_mask, row.to(tl.float32), 0.0)
382
+ inv_rms = tl.rsqrt(tl.sum(row_f32 * row_f32, axis=0) / C + eps)
383
+ tl.store(
384
+ normed_ptr + token * stride_output_t + channel,
385
+ (row_f32 * inv_rms * weight.to(tl.float32)).to(tl.bfloat16),
386
+ mask=channel_mask,
387
+ )
388
+ tl.store(
389
+ output_u32 + offsets,
390
+ tl.full([C_P2 // 2], _EMPTY_PAIR, tl.uint32),
391
+ mask=pair_mask,
392
+ )
393
+ if USE_PDL:
394
+ tl.extra.cuda.gdc_launch_dependents()
395
+
396
+
397
+@triton.jit
398
+def _validate_lamport_init_kernel(
399
+ input_peer_ptrs,
400
+ output_peer_ptrs,
401
+ bad_ptr,
402
+ num_pairs,
403
+ RANK: tl.constexpr,
404
+ BLOCK: tl.constexpr,
405
+):
406
+ """Validate both complete local allocations through their fabric pointers."""
407
+ offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
408
+ mask = offsets < num_pairs
409
+ ptrs_in = input_peer_ptrs.to(tl.pointer_type(tl.uint64))
410
+ ptrs_out = output_peer_ptrs.to(tl.pointer_type(tl.uint64))
411
+ local_in = tl.load(ptrs_in + RANK).to(tl.pointer_type(tl.uint32))
412
+ local_out = tl.load(ptrs_out + RANK).to(tl.pointer_type(tl.uint32))
413
+ value_in = tl.load(local_in + offsets, mask=mask, other=_EMPTY_PAIR)
414
+ value_out = tl.load(local_out + offsets, mask=mask, other=_EMPTY_PAIR)
415
+ bad = tl.max(
416
+ tl.where(mask & ((value_in != _EMPTY_PAIR) | (value_out != _EMPTY_PAIR)), 1, 0)
417
+ )
418
+ if bad != 0:
419
+ tl.atomic_max(bad_ptr, 1)
420
+
421
+
422
+class LamportRSConv:
423
+ """Persistent symmetric buffers for one TP group."""
424
+
425
+ def _initialize_lamport_buffers(self) -> None:
426
+ """Arm every slot and collectively verify the exact sentinel bits.
427
+
428
+ Do not replace this with a zero-fill: Lamport readiness distinguishes
429
+ the bf16 bit pattern for -0.0 (0x8000) from every published payload.
430
+ The validation is deliberately collective so one rank cannot enter the
431
+ first polling kernel while another rank still has an unarmed buffer.
432
+ """
433
+ if not self.buf_in.is_contiguous() or not self.buf_out.is_contiguous():
434
+ raise RuntimeError("Lamport symmetric buffers must be contiguous")
435
+ if self.buf_in.numel() % 2 or self.buf_out.numel() % 2:
436
+ raise RuntimeError("Lamport buffers must contain whole uint32 pairs")
437
+
438
+ # Fill through int16 so each bf16 lane receives the exact -0.0 bits.
439
+ # This covers all three generations and all max-token slots, including
440
+ # slots that a smaller first invocation does not touch.
441
+ self.buf_in.view(torch.int16).fill_(-0x8000)
442
+ self.buf_out.view(torch.int16).fill_(-0x8000)
443
+ torch.accelerator.synchronize(self.device)
444
+ self.tp.barrier()
445
+
446
+ # Scan the complete local allocations. Since every rank performs the
447
+ # scan and participates in MAX, success proves every symmetric backing
448
+ # allocation was armed before any rank is allowed to use generation 0.
449
+ bad = torch.logical_or(
450
+ self.buf_in.view(torch.int16).ne(-0x8000).any(),
451
+ self.buf_out.view(torch.int16).ne(-0x8000).any(),
452
+ ).to(dtype=torch.float32)
453
+ bad = self.tp.all_reduce(bad)
454
+ if int(bad.item()) != 0:
455
+ raise RuntimeError("Lamport sentinel initialization failed on a TP rank")
456
+ self.tp.barrier()
457
+
458
+ def _initialize_mnnvl_buffers(self) -> None:
459
+ """Initialize and validate FlashInfer fabric-mapped Lamport storage."""
460
+ self._mnnvl_input_handle.lamport_initialize(self.rank, torch.bfloat16)
461
+ self._mnnvl_output_handle.lamport_initialize(self.rank, torch.bfloat16)
462
+ torch.accelerator.synchronize(self.device)
463
+ self.tp.barrier()
464
+
465
+ bad = torch.zeros((), dtype=torch.int32, device=self.device)
466
+ num_pairs = self.num_buffers * self.max_tokens * self.hidden_size // 2
467
+ _validate_lamport_init_kernel[(triton.cdiv(num_pairs, 256),)](
468
+ self.input_peer_ptrs,
469
+ self.output_peer_ptrs,
470
+ bad,
471
+ num_pairs,
472
+ RANK=self.rank,
473
+ BLOCK=256,
474
+ num_warps=4,
475
+ )
476
+ bad = self.tp.all_reduce(bad.to(torch.float32))
477
+ if int(bad.item()) != 0:
478
+ raise RuntimeError("MNNVL Lamport sentinel initialization failed")
479
+ self.tp.barrier()
480
+
481
+ def __init__(
482
+ self, hidden_size: int, window_size: int, max_tokens: int = _MAX_TOKENS
483
+ ) -> None:
484
+ import torch.distributed._symmetric_memory as symm_mem
485
+
486
+ tp = get_tp_group()
487
+ self.tp = tp
488
+ self.group = tp.device_group
489
+ self.world_size = tp.world_size
490
+ self.rank = tp.rank_in_group
491
+ self.device = torch.device(tp.device)
492
+ if self.world_size not in (2, 4, 8):
493
+ raise ValueError(f"TP world size must be 2, 4, or 8, got {self.world_size}")
494
+ if hidden_size % (2 * self.world_size) != 0:
495
+ raise ValueError("hidden size must produce an even shard on every rank")
496
+ if window_size != 4:
497
+ raise ValueError(f"short-conv window size must be 4, got {window_size}")
498
+ if max_tokens < 1 or max_tokens > _MAX_TOKENS:
499
+ raise ValueError(f"max_tokens must be in [1, {_MAX_TOKENS}]")
500
+
501
+ is_cross_node = not all(in_the_same_node_as(tp.cpu_group))
502
+ self.hidden_size = hidden_size
503
+ self.window_size = window_size
504
+ self.max_tokens = max_tokens
505
+ self.shard_size = hidden_size // self.world_size
506
+ # Three generations follow FlashInfer's Lamport layout. A generation
507
+ # is reused only after two intervening collective calls.
508
+ self.num_buffers = 3
509
+ self.input_generation_bytes = max_tokens * hidden_size * 2
510
+ self.output_generation_bytes = max_tokens * hidden_size * 2
511
+ self.use_pdl = torch.cuda.get_device_capability(self.device)[0] >= 9
512
+ if is_cross_node:
513
+ try:
514
+ from flashinfer.comm.mnnvl import (
515
+ McastGPUBuffer,
516
+ TorchDistBackend,
517
+ is_mnnvl_fabric_supported,
518
+ )
519
+ except ImportError as error:
520
+ raise RuntimeError(
521
+ "cross-node TP requires FlashInfer MNNVL support"
522
+ ) from error
523
+
524
+ local_supported = int(
525
+ is_mnnvl_fabric_supported(torch.accelerator.current_device_index())
526
+ )
527
+ unsupported = torch.tensor(
528
+ 1 - local_supported, dtype=torch.float32, device=self.device
529
+ )
530
+ unsupported = tp.all_reduce(unsupported)
531
+ if int(unsupported.item()) != 0:
532
+ raise RuntimeError("cross-node TP is supported only on MNNVL fabric")
533
+
534
+ comm_backend = TorchDistBackend(self.group)
535
+ allocation_bytes = self.num_buffers * max_tokens * hidden_size * 2
536
+ self._mnnvl_input_handle = McastGPUBuffer(
537
+ allocation_bytes,
538
+ self.world_size,
539
+ self.rank,
540
+ self.device,
541
+ comm_backend,
542
+ )
543
+ self._mnnvl_output_handle = McastGPUBuffer(
544
+ allocation_bytes,
545
+ self.world_size,
546
+ self.rank,
547
+ self.device,
548
+ comm_backend,
549
+ )
550
+ self.input_peer_ptrs = self._mnnvl_input_handle.get_buffer_ptrs_dev()
551
+ self.output_peer_ptrs = self._mnnvl_output_handle.get_buffer_ptrs_dev()
552
+ self._initialize_mnnvl_buffers()
553
+ logger.info("using FlashInfer fabric-mapped MNNVL Lamport buffers")
554
+ else:
555
+ self.buf_in = symm_mem.empty(
556
+ self.num_buffers,
557
+ max_tokens,
558
+ self.world_size,
559
+ self.shard_size,
560
+ dtype=torch.bfloat16,
561
+ device=self.device,
562
+ )
563
+ self.buf_out = symm_mem.empty(
564
+ self.num_buffers,
565
+ max_tokens,
566
+ hidden_size,
567
+ dtype=torch.bfloat16,
568
+ device=self.device,
569
+ )
570
+ group_name = self.group.group_name
571
+ input_handle = symm_mem.rendezvous(self.buf_in, group_name)
572
+ output_handle = symm_mem.rendezvous(self.buf_out, group_name)
573
+ self.input_peer_ptrs = input_handle.buffer_ptrs_dev
574
+ self.output_peer_ptrs = output_handle.buffer_ptrs_dev
575
+ self._input_handle = input_handle
576
+ self._output_handle = output_handle
577
+ self._initialize_lamport_buffers()
578
+ self.generation = 0
579
+
580
+ def usable(self, num_tokens: int) -> bool:
581
+ return 0 < num_tokens <= self.max_tokens
582
+
583
+ def rs_sconv_ag_add_norm(
584
+ self,
585
+ input_tensor: torch.Tensor,
586
+ residual: torch.Tensor,
587
+ conv_weight: torch.Tensor,
588
+ norm_weight: torch.Tensor | None,
589
+ eps: float,
590
+ cache: torch.Tensor,
591
+ positions: torch.Tensor,
592
+ block_table: torch.Tensor,
593
+ seq_idx: torch.Tensor,
594
+ slot_mapping: torch.Tensor,
595
+ off_s: int,
596
+ ws: int,
597
+ block_size: int,
598
+ ) -> tuple[torch.Tensor | None, torch.Tensor]:
599
+ """Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``."""
600
+ tokens, hidden_size = residual.shape
601
+ if not self.usable(tokens):
602
+ raise ValueError(f"num_tokens must be in [1, {self.max_tokens}]")
603
+ if hidden_size != self.hidden_size or residual.dtype != torch.bfloat16:
604
+ raise ValueError("residual must be bf16 [T, 6144]")
605
+ if (
606
+ input_tensor.shape != residual.shape
607
+ or input_tensor.dtype != torch.bfloat16
608
+ or input_tensor.stride(1) != 1
609
+ ):
610
+ raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]")
611
+ shard_size = hidden_size // self.world_size
612
+ if conv_weight.shape != (shard_size, self.window_size):
613
+ raise ValueError(
614
+ f"conv_weight must have shape [{shard_size}, {self.window_size}]"
615
+ )
616
+ if (
617
+ conv_weight.dtype != torch.bfloat16
618
+ or conv_weight.stride(0) != self.window_size
619
+ ):
620
+ raise ValueError("conv_weight must be contiguous bf16")
621
+ if norm_weight is not None and (
622
+ norm_weight.shape != (hidden_size,) or norm_weight.dtype != torch.bfloat16
623
+ ):
624
+ raise ValueError("norm_weight must be bf16 [6144] or None")
625
+ if cache.dtype != torch.bfloat16 or cache.ndim != 4:
626
+ raise ValueError("cache must be a 4-D bf16 tensor")
627
+ if shard_size % ws != 0 or cache.shape[1] != shard_size // ws:
628
+ raise ValueError("cache head layout is inconsistent with ws")
629
+ if off_s < 0 or off_s + ws > cache.shape[3]:
630
+ raise ValueError("cache channel offset is out of bounds")
631
+
632
+ index = self.generation
633
+ input_offset = index * self.input_generation_bytes // 4
634
+ output_offset = index * self.output_generation_bytes // 4
635
+ normed = torch.empty_like(residual) if norm_weight is not None else None
636
+ residual_out = torch.empty_like(residual)
637
+ phase_splits = 8 if tokens == 1 and shard_size % 8 == 0 else 1
638
+ phase_tile_p2 = triton.next_power_of_2(shard_size // phase_splits)
639
+ phase_grid = (tokens, phase_splits)
640
+ # Wide CTAs win before the grid saturates; smaller CTAs reduce register
641
+ # pressure once high-throughput batches provide enough parallelism.
642
+ if 128 <= tokens <= 2048:
643
+ phase_warps = 16
644
+ elif tokens > 2048:
645
+ phase_warps = 8
646
+ else:
647
+ phase_warps = 4
648
+ gather_warps = 4 if tokens >= 256 else 8
649
+ _publish_input_kernel[phase_grid](
650
+ input_tensor,
651
+ self.input_peer_ptrs,
652
+ input_offset,
653
+ input_tensor.stride(0),
654
+ C=hidden_size,
655
+ CS=shard_size,
656
+ CS_P2=phase_tile_p2,
657
+ SPLITS=phase_splits,
658
+ RANK=self.rank,
659
+ WORLD=self.world_size,
660
+ USE_PDL=self.use_pdl,
661
+ launch_pdl=self.use_pdl,
662
+ num_warps=phase_warps,
663
+ )
664
+ _reduce_insert_kernel[phase_grid](
665
+ self.input_peer_ptrs,
666
+ input_offset,
667
+ cache,
668
+ slot_mapping,
669
+ cache.stride(0),
670
+ cache.stride(1),
671
+ cache.stride(2),
672
+ cache.stride(3),
673
+ block_size,
674
+ C=hidden_size,
675
+ CS=shard_size,
676
+ CS_P2=phase_tile_p2,
677
+ SPLITS=phase_splits,
678
+ HEAD_SIZE=ws,
679
+ CACHE_OFFSET=off_s,
680
+ RANK=self.rank,
681
+ WORLD=self.world_size,
682
+ USE_PDL=self.use_pdl,
683
+ launch_pdl=self.use_pdl,
684
+ num_warps=phase_warps,
685
+ )
686
+ _sconv_publish_kernel[phase_grid](
687
+ self.output_peer_ptrs,
688
+ output_offset,
689
+ residual,
690
+ conv_weight,
691
+ cache,
692
+ positions,
693
+ seq_idx,
694
+ slot_mapping,
695
+ block_table,
696
+ residual.stride(0),
697
+ cache.stride(0),
698
+ cache.stride(1),
699
+ cache.stride(2),
700
+ cache.stride(3),
701
+ block_table.stride(0),
702
+ block_table.shape[1],
703
+ block_size,
704
+ C=hidden_size,
705
+ CS=shard_size,
706
+ CS_P2=phase_tile_p2,
707
+ SPLITS=phase_splits,
708
+ HEAD_SIZE=ws,
709
+ CACHE_OFFSET=off_s,
710
+ RANK=self.rank,
711
+ WORLD=self.world_size,
712
+ WINDOW=self.window_size,
713
+ USE_PDL=self.use_pdl,
714
+ launch_pdl=self.use_pdl,
715
+ num_warps=phase_warps,
716
+ )
717
+ _gather_norm_kernel[(tokens,)](
718
+ self.output_peer_ptrs,
719
+ output_offset,
720
+ norm_weight if norm_weight is not None else residual,
721
+ normed if normed is not None else residual_out,
722
+ residual_out,
723
+ eps,
724
+ residual.stride(0),
725
+ C=hidden_size,
726
+ C_P2=triton.next_power_of_2(hidden_size),
727
+ RANK=self.rank,
728
+ HAS_NORM=norm_weight is not None,
729
+ USE_PDL=self.use_pdl,
730
+ launch_pdl=self.use_pdl,
731
+ num_warps=gather_warps,
732
+ )
733
+ self.generation = (index + 1) % self.num_buffers
734
+ return normed, residual_out
735
+
736
+
737
+_STATE: LamportRSConv | None = None
738
+_STATE_FAILED = False
739
+
740
+
741
+def initialize_lamport_rs_conv(
742
+ hidden_size: int, window_size: int, max_num_batched_tokens: int
743
+) -> None:
744
+ """Collectively initialize the TP-group state during model construction."""
745
+ global _STATE, _STATE_FAILED
746
+ if _STATE is not None:
747
+ if _STATE.hidden_size != hidden_size or _STATE.window_size != window_size:
748
+ raise RuntimeError("all Lamport users must share hidden and window sizes")
749
+ return
750
+ if _STATE_FAILED or os.environ.get("LAMPORT_RS_SCONV", "1") == "0":
751
+ return
752
+ try:
753
+ max_tokens = min(_MAX_TOKENS, max_num_batched_tokens)
754
+ _STATE = LamportRSConv(hidden_size, window_size, max_tokens=max_tokens)
755
+ except Exception:
756
+ _STATE_FAILED = True
757
+ logger.exception("fused collective unavailable; use the NCCL fallback")
758
+
759
+
760
+def get_lamport_rs_conv(hidden_size: int, window_size: int) -> LamportRSConv | None:
761
+ """Return the state initialized with the model, or ``None`` for fallback."""
762
+ if _STATE is not None and (
763
+ _STATE.hidden_size != hidden_size or _STATE.window_size != window_size
764
+ ):
765
+ raise RuntimeError("all Lamport users must share hidden and window sizes")
766
+ return _STATE
vllm/models/inkling/amd/ops/mm_towers.py
new
+190
@@ -0,0 +1,190 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Fused CUDA kernels for the Inkling vision/audio towers.
4
+
5
+Both kernels keep the reference paths' fp32 accumulation and per-op bf16
6
+rounding points (native ``rms_norm`` / ``F.gelu``); outputs are frequently
7
+bit-identical and otherwise differ by 1-2 bf16 ulps from reduction-order
8
+(real-checkpoint-weight cosine vs reference > 0.9999998).
9
+"""
10
+
11
+from __future__ import annotations
12
+
13
+import torch
14
+
15
+from vllm.triton_utils import tl, tldevice, triton
16
+
17
+from .norm import _get_num_warps_from_block_size
18
+
19
+
20
+@triton.jit
21
+def _dmel_embed_sum_norm_kernel(
22
+ idx_ptr, # [T, NB] int32 dMel bin indices (values in [0, VOCAB))
23
+ w_ptr, # [NB * VOCAB, D] bf16 embedding table
24
+ norm_w_ptr, # [D] (unused if HAS_NORM=False)
25
+ out_ptr, # [T, D] bf16
26
+ eps,
27
+ D,
28
+ stride_idx_t,
29
+ NB: tl.constexpr,
30
+ VOCAB: tl.constexpr,
31
+ D_P2: tl.constexpr,
32
+ HAS_NORM: tl.constexpr,
33
+):
34
+ t = tl.program_id(0).to(tl.int64)
35
+ offs = tl.arange(0, D_P2)
36
+ mask = offs < D
37
+ # One embedding row per mel bin (bin b uses table rows [b*VOCAB, (b+1)*VOCAB)),
38
+ # summed in fp32 (matches torch's fp32-accumulated bf16 .sum()).
39
+ acc = tl.zeros([D_P2], dtype=tl.float32)
40
+ for b in tl.static_range(NB):
41
+ v = tl.load(idx_ptr + t * stride_idx_t + b)
42
+ row = (b * VOCAB + v).to(tl.int64)
43
+ acc += tl.load(w_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32)
44
+ h = acc.to(tl.bfloat16)
45
+ if HAS_NORM:
46
+ # Match ir.ops.rms_norm: fp32 variance/normalize, then a single-rounded
47
+ # bf16 multiply with the bf16 weight.
48
+ x32 = h.to(tl.float32)
49
+ var = tl.sum(tl.where(mask, x32 * x32, 0.0), axis=0) / D
50
+ xn = (x32 * tl.math.rsqrt(var + eps)).to(tl.bfloat16)
51
+ w = tl.load(norm_w_ptr + offs, mask=mask, other=0.0)
52
+ h = xn * w
53
+ tl.store(out_ptr + t * D + offs, h, mask=mask)
54
+
55
+
56
+def dmel_embed_sum_norm(
57
+ dmel_idx: torch.Tensor, # [T, NB] int32
58
+ weight: torch.Tensor, # [NB * VOCAB, D] bf16
59
+ norm_weight: torch.Tensor | None,
60
+ eps: float,
61
+) -> torch.Tensor:
62
+ """``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no
63
+ [T, NB, D] intermediate)."""
64
+ T, nb = dmel_idx.shape
65
+ D = weight.shape[1]
66
+ assert weight.shape[0] % nb == 0
67
+ vocab = weight.shape[0] // nb
68
+ out = torch.empty((T, D), dtype=weight.dtype, device=weight.device)
69
+ if T == 0:
70
+ return out
71
+ d_p2 = triton.next_power_of_2(D)
72
+ _dmel_embed_sum_norm_kernel[(T,)](
73
+ dmel_idx,
74
+ weight,
75
+ norm_weight if norm_weight is not None else weight,
76
+ out,
77
+ eps,
78
+ D,
79
+ dmel_idx.stride(0),
80
+ NB=nb,
81
+ VOCAB=vocab,
82
+ D_P2=d_p2,
83
+ HAS_NORM=norm_weight is not None,
84
+ # Swept on GB200: 8 warps beats the block-size heuristic's 16 by ~4%
85
+ # at large T (the 80-row gather chain is latency- not lane-bound).
86
+ num_warps=8,
87
+ )
88
+ return out
89
+
90
+
91
+@triton.jit
92
+def _rmsnorm_gelu_kernel(
93
+ x_ptr, # [R, D] bf16
94
+ w_ptr, # [D]
95
+ out_ptr, # [R, D] bf16 (or the folded layout when FOLD)
96
+ eps,
97
+ R,
98
+ D,
99
+ D_P2: tl.constexpr,
100
+ BLOCK_M: tl.constexpr, # rows per block (>1 for small D)
101
+ HAS_GELU: tl.constexpr,
102
+ FOLD: tl.constexpr,
103
+ # fold geometry: input rows index [N, T, H, W]; the store scatters each
104
+ # row to (out_row, slot) of fold_timespace_to_depth's output layout.
105
+ FT: tl.constexpr,
106
+ FH: tl.constexpr,
107
+ FW: tl.constexpr,
108
+ TF: tl.constexpr,
109
+ HF: tl.constexpr,
110
+):
111
+ pid = tl.program_id(0).to(tl.int64)
112
+ rows = pid * BLOCK_M + tl.arange(0, BLOCK_M)
113
+ rmask = rows < R
114
+ offs = tl.arange(0, D_P2)
115
+ mask = rmask[:, None] & (offs < D)[None, :]
116
+ x32 = tl.load(x_ptr + rows[:, None] * D + offs[None, :], mask=mask, other=0.0).to(
117
+ tl.float32
118
+ )
119
+ var = tl.sum(x32 * x32, axis=1) / D
120
+ xn = (x32 * tl.math.rsqrt(var + eps)[:, None]).to(tl.bfloat16)
121
+ w = tl.load(w_ptr + offs, mask=offs < D, other=0.0)
122
+ h = xn * w[None, :] # bf16 multiply, matching ir.ops.rms_norm
123
+ if HAS_GELU:
124
+ # Exact (erf) GELU on the bf16-rounded norm output, fp32 math,
125
+ # matching F.gelu's opmath on a bf16 tensor.
126
+ g32 = h.to(tl.float32)
127
+ h = (0.5 * g32 * (1.0 + tldevice.erf(g32 * 0.7071067811865476))).to(tl.bfloat16)
128
+ if FOLD:
129
+ # Store directly into the next layer's folded layout (a pure
130
+ # permutation — replaces the separate fold copy pass).
131
+ t = (rows // (FH * FW)) % FT
132
+ hh = (rows // FW) % FH
133
+ ww = rows % FW
134
+ n = rows // (FT * FH * FW)
135
+ slot = ((t % TF) * HF + hh % HF) * HF + ww % HF
136
+ out_row = ((n * (FT // TF) + t // TF) * (FH // HF) + hh // HF) * (
137
+ FW // HF
138
+ ) + ww // HF
139
+ base = (out_row * (TF * HF * HF) + slot) * D
140
+ tl.store(out_ptr + base[:, None] + offs[None, :], h, mask=mask)
141
+ else:
142
+ tl.store(out_ptr + rows[:, None] * D + offs[None, :], h, mask=mask)
143
+
144
+
145
+def rmsnorm_gelu(
146
+ x: torch.Tensor, # [..., D] bf16 contiguous
147
+ weight: torch.Tensor,
148
+ eps: float,
149
+ gelu: bool = True,
150
+ fold: tuple[int, int] | None = None, # (t_fold, hw_fold) of the NEXT fold
151
+) -> torch.Tensor:
152
+ """Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block
153
+ when D is small. With ``fold``, x must be [N, T, H, W, D] and the output
154
+ comes back as ``fold_timespace_to_depth(result, *fold)``."""
155
+ D = x.shape[-1]
156
+ flat = x.reshape(-1, D)
157
+ assert flat.stride(1) == 1 and flat.stride(0) == D
158
+ R = flat.shape[0]
159
+ if fold is None:
160
+ out = torch.empty_like(flat)
161
+ ft = fh = fw = tf = hf = 1
162
+ out_shape = x.shape
163
+ else:
164
+ tf, hf = fold
165
+ N, ft, fh, fw, _ = x.shape
166
+ out_shape = (N, ft // tf, fh // hf, fw // hf, tf * hf * hf * D)
167
+ out = torch.empty(out_shape, dtype=x.dtype, device=x.device)
168
+ if R == 0:
169
+ return out.reshape(out_shape)
170
+ d_p2 = triton.next_power_of_2(D)
171
+ block_m = max(1, 4096 // d_p2)
172
+ _rmsnorm_gelu_kernel[(triton.cdiv(R, block_m),)](
173
+ flat,
174
+ weight,
175
+ out,
176
+ eps,
177
+ R,
178
+ D,
179
+ D_P2=d_p2,
180
+ BLOCK_M=block_m,
181
+ HAS_GELU=gelu,
182
+ FOLD=fold is not None,
183
+ FT=ft,
184
+ FH=fh,
185
+ FW=fw,
186
+ TF=tf,
187
+ HF=hf,
188
+ num_warps=_get_num_warps_from_block_size(d_p2 * block_m),
189
+ )
190
+ return out.reshape(out_shape)
vllm/models/inkling/amd/ops/norm.py
new
+414
@@ -0,0 +1,414 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from __future__ import annotations
4
+
5
+from functools import lru_cache
6
+
7
+import torch
8
+
9
+from vllm.triton_utils import tl, triton
10
+
11
+_MAX_FUSED_SIZE = 65536
12
+
13
+
14
+def _get_num_warps_from_block_size(block_size: int) -> int:
15
+ if block_size >= 32768:
16
+ return 32
17
+ if block_size >= 8192:
18
+ return 16
19
+ if block_size >= 2048:
20
+ return 8
21
+ return 4
22
+
23
+
24
+def _largest_power_of_2(n: int) -> int:
25
+ assert n > 0, f"{n=}"
26
+ return 1 << (n.bit_length() - 1)
27
+
28
+
29
+@lru_cache(maxsize=128)
30
+def _get_grid_size_for_mem_bw_kernel(device: torch.device, factor: int = 8) -> int:
31
+ num_sms = torch.cuda.get_device_properties(device).multi_processor_count
32
+ return _largest_power_of_2(num_sms) * factor
33
+
34
+
35
+@triton.jit
36
+def _rmsnorm_fwd_kernel(
37
+ x_ptr,
38
+ weight_ptr,
39
+ y_ptr,
40
+ rstd_ptr,
41
+ eps,
42
+ x_stride_0,
43
+ y_stride_0,
44
+ n_cols,
45
+ block_size_n: tl.constexpr,
46
+):
47
+ pid_m = tl.program_id(0).to(tl.int64)
48
+ offs_n = tl.arange(0, block_size_n)
49
+ mask_n = offs_n < n_cols
50
+
51
+ weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
52
+ x = tl.load(x_ptr + pid_m * x_stride_0 + offs_n, mask=mask_n, other=0.0).to(
53
+ tl.float32
54
+ )
55
+ row_var = tl.sum(x * x, axis=0) / n_cols
56
+ rstd = tl.math.rsqrt(row_var + eps)
57
+ tl.store(rstd_ptr + pid_m, rstd)
58
+
59
+ y = x * rstd * weight
60
+ tl.store(y_ptr + pid_m * y_stride_0 + offs_n, y, mask=mask_n)
61
+
62
+
63
+@triton.jit(do_not_specialize=["n_rows"])
64
+def _rmsnorm_fwd_kernel_block_m(
65
+ x_ptr,
66
+ weight_ptr,
67
+ y_ptr,
68
+ rstd_ptr,
69
+ eps,
70
+ x_stride_0,
71
+ y_stride_0,
72
+ n_rows,
73
+ n_cols,
74
+ block_size_m: tl.constexpr,
75
+ block_size_n: tl.constexpr,
76
+):
77
+ pid_m = tl.program_id(0).to(tl.int64)
78
+ offs_n = tl.arange(0, block_size_n)
79
+ mask_n = offs_n < n_cols
80
+
81
+ num_blocks_m = tl.cdiv(n_rows, block_size_m)
82
+ blocks_per_pid = tl.cdiv(num_blocks_m, tl.num_programs(0))
83
+ block_id_start = pid_m * blocks_per_pid
84
+ block_id_end = min(block_id_start + blocks_per_pid, num_blocks_m)
85
+
86
+ weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
87
+
88
+ for block_id in range(block_id_start, block_id_end):
89
+ offs_m = block_id * block_size_m + tl.arange(0, block_size_m)
90
+ mask_m = offs_m < n_rows
91
+ mask_mn = mask_m[:, None] & mask_n[None, :]
92
+ x = tl.load(
93
+ x_ptr + offs_m[:, None] * x_stride_0 + offs_n[None, :],
94
+ mask=mask_mn,
95
+ other=0.0,
96
+ ).to(tl.float32)
97
+ row_var = tl.sum(x * x, axis=1) / n_cols
98
+ rstd = tl.math.rsqrt(row_var + eps)
99
+ tl.store(rstd_ptr + offs_m, rstd, mask=mask_m)
100
+
101
+ y = x * rstd[:, None] * weight
102
+ tl.store(
103
+ y_ptr + offs_m[:, None] * y_stride_0 + offs_n[None, :],
104
+ y,
105
+ mask=mask_mn,
106
+ )
107
+
108
+
109
+@triton.jit
110
+def _add_rmsnorm_fwd_kernel(
111
+ res_ptr, # [T, N] residual (read)
112
+ delta_ptr, # [T, N] delta to add (read)
113
+ weight_ptr,
114
+ y_ptr, # [T, N] normed output
115
+ res_out_ptr, # [T, N] updated residual output
116
+ eps,
117
+ res_stride_0,
118
+ delta_stride_0,
119
+ y_stride_0,
120
+ res_out_stride_0,
121
+ n_cols,
122
+ block_size_n: tl.constexpr,
123
+):
124
+ pid_m = tl.program_id(0).to(tl.int64)
125
+ offs_n = tl.arange(0, block_size_n)
126
+ mask_n = offs_n < n_cols
127
+
128
+ weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
129
+ r = tl.load(res_ptr + pid_m * res_stride_0 + offs_n, mask=mask_n, other=0.0).to(
130
+ tl.float32
131
+ )
132
+ d = tl.load(delta_ptr + pid_m * delta_stride_0 + offs_n, mask=mask_n, other=0.0).to(
133
+ tl.float32
134
+ )
135
+ # Round the sum to the residual dtype first (matches the eager
136
+ # `residual + delta` then rmsnorm-on-bf16 sequence bit-for-bit).
137
+ s = (r + d).to(res_out_ptr.dtype.element_ty)
138
+ tl.store(res_out_ptr + pid_m * res_out_stride_0 + offs_n, s, mask=mask_n)
139
+ x = s.to(tl.float32)
140
+ row_var = tl.sum(x * x, axis=0) / n_cols
141
+ rstd = tl.math.rsqrt(row_var + eps)
142
+ y = x * rstd * weight
143
+ tl.store(y_ptr + pid_m * y_stride_0 + offs_n, y, mask=mask_n)
144
+
145
+
146
+def add_rmsnorm(
147
+ residual: torch.Tensor,
148
+ delta: torch.Tensor,
149
+ weight: torch.Tensor,
150
+ eps: float,
151
+) -> tuple[torch.Tensor, torch.Tensor]:
152
+ """Fused ``res = residual + delta; y = rmsnorm(res)``.
153
+
154
+ Returns ``(y, res)``; both are fresh tensors (cudagraph-friendly, no
155
+ in-place update of the inputs).
156
+ """
157
+ assert residual.ndim == 2 and delta.ndim == 2, (residual.shape, delta.shape)
158
+ n_rows, n_cols = residual.shape
159
+ assert weight.shape[0] == n_cols
160
+ y = torch.empty_like(residual)
161
+ res_out = torch.empty_like(residual)
162
+ if n_rows == 0:
163
+ return y, res_out
164
+
165
+ block_size_n = triton.next_power_of_2(n_cols)
166
+ max_block_size_n = _MAX_FUSED_SIZE // residual.element_size()
167
+ if max_block_size_n < block_size_n:
168
+ raise RuntimeError(f"Large {n_cols=} is not supported")
169
+ num_warps = _get_num_warps_from_block_size(block_size_n)
170
+ _add_rmsnorm_fwd_kernel[(n_rows,)](
171
+ residual,
172
+ delta,
173
+ weight,
174
+ y,
175
+ res_out,
176
+ eps,
177
+ residual.stride(0),
178
+ delta.stride(0),
179
+ y.stride(0),
180
+ res_out.stride(0),
181
+ n_cols,
182
+ block_size_n,
183
+ num_warps=num_warps,
184
+ )
185
+ return y, res_out
186
+
187
+
188
+@triton.jit
189
+def _embed_rmsnorm_kernel(
190
+ ids_ptr, # [T] token ids
191
+ table_ptr, # [V, N] embedding table
192
+ weight_ptr, # [N] (HAS_NORM only)
193
+ chain_weight_ptr, # [N] (HAS_CHAIN only)
194
+ out_ptr, # [T, N] rmsnorm(table[ids], weight)
195
+ chain_out_ptr, # [T, N] rmsnorm(out, chain_weight) (HAS_CHAIN only)
196
+ eps,
197
+ table_stride_0,
198
+ n_cols,
199
+ block_size_n: tl.constexpr,
200
+ HAS_NORM: tl.constexpr,
201
+ HAS_CHAIN: tl.constexpr,
202
+):
203
+ pid_m = tl.program_id(0).to(tl.int64)
204
+ offs_n = tl.arange(0, block_size_n)
205
+ mask_n = offs_n < n_cols
206
+ row = tl.load(ids_ptr + pid_m).to(tl.int64)
207
+ x = tl.load(table_ptr + row * table_stride_0 + offs_n, mask=mask_n, other=0.0)
208
+ if HAS_NORM:
209
+ xf = x.to(tl.float32)
210
+ rstd = tl.math.rsqrt(tl.sum(xf * xf, axis=0) / n_cols + eps)
211
+ w = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
212
+ # Round to the output dtype so the chained norm is bit-exact vs the
213
+ # unfused pair (which stores bf16 in between).
214
+ x = (xf * rstd * w).to(out_ptr.dtype.element_ty)
215
+ tl.store(out_ptr + pid_m * n_cols + offs_n, x, mask=mask_n)
216
+ if HAS_CHAIN:
217
+ xf = x.to(tl.float32)
218
+ rstd = tl.math.rsqrt(tl.sum(xf * xf, axis=0) / n_cols + eps)
219
+ w = tl.load(chain_weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
220
+ tl.store(
221
+ chain_out_ptr + pid_m * n_cols + offs_n,
222
+ (xf * rstd * w).to(chain_out_ptr.dtype.element_ty),
223
+ mask=mask_n,
224
+ )
225
+
226
+
227
+def embed_rmsnorm(
228
+ input_ids: torch.Tensor,
229
+ embed_table: torch.Tensor,
230
+ weight: torch.Tensor | None,
231
+ eps: float,
232
+ chain_weight: torch.Tensor | None = None,
233
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
234
+ """Fused ``rmsnorm(embed_table[input_ids], weight)`` row gather + norm.
235
+
236
+ Requires the full vocab on-rank (replicated or tp_size == 1).
237
+ ``weight=None`` skips the norm (``use_embed_norm=False``), leaving a pure
238
+ embedding-table gather.
239
+ ``chain_weight`` additionally emits ``rmsnorm(out, chain_weight)`` (the
240
+ first decoder layer's pre-attention norm) as a second output, still one
241
+ launch. Bit-exact vs the unfused module sequence."""
242
+ ids = input_ids.view(-1)
243
+ (T,) = ids.shape
244
+ n = embed_table.shape[1]
245
+ out = torch.empty(
246
+ (*input_ids.shape, n), dtype=embed_table.dtype, device=embed_table.device
247
+ )
248
+ chain_out = torch.empty_like(out) if chain_weight is not None else None
249
+ if T > 0:
250
+ block_size_n = triton.next_power_of_2(n)
251
+ _embed_rmsnorm_kernel[(T,)](
252
+ ids,
253
+ embed_table,
254
+ weight if weight is not None else embed_table,
255
+ chain_weight if chain_weight is not None else embed_table,
256
+ out,
257
+ chain_out if chain_out is not None else out,
258
+ eps,
259
+ embed_table.stride(0),
260
+ n,
261
+ block_size_n,
262
+ HAS_NORM=weight is not None,
263
+ HAS_CHAIN=chain_weight is not None,
264
+ num_warps=_get_num_warps_from_block_size(block_size_n),
265
+ )
266
+ if chain_out is not None:
267
+ return out, chain_out
268
+ return out
269
+
270
+
271
+@triton.jit
272
+def _embed_dual_rmsnorm_cat_kernel(
273
+ hidden_ptr, # [T, N]
274
+ emb_ptr, # [T, N] embeddings, or the [V, N] embedding table when GATHER
275
+ ids_ptr, # [T] token ids (GATHER only)
276
+ w_hidden_ptr, # [N]
277
+ w_pre_ptr, # [N] chained pre-norm on the embed side (HAS_PRE_NORM only)
278
+ w_embed_ptr, # [N]
279
+ out_ptr, # [T, 2N]: [rmsnorm(hidden) | rmsnorm(rmsnorm?(emb))]
280
+ eps,
281
+ hidden_stride_0,
282
+ emb_stride_0,
283
+ n_cols,
284
+ block_size_n: tl.constexpr,
285
+ GATHER: tl.constexpr,
286
+ HAS_PRE_NORM: tl.constexpr,
287
+):
288
+ pid_m = tl.program_id(0).to(tl.int64)
289
+ which = tl.program_id(1) # 0 -> hidden into cols [0, N); 1 -> emb into [N, 2N)
290
+ offs_n = tl.arange(0, block_size_n)
291
+ mask_n = offs_n < n_cols
292
+ if which == 0:
293
+ x = tl.load(
294
+ hidden_ptr + pid_m * hidden_stride_0 + offs_n, mask=mask_n, other=0.0
295
+ ).to(tl.float32)
296
+ w = tl.load(w_hidden_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
297
+ else:
298
+ row = tl.load(ids_ptr + pid_m).to(tl.int64) if GATHER else pid_m
299
+ x = tl.load(emb_ptr + row * emb_stride_0 + offs_n, mask=mask_n, other=0.0).to(
300
+ tl.float32
301
+ )
302
+ if HAS_PRE_NORM:
303
+ w_pre = tl.load(w_pre_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
304
+ rstd = tl.math.rsqrt(tl.sum(x * x, axis=0) / n_cols + eps)
305
+ # Round-trip through the output dtype so the chained norm is
306
+ # bit-exact vs the unfused pair (which stores bf16 in between).
307
+ x = (x * rstd * w_pre).to(out_ptr.dtype.element_ty).to(tl.float32)
308
+ w = tl.load(w_embed_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
309
+ rstd = tl.math.rsqrt(tl.sum(x * x, axis=0) / n_cols + eps)
310
+ tl.store(
311
+ out_ptr + pid_m * (2 * n_cols) + which * n_cols + offs_n,
312
+ (x * rstd * w).to(out_ptr.dtype.element_ty),
313
+ mask=mask_n,
314
+ )
315
+
316
+
317
+def embed_dual_rmsnorm_cat(
318
+ hidden: torch.Tensor,
319
+ hidden_weight: torch.Tensor,
320
+ embed_weight: torch.Tensor,
321
+ eps: float,
322
+ *,
323
+ embeds: torch.Tensor | None = None,
324
+ input_ids: torch.Tensor | None = None,
325
+ embed_table: torch.Tensor | None = None,
326
+ pre_norm_weight: torch.Tensor | None = None,
327
+) -> torch.Tensor:
328
+ """The MTP depth-layer input in one launch:
329
+ ``cat([rmsnorm(hidden, w_h), rmsnorm(pre?(emb), w_e)], -1)``.
330
+
331
+ The embed side is either a fused row gather ``embed_table[input_ids]``
332
+ (draft decode steps) or precomputed ``embeds`` ([T, N], the target-merged
333
+ multimodal embeddings at draft prefill); ``pre_norm_weight`` chains the
334
+ backbone embed_norm in front of the depth embed_norm (bit-exact vs the
335
+ unfused sequence). The concat copies collapse into direct writes."""
336
+ T, n = hidden.shape
337
+ if embeds is not None:
338
+ assert embeds.shape == hidden.shape
339
+ src, ids, src_stride = embeds, embeds, embeds.stride(0)
340
+ gather = False
341
+ else:
342
+ assert input_ids is not None and embed_table is not None
343
+ assert input_ids.shape == (T,) and embed_table.shape[1] == n
344
+ src, ids, src_stride = embed_table, input_ids, embed_table.stride(0)
345
+ gather = True
346
+ out = torch.empty((T, 2 * n), dtype=hidden.dtype, device=hidden.device)
347
+ if T == 0:
348
+ return out
349
+ block_size_n = triton.next_power_of_2(n)
350
+ _embed_dual_rmsnorm_cat_kernel[(T, 2)](
351
+ hidden,
352
+ src,
353
+ ids,
354
+ hidden_weight,
355
+ pre_norm_weight if pre_norm_weight is not None else embed_weight,
356
+ embed_weight,
357
+ out,
358
+ eps,
359
+ hidden.stride(0),
360
+ src_stride,
361
+ n,
362
+ block_size_n,
363
+ GATHER=gather,
364
+ HAS_PRE_NORM=pre_norm_weight is not None,
365
+ num_warps=_get_num_warps_from_block_size(block_size_n),
366
+ )
367
+ return out
368
+
369
+
370
+def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
371
+ assert x.ndim == 2, f"{x.shape=}"
372
+ assert weight.ndim == 1, f"{weight.shape=}"
373
+ n_rows, n_cols = x.shape
374
+ assert weight.shape[0] == n_cols, f"{weight.shape=} {x.shape=}"
375
+ y = torch.empty_like(x)
376
+ rstd = torch.empty((n_rows,), dtype=torch.float32, device=x.device)
377
+
378
+ block_size_n = triton.next_power_of_2(n_cols)
379
+ max_block_size_n = _MAX_FUSED_SIZE // x.element_size()
380
+ if max_block_size_n < block_size_n:
381
+ raise RuntimeError(f"Large {n_cols=} is not supported")
382
+ block_size_m = max(1, 4096 // block_size_n)
383
+ num_warps = _get_num_warps_from_block_size(block_size_n)
384
+
385
+ if block_size_m == 1:
386
+ _rmsnorm_fwd_kernel[(n_rows,)](
387
+ x,
388
+ weight,
389
+ y,
390
+ rstd,
391
+ eps,
392
+ x.stride(0),
393
+ y.stride(0),
394
+ n_cols,
395
+ block_size_n,
396
+ num_warps=num_warps,
397
+ )
398
+ else:
399
+ grid_size = _get_grid_size_for_mem_bw_kernel(x.device)
400
+ _rmsnorm_fwd_kernel_block_m[(grid_size,)](
401
+ x,
402
+ weight,
403
+ y,
404
+ rstd,
405
+ eps,
406
+ x.stride(0),
407
+ y.stride(0),
408
+ n_rows,
409
+ n_cols,
410
+ block_size_m,
411
+ block_size_n,
412
+ num_warps=num_warps,
413
+ )
414
+ return y
vllm/models/inkling/amd/ops/qkvr_prep.py
new
+918
@@ -0,0 +1,918 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.triton_utils import tl, triton
7
+from vllm.utils.torch_utils import aux_stream
8
+
9
+LOW_BLOCK_M = 32
10
+LOW_BLOCK_N = 64
11
+LOW_NUM_WARPS = 4
12
+THROUGHPUT_BLOCK_M = 32
13
+THROUGHPUT_BLOCK_N = 128
14
+THROUGHPUT_GROUP_M = 2
15
+THROUGHPUT_NUM_WARPS = 4
16
+SMALL_TOKEN_THRESHOLD = 128
17
+SMALL_NUM_WARPS = 2
18
+Q_BLOCK_ROWS = 8
19
+Q_NUM_WARPS = 2
20
+KV_BLOCK_ROWS = 4
21
+KV_NUM_WARPS = 2
22
+
23
+
24
+@triton.jit(do_not_specialize=["rows"])
25
+def _rel_proj_low_latency_kernel(
26
+ qkvr_ptr,
27
+ rel_proj_ptr,
28
+ rel_out_ptr,
29
+ log_scaling_ptr,
30
+ rows,
31
+ stride_x_t,
32
+ R_OFFSET: tl.constexpr,
33
+ NUM_Q_HEADS: tl.constexpr,
34
+ REL_EXTENT: tl.constexpr,
35
+ BLOCK_M: tl.constexpr,
36
+ BLOCK_N: tl.constexpr,
37
+ APPLY_LOG_SCALING: tl.constexpr,
38
+):
39
+ row = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
40
+ col = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
41
+ inner = tl.arange(0, 16)
42
+ token = row // NUM_Q_HEADS
43
+ head = row % NUM_Q_HEADS
44
+ relative = tl.load(
45
+ qkvr_ptr
46
+ + token[:, None] * stride_x_t
47
+ + R_OFFSET
48
+ + head[:, None] * 16
49
+ + inner[None, :],
50
+ mask=row[:, None] < rows,
51
+ other=0.0,
52
+ )
53
+ projection = tl.load(
54
+ rel_proj_ptr + inner[:, None] * REL_EXTENT + col[None, :],
55
+ mask=col[None, :] < REL_EXTENT,
56
+ other=0.0,
57
+ )
58
+ values = tl.dot(relative, projection, out_dtype=tl.float32).to(
59
+ rel_out_ptr.dtype.element_ty
60
+ )
61
+ values = values.to(tl.float32)
62
+ if APPLY_LOG_SCALING:
63
+ values *= tl.load(log_scaling_ptr + token, mask=row < rows, other=1.0)[:, None]
64
+ tl.store(
65
+ rel_out_ptr + row[:, None] * REL_EXTENT + col[None, :],
66
+ values.to(rel_out_ptr.dtype.element_ty),
67
+ mask=(row[:, None] < rows) & (col[None, :] < REL_EXTENT),
68
+ )
69
+
70
+
71
+@triton.jit(do_not_specialize=["rows"])
72
+def _rel_proj_throughput_kernel(
73
+ qkvr_ptr,
74
+ rel_proj_ptr,
75
+ rel_out_ptr,
76
+ log_scaling_ptr,
77
+ rows,
78
+ stride_x_t,
79
+ R_OFFSET: tl.constexpr,
80
+ NUM_Q_HEADS: tl.constexpr,
81
+ REL_EXTENT: tl.constexpr,
82
+ BLOCK_M: tl.constexpr,
83
+ BLOCK_N: tl.constexpr,
84
+ GROUP_M: tl.constexpr,
85
+ APPLY_LOG_SCALING: tl.constexpr,
86
+):
87
+ row_group = tl.program_id(0)
88
+ col = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
89
+ inner = tl.arange(0, 16)
90
+ projection = tl.load(
91
+ rel_proj_ptr + inner[:, None] * REL_EXTENT + col[None, :],
92
+ mask=col[None, :] < REL_EXTENT,
93
+ other=0.0,
94
+ )
95
+ row_offsets = tl.arange(0, BLOCK_M)
96
+ for group_offset in tl.static_range(GROUP_M):
97
+ row = (row_group * GROUP_M + group_offset) * BLOCK_M + row_offsets
98
+ token = row // NUM_Q_HEADS
99
+ head = row % NUM_Q_HEADS
100
+ relative = tl.load(
101
+ qkvr_ptr
102
+ + token[:, None] * stride_x_t
103
+ + R_OFFSET
104
+ + head[:, None] * 16
105
+ + inner[None, :],
106
+ mask=row[:, None] < rows,
107
+ other=0.0,
108
+ )
109
+ values = tl.dot(relative, projection, out_dtype=tl.float32).to(
110
+ rel_out_ptr.dtype.element_ty
111
+ )
112
+ values = values.to(tl.float32)
113
+ if APPLY_LOG_SCALING:
114
+ values *= tl.load(log_scaling_ptr + token, mask=row < rows, other=1.0)[
115
+ :, None
116
+ ]
117
+ tl.store(
118
+ rel_out_ptr + row[:, None] * REL_EXTENT + col[None, :],
119
+ values.to(rel_out_ptr.dtype.element_ty),
120
+ mask=(row[:, None] < rows) & (col[None, :] < REL_EXTENT),
121
+ )
122
+
123
+
124
+def use_rel_proj_throughput(rows: int, rel_extent: int) -> bool:
125
+ min_rows = 8192 if rel_extent == 512 else 2048
126
+ return rows >= min_rows
127
+
128
+
129
+def qkvr_rel_proj(
130
+ qkvr: torch.Tensor,
131
+ rel_proj: torch.Tensor,
132
+ rel_out: torch.Tensor,
133
+ log_scaling: torch.Tensor | None,
134
+ *,
135
+ num_q_heads: int,
136
+ num_kv_heads: int,
137
+ head_dim: int,
138
+ d_rel: int,
139
+) -> None:
140
+ rows = qkvr.shape[0] * num_q_heads
141
+ rel_extent = rel_proj.shape[1]
142
+ assert d_rel == 16 and rel_proj.shape[0] == 16
143
+ r_offset = num_q_heads * head_dim + 2 * num_kv_heads * head_dim
144
+ log_scaling_ptr = log_scaling if log_scaling is not None else qkvr
145
+ common = dict(
146
+ R_OFFSET=r_offset,
147
+ NUM_Q_HEADS=num_q_heads,
148
+ REL_EXTENT=rel_extent,
149
+ APPLY_LOG_SCALING=log_scaling is not None,
150
+ )
151
+
152
+ if use_rel_proj_throughput(rows, rel_extent):
153
+ grid = (
154
+ triton.cdiv(rows, THROUGHPUT_BLOCK_M * THROUGHPUT_GROUP_M),
155
+ triton.cdiv(rel_extent, THROUGHPUT_BLOCK_N),
156
+ )
157
+ _rel_proj_throughput_kernel[grid](
158
+ qkvr,
159
+ rel_proj,
160
+ rel_out,
161
+ log_scaling_ptr,
162
+ rows,
163
+ qkvr.stride(0),
164
+ BLOCK_M=THROUGHPUT_BLOCK_M,
165
+ BLOCK_N=THROUGHPUT_BLOCK_N,
166
+ GROUP_M=THROUGHPUT_GROUP_M,
167
+ num_warps=THROUGHPUT_NUM_WARPS,
168
+ **common,
169
+ )
170
+ return
171
+
172
+ grid = (
173
+ triton.cdiv(rows, LOW_BLOCK_M),
174
+ triton.cdiv(rel_extent, LOW_BLOCK_N),
175
+ )
176
+ _rel_proj_low_latency_kernel[grid](
177
+ qkvr,
178
+ rel_proj,
179
+ rel_out,
180
+ log_scaling_ptr,
181
+ rows,
182
+ qkvr.stride(0),
183
+ BLOCK_M=LOW_BLOCK_M,
184
+ BLOCK_N=LOW_BLOCK_N,
185
+ num_warps=LOW_NUM_WARPS,
186
+ **common,
187
+ )
188
+
189
+
190
+@triton.jit(do_not_specialize=["tokens", "stride_block_table_req", "max_blocks"])
191
+def _qkvr_qkv_kernel(
192
+ qkvr_ptr,
193
+ q_norm_weight_ptr,
194
+ q_out_ptr,
195
+ rel_proj_ptr,
196
+ rel_out_ptr,
197
+ k_weight_ptr,
198
+ v_weight_ptr,
199
+ k_norm_weight_ptr,
200
+ conv_cache_ptr,
201
+ key_cache_ptr,
202
+ value_cache_ptr,
203
+ positions_ptr,
204
+ seq_idx_ptr,
205
+ conv_slot_mapping_ptr,
206
+ conv_block_table_ptr,
207
+ query_start_ptr,
208
+ attention_slot_mapping_ptr,
209
+ log_scaling_ptr,
210
+ tokens,
211
+ eps,
212
+ stride_x_t,
213
+ stride_cc_block,
214
+ stride_cc_head,
215
+ stride_cc_token,
216
+ stride_cc_dim,
217
+ stride_kc_block,
218
+ stride_kc_token,
219
+ stride_kc_head,
220
+ stride_vc_block,
221
+ stride_vc_token,
222
+ stride_vc_head,
223
+ stride_block_table_req,
224
+ max_blocks,
225
+ conv_block_size,
226
+ attention_page_size,
227
+ Q_WIDTH: tl.constexpr,
228
+ KV_WIDTH: tl.constexpr,
229
+ NUM_Q_HEADS: tl.constexpr,
230
+ NUM_KV_HEADS: tl.constexpr,
231
+ HEAD_DIM: tl.constexpr,
232
+ WINDOW_SIZE: tl.constexpr,
233
+ OFF_K: tl.constexpr,
234
+ OFF_V: tl.constexpr,
235
+ APPLY_LOG_SCALING: tl.constexpr,
236
+ D_REL: tl.constexpr,
237
+ REL_EXTENT: tl.constexpr,
238
+ REL_EXTENT_PADDED: tl.constexpr,
239
+):
240
+ block = tl.program_id(0)
241
+ num_q_rows = tokens * NUM_Q_HEADS
242
+ dims = tl.arange(0, HEAD_DIM)
243
+
244
+ if block < num_q_rows:
245
+ row = block
246
+ token = row // NUM_Q_HEADS
247
+ head = row % NUM_Q_HEADS
248
+ values = tl.load(
249
+ qkvr_ptr + token * stride_x_t + head * HEAD_DIM + dims,
250
+ ).to(tl.float32)
251
+ weight = tl.load(q_norm_weight_ptr + dims).to(tl.float32)
252
+ rstd = tl.rsqrt(tl.sum(values * values, axis=0) / HEAD_DIM + eps)
253
+ normalized = values * rstd * weight
254
+ if APPLY_LOG_SCALING:
255
+ normalized = normalized.to(q_out_ptr.dtype.element_ty).to(tl.float32)
256
+ normalized *= tl.load(log_scaling_ptr + token)
257
+ tl.store(
258
+ q_out_ptr + row * HEAD_DIM + dims,
259
+ normalized.to(q_out_ptr.dtype.element_ty),
260
+ )
261
+ rel_cols = tl.arange(0, REL_EXTENT_PADDED)
262
+ rel_mask = rel_cols < REL_EXTENT
263
+ projected = tl.zeros([REL_EXTENT_PADDED], dtype=tl.float32)
264
+ rel_offset = Q_WIDTH + 2 * KV_WIDTH + head * D_REL
265
+ for rel_dim in tl.static_range(D_REL):
266
+ rel_value = tl.load(
267
+ qkvr_ptr + token * stride_x_t + rel_offset + rel_dim
268
+ ).to(tl.float32)
269
+ proj = tl.load(
270
+ rel_proj_ptr + rel_dim * REL_EXTENT + rel_cols,
271
+ mask=rel_mask,
272
+ other=0.0,
273
+ ).to(tl.float32)
274
+ projected += rel_value * proj
275
+ projected = projected.to(rel_out_ptr.dtype.element_ty).to(tl.float32)
276
+ if APPLY_LOG_SCALING:
277
+ projected *= tl.load(log_scaling_ptr + token)
278
+ tl.store(
279
+ rel_out_ptr + row * REL_EXTENT + rel_cols,
280
+ projected.to(rel_out_ptr.dtype.element_ty),
281
+ mask=rel_mask,
282
+ )
283
+ else:
284
+ row = block - num_q_rows
285
+ if row < tokens * NUM_KV_HEADS:
286
+ token = row // NUM_KV_HEADS
287
+ head = row % NUM_KV_HEADS
288
+ position = tl.load(positions_ptr + token)
289
+ request = tl.load(seq_idx_ptr + token)
290
+ conv_slot = tl.load(conv_slot_mapping_ptr + token)
291
+ query_start = tl.load(query_start_ptr + token)
292
+ attention_slot = tl.load(attention_slot_mapping_ptr + token)
293
+ valid = conv_slot >= 0
294
+
295
+ k_col = Q_WIDTH + head * HEAD_DIM
296
+ v_col = Q_WIDTH + KV_WIDTH + head * HEAD_DIM
297
+ k_value = tl.load(qkvr_ptr + token * stride_x_t + k_col + dims)
298
+ v_value = tl.load(qkvr_ptr + token * stride_x_t + v_col + dims)
299
+
300
+ safe_slot = tl.maximum(conv_slot, 0)
301
+ cache_base = (
302
+ conv_cache_ptr
303
+ + (safe_slot // conv_block_size) * stride_cc_block
304
+ + head * stride_cc_head
305
+ + (safe_slot % conv_block_size) * stride_cc_token
306
+ )
307
+ tl.store(
308
+ cache_base + (OFF_K + dims) * stride_cc_dim,
309
+ k_value,
310
+ mask=valid,
311
+ )
312
+ tl.store(
313
+ cache_base + (OFF_V + dims) * stride_cc_dim,
314
+ v_value,
315
+ mask=valid,
316
+ )
317
+
318
+ acc_k = tl.zeros([HEAD_DIM], dtype=tl.float32)
319
+ acc_v = tl.zeros([HEAD_DIM], dtype=tl.float32)
320
+ for tap in tl.static_range(WINDOW_SIZE):
321
+ source_position = position - (WINDOW_SIZE - 1) + tap
322
+ source_row = token - (WINDOW_SIZE - 1) + tap
323
+ in_window = valid & (source_position >= 0)
324
+ intra = in_window & (source_row >= query_start)
325
+ cached = in_window & (source_row < query_start)
326
+ safe_row = tl.maximum(source_row, 0)
327
+ source_k = tl.load(
328
+ qkvr_ptr + safe_row * stride_x_t + k_col + dims,
329
+ mask=intra,
330
+ other=0.0,
331
+ ).to(tl.float32)
332
+ source_v = tl.load(
333
+ qkvr_ptr + safe_row * stride_x_t + v_col + dims,
334
+ mask=intra,
335
+ other=0.0,
336
+ ).to(tl.float32)
337
+ safe_position = tl.maximum(source_position, 0)
338
+ logical_block = tl.minimum(
339
+ safe_position // conv_block_size, max_blocks - 1
340
+ )
341
+ physical_block = tl.load(
342
+ conv_block_table_ptr
343
+ + request * stride_block_table_req
344
+ + logical_block,
345
+ mask=cached,
346
+ other=0,
347
+ ).to(tl.int64)
348
+ tap_base = (
349
+ conv_cache_ptr
350
+ + physical_block * stride_cc_block
351
+ + head * stride_cc_head
352
+ + (safe_position % conv_block_size) * stride_cc_token
353
+ )
354
+ source_k += tl.load(
355
+ tap_base + (OFF_K + dims) * stride_cc_dim,
356
+ mask=cached,
357
+ other=0.0,
358
+ ).to(tl.float32)
359
+ source_v += tl.load(
360
+ tap_base + (OFF_V + dims) * stride_cc_dim,
361
+ mask=cached,
362
+ other=0.0,
363
+ ).to(tl.float32)
364
+ k_weight = tl.load(
365
+ k_weight_ptr + (head * HEAD_DIM + dims) * WINDOW_SIZE + tap
366
+ ).to(tl.float32)
367
+ v_weight = tl.load(
368
+ v_weight_ptr + (head * HEAD_DIM + dims) * WINDOW_SIZE + tap
369
+ ).to(tl.float32)
370
+ acc_k += source_k * k_weight
371
+ acc_v += source_v * v_weight
372
+
373
+ k_rounded = (acc_k + k_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty)
374
+ v_rounded = (acc_v + v_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty)
375
+ k_float = k_rounded.to(tl.float32)
376
+ k_norm_weight = tl.load(k_norm_weight_ptr + dims).to(tl.float32)
377
+ rstd = tl.rsqrt(tl.sum(k_float * k_float, axis=0) / HEAD_DIM + eps)
378
+ k_normalized = (k_float * rstd * k_norm_weight).to(
379
+ qkvr_ptr.dtype.element_ty
380
+ )
381
+
382
+ safe_attention_slot = tl.maximum(attention_slot, 0)
383
+ attention_block = safe_attention_slot // attention_page_size
384
+ attention_offset = safe_attention_slot % attention_page_size
385
+ tl.store(
386
+ key_cache_ptr
387
+ + attention_block * stride_kc_block
388
+ + attention_offset * stride_kc_token
389
+ + head * stride_kc_head
390
+ + dims,
391
+ k_normalized,
392
+ mask=attention_slot >= 0,
393
+ )
394
+ tl.store(
395
+ value_cache_ptr
396
+ + attention_block * stride_vc_block
397
+ + attention_offset * stride_vc_token
398
+ + head * stride_vc_head
399
+ + dims,
400
+ v_rounded,
401
+ mask=attention_slot >= 0,
402
+ )
403
+
404
+
405
+@triton.jit(do_not_specialize=["num_rows"])
406
+def _q_kernel(
407
+ qkvr_ptr,
408
+ q_norm_weight_ptr,
409
+ q_out_ptr,
410
+ log_scaling_ptr,
411
+ num_rows,
412
+ stride_x_t,
413
+ eps,
414
+ NUM_Q_HEADS: tl.constexpr,
415
+ HEAD_DIM: tl.constexpr,
416
+ BLOCK_ROWS: tl.constexpr,
417
+ APPLY_LOG_SCALING: tl.constexpr,
418
+):
419
+ rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
420
+ dims = tl.arange(0, HEAD_DIM)
421
+ row_mask = rows < num_rows
422
+ tokens = rows // NUM_Q_HEADS
423
+ heads = rows % NUM_Q_HEADS
424
+ values = tl.load(
425
+ qkvr_ptr
426
+ + tokens[:, None] * stride_x_t
427
+ + heads[:, None] * HEAD_DIM
428
+ + dims[None, :],
429
+ mask=row_mask[:, None],
430
+ other=0.0,
431
+ ).to(tl.float32)
432
+ weight = tl.load(q_norm_weight_ptr + dims).to(tl.float32)
433
+ rstd = tl.rsqrt(tl.sum(values * values, axis=1) / HEAD_DIM + eps)
434
+ normalized = values * rstd[:, None] * weight[None, :]
435
+ if APPLY_LOG_SCALING:
436
+ normalized = normalized.to(q_out_ptr.dtype.element_ty).to(tl.float32)
437
+ tau = tl.load(log_scaling_ptr + tokens, mask=row_mask, other=1.0)
438
+ normalized *= tau[:, None]
439
+ tl.store(
440
+ q_out_ptr + rows[:, None] * HEAD_DIM + dims[None, :],
441
+ normalized.to(q_out_ptr.dtype.element_ty),
442
+ mask=row_mask[:, None],
443
+ )
444
+
445
+
446
+@triton.jit(do_not_specialize=["tokens", "stride_block_table_req", "max_blocks"])
447
+def _kv_kernel(
448
+ qkvr_ptr,
449
+ k_weight_ptr,
450
+ v_weight_ptr,
451
+ k_norm_weight_ptr,
452
+ conv_cache_ptr,
453
+ key_cache_ptr,
454
+ value_cache_ptr,
455
+ positions_ptr,
456
+ seq_idx_ptr,
457
+ conv_slot_mapping_ptr,
458
+ conv_block_table_ptr,
459
+ query_start_ptr,
460
+ attention_slot_mapping_ptr,
461
+ tokens,
462
+ eps,
463
+ stride_x_t,
464
+ stride_cc_block,
465
+ stride_cc_head,
466
+ stride_cc_token,
467
+ stride_cc_dim,
468
+ stride_kc_block,
469
+ stride_kc_token,
470
+ stride_kc_head,
471
+ stride_vc_block,
472
+ stride_vc_token,
473
+ stride_vc_head,
474
+ stride_block_table_req,
475
+ max_blocks,
476
+ conv_block_size,
477
+ attention_page_size,
478
+ Q_WIDTH: tl.constexpr,
479
+ KV_WIDTH: tl.constexpr,
480
+ NUM_KV_HEADS: tl.constexpr,
481
+ HEAD_DIM: tl.constexpr,
482
+ WINDOW_SIZE: tl.constexpr,
483
+ OFF_K: tl.constexpr,
484
+ OFF_V: tl.constexpr,
485
+ BLOCK_ROWS: tl.constexpr,
486
+):
487
+ token = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
488
+ dims = tl.arange(0, HEAD_DIM)
489
+ row_mask = token < tokens
490
+ head_id = tl.program_id(1)
491
+ head = tl.full([BLOCK_ROWS], head_id, tl.int64)
492
+ position = tl.load(positions_ptr + token, mask=row_mask, other=0)
493
+ request = tl.load(seq_idx_ptr + token, mask=row_mask, other=0)
494
+ conv_slot = tl.load(conv_slot_mapping_ptr + token, mask=row_mask, other=-1)
495
+ query_start = tl.load(query_start_ptr + token, mask=row_mask, other=0)
496
+ attention_slot = tl.load(
497
+ attention_slot_mapping_ptr + token, mask=row_mask, other=-1
498
+ )
499
+ valid = row_mask & (conv_slot >= 0)
500
+
501
+ k_col = Q_WIDTH + head * HEAD_DIM
502
+ v_col = Q_WIDTH + KV_WIDTH + head * HEAD_DIM
503
+ k_value = tl.load(
504
+ qkvr_ptr + token[:, None] * stride_x_t + k_col[:, None] + dims[None, :],
505
+ mask=row_mask[:, None],
506
+ other=0.0,
507
+ )
508
+ v_value = tl.load(
509
+ qkvr_ptr + token[:, None] * stride_x_t + v_col[:, None] + dims[None, :],
510
+ mask=row_mask[:, None],
511
+ other=0.0,
512
+ )
513
+ safe_slot = tl.maximum(conv_slot, 0)
514
+ cache_base = (
515
+ conv_cache_ptr
516
+ + (safe_slot // conv_block_size) * stride_cc_block
517
+ + head * stride_cc_head
518
+ + (safe_slot % conv_block_size) * stride_cc_token
519
+ )
520
+ tl.store(
521
+ cache_base[:, None] + (OFF_K + dims[None, :]) * stride_cc_dim,
522
+ k_value,
523
+ mask=valid[:, None],
524
+ )
525
+ tl.store(
526
+ cache_base[:, None] + (OFF_V + dims[None, :]) * stride_cc_dim,
527
+ v_value,
528
+ mask=valid[:, None],
529
+ )
530
+
531
+ acc_k = tl.zeros([BLOCK_ROWS, HEAD_DIM], dtype=tl.float32)
532
+ acc_v = tl.zeros([BLOCK_ROWS, HEAD_DIM], dtype=tl.float32)
533
+ for tap in tl.static_range(WINDOW_SIZE):
534
+ source_position = position - (WINDOW_SIZE - 1) + tap
535
+ source_row = token - (WINDOW_SIZE - 1) + tap
536
+ in_window = valid & (source_position >= 0)
537
+ intra = in_window & (source_row >= query_start)
538
+ cached = in_window & (source_row < query_start)
539
+ safe_row = tl.maximum(source_row, 0)
540
+ source_k = tl.load(
541
+ qkvr_ptr + safe_row[:, None] * stride_x_t + k_col[:, None] + dims[None, :],
542
+ mask=intra[:, None],
543
+ other=0.0,
544
+ ).to(tl.float32)
545
+ source_v = tl.load(
546
+ qkvr_ptr + safe_row[:, None] * stride_x_t + v_col[:, None] + dims[None, :],
547
+ mask=intra[:, None],
548
+ other=0.0,
549
+ ).to(tl.float32)
550
+ safe_position = tl.maximum(source_position, 0)
551
+ logical_block = tl.minimum(safe_position // conv_block_size, max_blocks - 1)
552
+ physical_block = tl.load(
553
+ conv_block_table_ptr + request * stride_block_table_req + logical_block,
554
+ mask=cached,
555
+ other=0,
556
+ ).to(tl.int64)
557
+ tap_base = (
558
+ conv_cache_ptr
559
+ + physical_block * stride_cc_block
560
+ + head * stride_cc_head
561
+ + (safe_position % conv_block_size) * stride_cc_token
562
+ )
563
+ source_k += tl.load(
564
+ tap_base[:, None] + (OFF_K + dims[None, :]) * stride_cc_dim,
565
+ mask=cached[:, None],
566
+ other=0.0,
567
+ ).to(tl.float32)
568
+ source_v += tl.load(
569
+ tap_base[:, None] + (OFF_V + dims[None, :]) * stride_cc_dim,
570
+ mask=cached[:, None],
571
+ other=0.0,
572
+ ).to(tl.float32)
573
+ k_weight = tl.load(
574
+ k_weight_ptr + (head_id * HEAD_DIM + dims) * WINDOW_SIZE + tap
575
+ ).to(tl.float32)
576
+ v_weight = tl.load(
577
+ v_weight_ptr + (head_id * HEAD_DIM + dims) * WINDOW_SIZE + tap
578
+ ).to(tl.float32)
579
+ acc_k += source_k * k_weight[None, :]
580
+ acc_v += source_v * v_weight[None, :]
581
+
582
+ k_rounded = (acc_k + k_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty)
583
+ v_rounded = (acc_v + v_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty)
584
+ k_float = k_rounded.to(tl.float32)
585
+ k_norm_weight = tl.load(k_norm_weight_ptr + dims).to(tl.float32)
586
+ rstd = tl.rsqrt(tl.sum(k_float * k_float, axis=1) / HEAD_DIM + eps)
587
+ k_normalized = (k_float * rstd[:, None] * k_norm_weight[None, :]).to(
588
+ qkvr_ptr.dtype.element_ty
589
+ )
590
+
591
+ safe_attention_slot = tl.maximum(attention_slot, 0)
592
+ attention_block = safe_attention_slot // attention_page_size
593
+ attention_offset = safe_attention_slot % attention_page_size
594
+ attention_mask = row_mask & (attention_slot >= 0)
595
+ tl.store(
596
+ key_cache_ptr
597
+ + attention_block[:, None] * stride_kc_block
598
+ + attention_offset[:, None] * stride_kc_token
599
+ + head[:, None] * stride_kc_head
600
+ + dims[None, :],
601
+ k_normalized,
602
+ mask=attention_mask[:, None],
603
+ )
604
+ tl.store(
605
+ value_cache_ptr
606
+ + attention_block[:, None] * stride_vc_block
607
+ + attention_offset[:, None] * stride_vc_token
608
+ + head[:, None] * stride_vc_head
609
+ + dims[None, :],
610
+ v_rounded,
611
+ mask=attention_mask[:, None],
612
+ )
613
+
614
+
615
+def _run_tiled_q(
616
+ qkvr: torch.Tensor,
617
+ q_norm_weight: torch.Tensor,
618
+ q_out: torch.Tensor,
619
+ positions: torch.Tensor,
620
+ *,
621
+ eps: float,
622
+ num_q_heads: int,
623
+ head_dim: int,
624
+ log_scaling: torch.Tensor | None,
625
+) -> None:
626
+ num_rows = qkvr.shape[0] * num_q_heads
627
+ _q_kernel[(triton.cdiv(num_rows, Q_BLOCK_ROWS),)](
628
+ qkvr,
629
+ q_norm_weight,
630
+ q_out,
631
+ log_scaling if log_scaling is not None else positions,
632
+ num_rows,
633
+ qkvr.stride(0),
634
+ eps,
635
+ NUM_Q_HEADS=num_q_heads,
636
+ HEAD_DIM=head_dim,
637
+ BLOCK_ROWS=Q_BLOCK_ROWS,
638
+ APPLY_LOG_SCALING=log_scaling is not None,
639
+ num_warps=Q_NUM_WARPS,
640
+ )
641
+
642
+
643
+def _run_tiled_kv(
644
+ qkvr: torch.Tensor,
645
+ k_weight: torch.Tensor,
646
+ v_weight: torch.Tensor,
647
+ k_norm_weight: torch.Tensor,
648
+ conv_cache: torch.Tensor,
649
+ key_cache: torch.Tensor,
650
+ value_cache: torch.Tensor,
651
+ positions: torch.Tensor,
652
+ seq_idx: torch.Tensor,
653
+ conv_slot_mapping: torch.Tensor,
654
+ conv_block_table: torch.Tensor,
655
+ query_start: torch.Tensor,
656
+ attention_slot_mapping: torch.Tensor,
657
+ *,
658
+ eps: float,
659
+ num_q_heads: int,
660
+ num_kv_heads: int,
661
+ head_dim: int,
662
+ off_k: int,
663
+ off_v: int,
664
+ conv_block_size: int,
665
+) -> None:
666
+ tokens = qkvr.shape[0]
667
+ _kv_kernel[(triton.cdiv(tokens, KV_BLOCK_ROWS), num_kv_heads)](
668
+ qkvr,
669
+ k_weight,
670
+ v_weight,
671
+ k_norm_weight,
672
+ conv_cache,
673
+ key_cache,
674
+ value_cache,
675
+ positions,
676
+ seq_idx,
677
+ conv_slot_mapping,
678
+ conv_block_table,
679
+ query_start,
680
+ attention_slot_mapping,
681
+ tokens,
682
+ eps,
683
+ qkvr.stride(0),
684
+ conv_cache.stride(0),
685
+ conv_cache.stride(1),
686
+ conv_cache.stride(2),
687
+ conv_cache.stride(3),
688
+ key_cache.stride(0),
689
+ key_cache.stride(1),
690
+ key_cache.stride(2),
691
+ value_cache.stride(0),
692
+ value_cache.stride(1),
693
+ value_cache.stride(2),
694
+ conv_block_table.stride(0),
695
+ conv_block_table.shape[1],
696
+ conv_block_size,
697
+ key_cache.shape[1],
698
+ Q_WIDTH=num_q_heads * head_dim,
699
+ KV_WIDTH=num_kv_heads * head_dim,
700
+ NUM_KV_HEADS=num_kv_heads,
701
+ HEAD_DIM=head_dim,
702
+ WINDOW_SIZE=k_weight.shape[1],
703
+ OFF_K=off_k,
704
+ OFF_V=off_v,
705
+ BLOCK_ROWS=KV_BLOCK_ROWS,
706
+ num_warps=KV_NUM_WARPS,
707
+ )
708
+
709
+
710
+def _run_fused_small(
711
+ qkvr: torch.Tensor,
712
+ q_norm_weight: torch.Tensor,
713
+ q_out: torch.Tensor,
714
+ rel_proj: torch.Tensor,
715
+ rel_out: torch.Tensor,
716
+ k_weight: torch.Tensor,
717
+ v_weight: torch.Tensor,
718
+ k_norm_weight: torch.Tensor,
719
+ conv_cache: torch.Tensor,
720
+ key_cache: torch.Tensor,
721
+ value_cache: torch.Tensor,
722
+ positions: torch.Tensor,
723
+ seq_idx: torch.Tensor,
724
+ conv_slot_mapping: torch.Tensor,
725
+ conv_block_table: torch.Tensor,
726
+ query_start: torch.Tensor,
727
+ attention_slot_mapping: torch.Tensor,
728
+ *,
729
+ eps: float,
730
+ num_q_heads: int,
731
+ num_kv_heads: int,
732
+ head_dim: int,
733
+ off_k: int,
734
+ off_v: int,
735
+ conv_block_size: int,
736
+ log_scaling: torch.Tensor | None,
737
+) -> None:
738
+ tokens = qkvr.shape[0]
739
+
740
+ num_q_rows = tokens * num_q_heads
741
+ grid = (num_q_rows + tokens * num_kv_heads,)
742
+ _qkvr_qkv_kernel[grid](
743
+ qkvr,
744
+ q_norm_weight,
745
+ q_out,
746
+ rel_proj,
747
+ rel_out,
748
+ k_weight,
749
+ v_weight,
750
+ k_norm_weight,
751
+ conv_cache,
752
+ key_cache,
753
+ value_cache,
754
+ positions,
755
+ seq_idx,
756
+ conv_slot_mapping,
757
+ conv_block_table,
758
+ query_start,
759
+ attention_slot_mapping,
760
+ log_scaling if log_scaling is not None else positions,
761
+ tokens,
762
+ eps,
763
+ qkvr.stride(0),
764
+ conv_cache.stride(0),
765
+ conv_cache.stride(1),
766
+ conv_cache.stride(2),
767
+ conv_cache.stride(3),
768
+ key_cache.stride(0),
769
+ key_cache.stride(1),
770
+ key_cache.stride(2),
771
+ value_cache.stride(0),
772
+ value_cache.stride(1),
773
+ value_cache.stride(2),
774
+ conv_block_table.stride(0),
775
+ conv_block_table.shape[1],
776
+ conv_block_size,
777
+ key_cache.shape[1],
778
+ Q_WIDTH=num_q_heads * head_dim,
779
+ KV_WIDTH=num_kv_heads * head_dim,
780
+ NUM_Q_HEADS=num_q_heads,
781
+ NUM_KV_HEADS=num_kv_heads,
782
+ HEAD_DIM=head_dim,
783
+ WINDOW_SIZE=k_weight.shape[1],
784
+ OFF_K=off_k,
785
+ OFF_V=off_v,
786
+ APPLY_LOG_SCALING=log_scaling is not None,
787
+ D_REL=16,
788
+ REL_EXTENT=rel_proj.shape[1],
789
+ REL_EXTENT_PADDED=triton.next_power_of_2(rel_proj.shape[1]),
790
+ num_warps=SMALL_NUM_WARPS,
791
+ )
792
+
793
+
794
+def fused_qkvr_prep(
795
+ qkvr: torch.Tensor,
796
+ k_weight: torch.Tensor,
797
+ v_weight: torch.Tensor,
798
+ q_norm_weight: torch.Tensor,
799
+ k_norm_weight: torch.Tensor,
800
+ rel_proj: torch.Tensor,
801
+ eps: float,
802
+ num_q_heads: int,
803
+ num_kv_heads: int,
804
+ head_dim: int,
805
+ d_rel: int,
806
+ conv_cache: torch.Tensor,
807
+ key_cache: torch.Tensor,
808
+ value_cache: torch.Tensor,
809
+ positions: torch.Tensor,
810
+ conv_block_table: torch.Tensor,
811
+ seq_idx: torch.Tensor,
812
+ conv_slot_mapping: torch.Tensor,
813
+ query_start: torch.Tensor,
814
+ attention_slot_mapping: torch.Tensor,
815
+ off_k: int,
816
+ off_v: int,
817
+ conv_block_size: int,
818
+ log_scaling: torch.Tensor | None = None,
819
+) -> tuple[torch.Tensor, torch.Tensor]:
820
+ assert d_rel == 16 and rel_proj.shape[0] == 16
821
+ assert head_dim == 128
822
+ assert qkvr.is_contiguous()
823
+ assert k_weight.stride() == (k_weight.shape[1], 1)
824
+ assert v_weight.stride() == (v_weight.shape[1], 1)
825
+ assert rel_proj.stride() == (rel_proj.shape[1], 1)
826
+ assert conv_cache.stride(3) == 1
827
+ assert key_cache.stride(3) == 1 and value_cache.stride(3) == 1
828
+ tokens = qkvr.shape[0]
829
+ q_out = torch.empty(
830
+ (tokens, num_q_heads * head_dim), dtype=qkvr.dtype, device=qkvr.device
831
+ )
832
+ rel_out = torch.empty(
833
+ (tokens, num_q_heads, rel_proj.shape[1]),
834
+ dtype=qkvr.dtype,
835
+ device=qkvr.device,
836
+ )
837
+ if tokens == 0:
838
+ return q_out, rel_out
839
+
840
+ if tokens < SMALL_TOKEN_THRESHOLD:
841
+ _run_fused_small(
842
+ qkvr,
843
+ q_norm_weight,
844
+ q_out,
845
+ rel_proj,
846
+ rel_out,
847
+ k_weight,
848
+ v_weight,
849
+ k_norm_weight,
850
+ conv_cache,
851
+ key_cache,
852
+ value_cache,
853
+ positions,
854
+ seq_idx,
855
+ conv_slot_mapping,
856
+ conv_block_table,
857
+ query_start,
858
+ attention_slot_mapping,
859
+ eps=eps,
860
+ num_q_heads=num_q_heads,
861
+ num_kv_heads=num_kv_heads,
862
+ head_dim=head_dim,
863
+ off_k=off_k,
864
+ off_v=off_v,
865
+ conv_block_size=conv_block_size,
866
+ log_scaling=log_scaling,
867
+ )
868
+ return q_out, rel_out
869
+
870
+ kv_stream = aux_stream()
871
+ assert kv_stream is not None
872
+ current_stream = torch.cuda.current_stream()
873
+ kv_stream.wait_stream(current_stream)
874
+ with torch.cuda.stream(kv_stream):
875
+ _run_tiled_kv(
876
+ qkvr,
877
+ k_weight,
878
+ v_weight,
879
+ k_norm_weight,
880
+ conv_cache,
881
+ key_cache,
882
+ value_cache,
883
+ positions,
884
+ seq_idx,
885
+ conv_slot_mapping,
886
+ conv_block_table,
887
+ query_start,
888
+ attention_slot_mapping,
889
+ eps=eps,
890
+ num_q_heads=num_q_heads,
891
+ num_kv_heads=num_kv_heads,
892
+ head_dim=head_dim,
893
+ off_k=off_k,
894
+ off_v=off_v,
895
+ conv_block_size=conv_block_size,
896
+ )
897
+ _run_tiled_q(
898
+ qkvr,
899
+ q_norm_weight,
900
+ q_out,
901
+ positions,
902
+ eps=eps,
903
+ num_q_heads=num_q_heads,
904
+ head_dim=head_dim,
905
+ log_scaling=log_scaling,
906
+ )
907
+ qkvr_rel_proj(
908
+ qkvr,
909
+ rel_proj,
910
+ rel_out,
911
+ log_scaling,
912
+ num_q_heads=num_q_heads,
913
+ num_kv_heads=num_kv_heads,
914
+ head_dim=head_dim,
915
+ d_rel=d_rel,
916
+ )
917
+ current_stream.wait_stream(kv_stream)
918
+ return q_out, rel_out
vllm/models/inkling/amd/ops/rel_attention_decode.py
new
+403
@@ -0,0 +1,403 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# SPDX-License-Identifier: MIT
5
+# Copyright (c) 2026 LightSeek Foundation
6
+#
7
+# Permission is hereby granted, free of charge, to any person obtaining a copy
8
+# of this software and associated documentation files (the "Software"), to deal
9
+# in the Software without restriction, including without limitation the rights
10
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11
+# copies of the Software, and to permit persons to whom the Software is
12
+# furnished to do so, subject to the following conditions:
13
+#
14
+# The above copyright notice and this permission notice shall be included in
15
+# all copies or substantial portions of the Software.
16
+#
17
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23
+# SOFTWARE.
24
+"""Split-KV decode for Inkling relative attention on ROCm.
25
+
26
+Adapted from LightSeek TokenSpeed's portable Triton relative-MHA decode.
27
+"""
28
+
29
+from __future__ import annotations
30
+
31
+import os
32
+
33
+import torch
34
+
35
+from vllm.triton_utils import tl, triton
36
+
37
+_MIN_BLOCK_KV = tl.constexpr(32)
38
+
39
+
40
+def decode_split_count(max_kv_len: int, window_left: int) -> int:
41
+ """Return the number of parallel KV partitions for decode."""
42
+ if window_left >= 0:
43
+ effective_len = max(1, min(max_kv_len, window_left + 1))
44
+ return min(4, max(1, triton.cdiv(effective_len, 128)))
45
+ return min(32, max(1, triton.cdiv(max(1, max_kv_len), 2048)))
46
+
47
+
48
+def use_split_kv_decode(
49
+ *,
50
+ max_query_len: int,
51
+ max_kv_len: int,
52
+ page_size: int,
53
+ window_left: int,
54
+) -> bool:
55
+ """Select split-KV only where it outperforms the single-pass kernel."""
56
+ if os.getenv("INKLING_SPLIT_KV", "1") != "1":
57
+ return False
58
+ if max_query_len != 1:
59
+ return False
60
+ if page_size >= 64:
61
+ return True
62
+ if window_left >= 0:
63
+ return page_size >= 32
64
+ return max_kv_len >= 8192
65
+
66
+
67
+@triton.jit
68
+def _split_kv_stage1(
69
+ q_ptr,
70
+ rel_ptr,
71
+ k_ptr,
72
+ v_ptr,
73
+ block_table_ptr,
74
+ cache_seqlens,
75
+ mid_out_ptr,
76
+ mid_lse_ptr,
77
+ softmax_scale,
78
+ stride_q_t,
79
+ stride_q_h,
80
+ stride_k_b,
81
+ stride_k_p,
82
+ stride_k_h,
83
+ stride_k_d,
84
+ stride_v_b,
85
+ stride_v_p,
86
+ stride_v_h,
87
+ stride_v_d,
88
+ stride_r_t,
89
+ stride_r_h,
90
+ stride_r_e,
91
+ stride_mo_t,
92
+ stride_mo_h,
93
+ stride_mo_s,
94
+ stride_ml_t,
95
+ stride_ml_h,
96
+ stride_ml_s,
97
+ stride_bt_b: tl.constexpr,
98
+ page_size: tl.constexpr,
99
+ window_left: tl.constexpr,
100
+ gqa_group_size: tl.constexpr,
101
+ num_q_heads: tl.constexpr,
102
+ head_dim: tl.constexpr,
103
+ rel_extent: tl.constexpr,
104
+ max_kv_splits: tl.constexpr,
105
+ BLOCK_D: tl.constexpr,
106
+ BLOCK_H: tl.constexpr,
107
+ BLOCK_N: tl.constexpr,
108
+):
109
+ pid_t = tl.program_id(0)
110
+ pid_h = tl.program_id(1)
111
+ pid_s = tl.program_id(2)
112
+
113
+ head_blocks_per_kv = tl.cdiv(gqa_group_size, BLOCK_H)
114
+ kv_head = pid_h // head_blocks_per_kv
115
+ valid_block_h: tl.constexpr = min(BLOCK_H, gqa_group_size)
116
+ off_h = pid_h * valid_block_h + tl.arange(0, BLOCK_H)
117
+ h_valid = (off_h < (pid_h + 1) * valid_block_h) & (off_h < num_q_heads)
118
+ off_d = tl.arange(0, BLOCK_D)
119
+ d_valid = off_d < head_dim
120
+
121
+ cache_len = tl.load(cache_seqlens + pid_t)
122
+ effective_len = (
123
+ tl.minimum(cache_len, window_left + 1) if window_left >= 0 else cache_len
124
+ )
125
+ kv_offset = cache_len - effective_len
126
+ split_len = (
127
+ tl.cdiv(tl.cdiv(effective_len, max_kv_splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV
128
+ )
129
+ split_start = split_len * pid_s
130
+ split_end = tl.minimum(split_start + split_len, effective_len)
131
+
132
+ q_offsets = pid_t * stride_q_t + off_h[:, None] * stride_q_h + off_d[None, :]
133
+ row_max = tl.full((BLOCK_H,), float("-inf"), dtype=tl.float32)
134
+ row_sum = tl.zeros((BLOCK_H,), dtype=tl.float32)
135
+ acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32)
136
+
137
+ if split_end > split_start:
138
+ q = tl.load(
139
+ q_ptr + q_offsets,
140
+ mask=h_valid[:, None] & d_valid[None, :],
141
+ other=0.0,
142
+ )
143
+ for start_n in range(split_start, split_end, BLOCK_N):
144
+ off_n = start_n + tl.arange(0, BLOCK_N)
145
+ n_valid = off_n < split_end
146
+ token_idx = kv_offset + off_n
147
+ logical_page = token_idx // page_size
148
+ page_offset = token_idx % page_size
149
+ physical_page = tl.load(
150
+ block_table_ptr + pid_t * stride_bt_b + logical_page,
151
+ mask=n_valid,
152
+ other=0,
153
+ )
154
+ k_offsets = (
155
+ physical_page[None, :].to(tl.int64) * stride_k_b
156
+ + page_offset[None, :] * stride_k_p
157
+ + kv_head * stride_k_h
158
+ + off_d[:, None] * stride_k_d
159
+ )
160
+ k = tl.load(
161
+ k_ptr + k_offsets,
162
+ mask=d_valid[:, None] & n_valid[None, :],
163
+ other=0.0,
164
+ )
165
+ scores = tl.dot(q, k.to(q.dtype)) * softmax_scale
166
+
167
+ rel_dist = cache_len - 1 - token_idx
168
+ rel_valid = (rel_dist >= 0) & (rel_dist < rel_extent)
169
+ rel_idx = tl.maximum(0, tl.minimum(rel_dist, rel_extent - 1))
170
+ rel_offsets = (
171
+ pid_t * stride_r_t
172
+ + off_h[:, None] * stride_r_h
173
+ + rel_idx[None, :] * stride_r_e
174
+ )
175
+ rel_bias = tl.load(
176
+ rel_ptr + rel_offsets,
177
+ mask=h_valid[:, None] & rel_valid[None, :] & n_valid[None, :],
178
+ other=0.0,
179
+ )
180
+ scores += rel_bias.to(tl.float32)
181
+ scores = tl.where(
182
+ h_valid[:, None] & n_valid[None, :],
183
+ scores,
184
+ float("-inf"),
185
+ )
186
+
187
+ v_offsets = (
188
+ physical_page[:, None].to(tl.int64) * stride_v_b
189
+ + page_offset[:, None] * stride_v_p
190
+ + kv_head * stride_v_h
191
+ + off_d[None, :] * stride_v_d
192
+ )
193
+ v = tl.load(
194
+ v_ptr + v_offsets,
195
+ mask=n_valid[:, None] & d_valid[None, :],
196
+ other=0.0,
197
+ )
198
+
199
+ next_max = tl.maximum(tl.max(scores, axis=1), row_max)
200
+ old_scale = tl.exp(row_max - next_max)
201
+ probs = tl.exp(scores - next_max[:, None])
202
+ acc *= old_scale[:, None]
203
+ acc += tl.dot(probs.to(v.dtype), v)
204
+ row_sum = row_sum * old_scale + tl.sum(probs, axis=1)
205
+ row_max = next_max
206
+
207
+ mid_out_offsets = (
208
+ pid_t * stride_mo_t
209
+ + off_h[:, None] * stride_mo_h
210
+ + pid_s * stride_mo_s
211
+ + off_d[None, :]
212
+ )
213
+ tl.store(
214
+ mid_out_ptr + mid_out_offsets,
215
+ acc / row_sum[:, None],
216
+ mask=h_valid[:, None] & d_valid[None, :],
217
+ )
218
+ mid_lse_offsets = (
219
+ pid_t * stride_ml_t + off_h * stride_ml_h + pid_s * stride_ml_s
220
+ )
221
+ tl.store(
222
+ mid_lse_ptr + mid_lse_offsets,
223
+ row_max + tl.log(row_sum),
224
+ mask=h_valid,
225
+ )
226
+
227
+
228
+@triton.jit
229
+def _split_kv_stage2(
230
+ mid_out_ptr,
231
+ mid_lse_ptr,
232
+ out_ptr,
233
+ cache_seqlens,
234
+ stride_mo_t,
235
+ stride_mo_h,
236
+ stride_mo_s,
237
+ stride_ml_t,
238
+ stride_ml_h,
239
+ stride_ml_s,
240
+ stride_o_t,
241
+ stride_o_h,
242
+ window_left: tl.constexpr,
243
+ head_dim: tl.constexpr,
244
+ max_kv_splits: tl.constexpr,
245
+ BLOCK_D: tl.constexpr,
246
+):
247
+ pid_t = tl.program_id(0)
248
+ pid_h = tl.program_id(1)
249
+
250
+ cache_len = tl.load(cache_seqlens + pid_t)
251
+ effective_len = (
252
+ tl.minimum(cache_len, window_left + 1) if window_left >= 0 else cache_len
253
+ )
254
+ split_len = (
255
+ tl.cdiv(tl.cdiv(effective_len, max_kv_splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV
256
+ )
257
+
258
+ off_d = tl.arange(0, BLOCK_D)
259
+ d_valid = off_d < head_dim
260
+ row_max = -float("inf")
261
+ row_sum = 0.0
262
+ acc = tl.zeros((BLOCK_D,), dtype=tl.float32)
263
+
264
+ for split_id in range(max_kv_splits):
265
+ split_start = split_len * split_id
266
+ split_end = tl.minimum(split_start + split_len, effective_len)
267
+ if split_end > split_start:
268
+ value = tl.load(
269
+ mid_out_ptr
270
+ + pid_t * stride_mo_t
271
+ + pid_h * stride_mo_h
272
+ + split_id * stride_mo_s
273
+ + off_d,
274
+ mask=d_valid,
275
+ other=0.0,
276
+ )
277
+ split_lse = tl.load(
278
+ mid_lse_ptr
279
+ + pid_t * stride_ml_t
280
+ + pid_h * stride_ml_h
281
+ + split_id * stride_ml_s
282
+ )
283
+ next_max = tl.maximum(split_lse, row_max)
284
+ old_scale = tl.exp(row_max - next_max)
285
+ split_scale = tl.exp(split_lse - next_max)
286
+ acc = acc * old_scale + value * split_scale
287
+ row_sum = row_sum * old_scale + split_scale
288
+ row_max = next_max
289
+
290
+ tl.store(
291
+ out_ptr + pid_t * stride_o_t + pid_h * stride_o_h + off_d,
292
+ acc / row_sum,
293
+ mask=d_valid,
294
+ )
295
+
296
+
297
+@torch.no_grad()
298
+def inkling_rel_attention_split_kv_decode(
299
+ q: torch.Tensor,
300
+ key_cache: torch.Tensor,
301
+ value_cache: torch.Tensor,
302
+ *,
303
+ block_table: torch.Tensor,
304
+ cache_seqlens: torch.Tensor,
305
+ softmax_scale: float,
306
+ window_left: int,
307
+ rel_extent: int,
308
+ rel_logits: torch.Tensor,
309
+ max_kv_len: int,
310
+ out: torch.Tensor,
311
+) -> torch.Tensor:
312
+ """Run split-KV relative attention for single-token decode."""
313
+ num_kv_heads = key_cache.shape[2]
314
+ gqa_group_size = q.shape[1] // num_kv_heads
315
+ max_kv_splits = decode_split_count(max_kv_len, window_left)
316
+ block_d = triton.next_power_of_2(q.shape[2])
317
+ block_h = min(8, gqa_group_size)
318
+
319
+ mid_out = torch.empty(
320
+ q.shape[0],
321
+ q.shape[1],
322
+ max_kv_splits,
323
+ q.shape[2],
324
+ dtype=torch.float32,
325
+ device=q.device,
326
+ )
327
+ mid_lse = torch.empty(
328
+ q.shape[0],
329
+ q.shape[1],
330
+ max_kv_splits,
331
+ dtype=torch.float32,
332
+ device=q.device,
333
+ )
334
+ stage1_grid = (
335
+ q.shape[0],
336
+ triton.cdiv(q.shape[1], block_h),
337
+ max_kv_splits,
338
+ )
339
+ _split_kv_stage1[stage1_grid](
340
+ q,
341
+ rel_logits,
342
+ key_cache,
343
+ value_cache,
344
+ block_table,
345
+ cache_seqlens,
346
+ mid_out,
347
+ mid_lse,
348
+ softmax_scale,
349
+ q.stride(0),
350
+ q.stride(1),
351
+ key_cache.stride(0),
352
+ key_cache.stride(1),
353
+ key_cache.stride(2),
354
+ key_cache.stride(3),
355
+ value_cache.stride(0),
356
+ value_cache.stride(1),
357
+ value_cache.stride(2),
358
+ value_cache.stride(3),
359
+ rel_logits.stride(0),
360
+ rel_logits.stride(1),
361
+ rel_logits.stride(2),
362
+ mid_out.stride(0),
363
+ mid_out.stride(1),
364
+ mid_out.stride(2),
365
+ mid_lse.stride(0),
366
+ mid_lse.stride(1),
367
+ mid_lse.stride(2),
368
+ block_table.stride(0),
369
+ page_size=key_cache.shape[1],
370
+ window_left=window_left,
371
+ gqa_group_size=gqa_group_size,
372
+ num_q_heads=q.shape[1],
373
+ head_dim=q.shape[2],
374
+ rel_extent=rel_extent,
375
+ max_kv_splits=max_kv_splits,
376
+ BLOCK_D=block_d,
377
+ BLOCK_H=block_h,
378
+ BLOCK_N=key_cache.shape[1],
379
+ num_warps=4,
380
+ num_stages=1,
381
+ )
382
+ stage2_grid = (q.shape[0], q.shape[1])
383
+ _split_kv_stage2[stage2_grid](
384
+ mid_out,
385
+ mid_lse,
386
+ out,
387
+ cache_seqlens,
388
+ mid_out.stride(0),
389
+ mid_out.stride(1),
390
+ mid_out.stride(2),
391
+ mid_lse.stride(0),
392
+ mid_lse.stride(1),
393
+ mid_lse.stride(2),
394
+ out.stride(0),
395
+ out.stride(1),
396
+ window_left=window_left,
397
+ head_dim=q.shape[2],
398
+ max_kv_splits=max_kv_splits,
399
+ BLOCK_D=block_d,
400
+ num_warps=4,
401
+ num_stages=2,
402
+ )
403
+ return out
vllm/models/inkling/amd/ops/sconv.py
new
+290
@@ -0,0 +1,290 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inkling short-convolution kernels backed by a paged sliding-window state cache.
4
+
5
+Each layer's 4 conv streams (K, V, attn-output, mlp-output) share one paged KV
6
+cache ``[num_blocks, H, N, D]`` (head-major; see ``sconv_swa_attn.py``). A stream
7
+occupies the contiguous D-sub-range ``[off_s, off_s + ws)`` across all ``H``
8
+heads, so its flat per-token width is ``H * ws`` and that is the conv channel
9
+dim. The cache stores the conv *input* at every absolute position.
10
+
11
+``fused_sconv`` is the single-launch path used by the model: per token it writes
12
+the current input to its slot and convolves the ``W`` taps ending at its
13
+absolute position. A tap landing inside the current forward is read from the
14
+immutable input ``x`` (row ``src - pos + pid_t``); only pre-forward taps are read
15
+from the paged cache (window position ``src`` -> physical block via
16
+``block_table[req, src // N]``). The just-written slot is never read back this
17
+step, so there is no write/read hazard within or across programs -- which is
18
+why this needs no decode-vs-prefill split and is valid for prefill / decode /
19
+spec alike.
20
+
21
+All kernels address the cache purely by ``(slot, absolute_position)`` and
22
+allocate nothing inside the captured region; their grids depend only on the
23
+token count (``fused_sconv`` on a fixed token/channel tiling), so the same
24
+path replays correctly under eager, breakable PIECEWISE, and FULL cudagraphs
25
+without any data-dependent shape or branch.
26
+"""
27
+
28
+from __future__ import annotations
29
+
30
+import torch
31
+
32
+from vllm.triton_utils import tl, triton
33
+
34
+
35
+@triton.jit
36
+def _fused_sconv_kernel(
37
+ x_ptr, # [T, H*WS] head-major current-token inputs (also residual)
38
+ cache_ptr, # [num_blocks, H, N, D] paged (page-strided view)
39
+ weight_ptr, # [H*WS, W]
40
+ out_ptr, # [T, H*WS]
41
+ pos_ptr, # [T] int64 absolute position per token
42
+ seq_idx_ptr, # [T] int32 token -> batch request
43
+ slot_ptr, # [T] int64 flat slot (block*N + blk_off); < 0 => PAD
44
+ block_table_ptr, # [num_reqs, max_blocks] int32 block_table
45
+ qstart_ptr, # [T] int32 first x-row of the token's request
46
+ T, # num tokens
47
+ stride_x_t,
48
+ stride_c_blk,
49
+ stride_c_h,
50
+ stride_c_n,
51
+ stride_c_d,
52
+ stride_w_d,
53
+ stride_w_w,
54
+ stride_bt_r,
55
+ MAX_BLOCKS,
56
+ N, # block_size
57
+ W: tl.constexpr,
58
+ USE_SILU: tl.constexpr,
59
+ USE_RESIDUAL: tl.constexpr,
60
+ OFF_S: tl.constexpr,
61
+ WS: tl.constexpr,
62
+ H: tl.constexpr,
63
+ BT: tl.constexpr,
64
+ BLOCK_C: tl.constexpr,
65
+):
66
+ # Each program owns a [BT tokens, BLOCK_C channels] tile. The flat channel
67
+ # index packs all H heads head-major (head = c // WS, in-stream offset =
68
+ # c % WS), so one program spans heads -- no per-head launch and no
69
+ # next_power_of_2(WS) lane waste.
70
+ pid_t = tl.program_id(0)
71
+ pid_c = tl.program_id(1)
72
+ toff = pid_t * BT + tl.arange(0, BT) # [BT] token rows
73
+ coff = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) # [BLOCK_C] flat channels
74
+ C = H * WS
75
+ t_mask = toff < T
76
+ c_mask = coff < C
77
+ head = tl.minimum(coff // WS, H - 1) # clamp keeps masked lanes in-buffer
78
+ cd = OFF_S + coff % WS # cache D-index of the channel's stream slot
79
+
80
+ slot = tl.load(slot_ptr + toff, mask=t_mask, other=-1) # [BT]
81
+ valid = slot >= 0
82
+ pos = tl.load(pos_ptr + toff, mask=t_mask, other=0)
83
+ req = tl.load(seq_idx_ptr + toff, mask=t_mask, other=0)
84
+ qstart = tl.load(qstart_ptr + toff, mask=t_mask, other=0)
85
+
86
+ tc_mask = t_mask[:, None] & c_mask[None, :]
87
+
88
+ # 1) Insert each token's input into its paged slot (skip PAD rows).
89
+ xv = tl.load(x_ptr + toff[:, None] * stride_x_t + coff[None, :], mask=tc_mask)
90
+ safe_slot = tl.maximum(slot, 0)
91
+ dst = (
92
+ cache_ptr
93
+ + (safe_slot // N)[:, None] * stride_c_blk
94
+ + head[None, :] * stride_c_h
95
+ + (safe_slot % N)[:, None] * stride_c_n
96
+ + cd[None, :] * stride_c_d
97
+ )
98
+ tl.store(dst, xv, mask=tc_mask & valid[:, None])
99
+
100
+ # 2) Convolve the W taps ending at each token's `pos`. Each tap is read from
101
+ # `x` if it falls inside this forward (row >= the request's first row), else
102
+ # from the paged cache. Exactly one source is unmasked per tap, so we sum both.
103
+ acc = tl.zeros([BT, BLOCK_C], dtype=tl.float32)
104
+ for iw in tl.static_range(W):
105
+ src = pos - (W - 1) + iw # [BT] absolute window position
106
+ row = toff - (W - 1) + iw # [BT] x-row of `src` (== src - pos + token)
107
+ in_win = valid & (src >= 0)
108
+ intra = in_win & (row >= qstart)
109
+ cached = in_win & (row < qstart)
110
+ # intra-forward tap: read the immutable input x (never the slot we just
111
+ # wrote), so there is no write/read hazard.
112
+ safe_row = tl.maximum(row, 0)
113
+ xt = tl.load(
114
+ x_ptr + safe_row[:, None] * stride_x_t + coff[None, :],
115
+ mask=c_mask[None, :] & intra[:, None],
116
+ other=0.0,
117
+ ).to(tl.float32)
118
+ # pre-forward tap: read from the paged cache via the block table. Clamp
119
+ # addressing terms; the load is masked off when out of window.
120
+ safe_src = tl.maximum(src, 0)
121
+ safe_lblk = tl.minimum(safe_src // N, MAX_BLOCKS - 1)
122
+ blk = tl.load(
123
+ block_table_ptr + req * stride_bt_r + safe_lblk, mask=cached, other=0
124
+ ).to(tl.int64)
125
+ cbase = (
126
+ cache_ptr
127
+ + blk[:, None] * stride_c_blk
128
+ + head[None, :] * stride_c_h
129
+ + (safe_src % N)[:, None] * stride_c_n
130
+ + cd[None, :] * stride_c_d
131
+ )
132
+ cv = tl.load(cbase, mask=c_mask[None, :] & cached[:, None], other=0.0).to(
133
+ tl.float32
134
+ )
135
+ wv = tl.load(
136
+ weight_ptr + coff * stride_w_d + iw * stride_w_w, mask=c_mask, other=0.0
137
+ ).to(tl.float32)
138
+ acc += (xt + cv) * wv[None, :]
139
+
140
+ if USE_SILU:
141
+ acc = acc * tl.sigmoid(acc)
142
+ if USE_RESIDUAL:
143
+ acc += xv.to(tl.float32)
144
+
145
+ tl.store(
146
+ out_ptr + toff[:, None] * stride_x_t + coff[None, :],
147
+ acc.to(out_ptr.dtype.element_ty),
148
+ mask=tc_mask,
149
+ )
150
+
151
+
152
+def fused_sconv(
153
+ x: torch.Tensor, # [T, H*ws] head-major current-token inputs
154
+ weight: torch.Tensor, # [H*ws, W]
155
+ cache: torch.Tensor, # [num_blocks, H, N, D] paged
156
+ positions: torch.Tensor, # [T] int64 absolute position per token
157
+ block_table: torch.Tensor, # [num_reqs, max_blocks] int32
158
+ seq_idx: torch.Tensor, # [T] int32 token -> batch request
159
+ slot_mapping: torch.Tensor, # [T] int64 flat slot (PAD = -1 => skip)
160
+ query_start: torch.Tensor, # [T] int32 first x-row of the token's request
161
+ off_s: int,
162
+ ws: int,
163
+ block_size: int,
164
+ activation: str | None = None,
165
+ use_residual: bool = True,
166
+) -> torch.Tensor:
167
+ """Single-launch insert + depthwise causal conv1d over the paged cache.
168
+
169
+ Reads same-forward taps from ``x`` and pre-forward taps from the cache, so
170
+ it is race-free in one launch for prefill / decode / spec and cudagraph-safe
171
+ under eager / piecewise / full capture.
172
+ """
173
+ T = x.shape[0]
174
+ out = torch.empty_like(x)
175
+ if T == 0:
176
+ return out
177
+ assert x.is_contiguous()
178
+ assert cache.stride(3) == 1, "cache D-dim must be contiguous"
179
+ H = cache.shape[1]
180
+ W = weight.shape[1]
181
+ C = H * ws # flat conv channel dim (all heads, head-major)
182
+ # Tile BT tokens x BLOCK_C channels per program: enough work per CTA to
183
+ # amortize the per-token addressing, while keeping the grid large for
184
+ # prefill. BLOCK_C spans heads so there is no per-head launch. A ~2K-element
185
+ # tile at 4 warps measured best on Blackwell; larger tiles spill registers.
186
+ BLOCK_C = min(triton.next_power_of_2(C), 256)
187
+ BT = 8
188
+ grid = (triton.cdiv(T, BT), triton.cdiv(C, BLOCK_C))
189
+ _fused_sconv_kernel[grid](
190
+ x,
191
+ cache,
192
+ weight,
193
+ out,
194
+ positions,
195
+ seq_idx,
196
+ slot_mapping,
197
+ block_table,
198
+ query_start,
199
+ T,
200
+ x.stride(0),
201
+ cache.stride(0),
202
+ cache.stride(1),
203
+ cache.stride(2),
204
+ cache.stride(3),
205
+ weight.stride(0),
206
+ weight.stride(1),
207
+ block_table.stride(0),
208
+ block_table.shape[1],
209
+ block_size,
210
+ W=W,
211
+ USE_SILU=activation in ("silu", "swish"),
212
+ USE_RESIDUAL=use_residual,
213
+ OFF_S=off_s,
214
+ WS=ws,
215
+ H=H,
216
+ BT=BT,
217
+ BLOCK_C=BLOCK_C,
218
+ num_warps=4,
219
+ )
220
+ return out
221
+
222
+
223
+@triton.jit
224
+def _seq_metadata_kernel(
225
+ qsl_ptr, # [num_reqs + 1] int32 cumulative query start rows
226
+ seq_idx_ptr, # [T] int32 out: token -> owning request
227
+ query_start_ptr, # [T] int32 out: first x-row of the token's request
228
+ num_reqs,
229
+ num_actual_tokens,
230
+ num_padded_tokens,
231
+ n_iters, # ceil(log2(num_reqs)): binary-search depth
232
+ BLOCK: tl.constexpr,
233
+):
234
+ offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
235
+ tok = offs.to(tl.int32)
236
+ # Largest j in [0, num_reqs) with qsl[j] <= tok.
237
+ lo = tl.zeros([BLOCK], tl.int32)
238
+ hi = tl.full([BLOCK], num_reqs - 1, tl.int32)
239
+ for _ in range(n_iters):
240
+ mid = (lo + hi + 1) // 2
241
+ below = tl.load(qsl_ptr + mid) <= tok
242
+ lo = tl.where(below, mid, lo)
243
+ hi = tl.where(below, hi, mid - 1)
244
+ actual = offs < num_actual_tokens
245
+ padded = offs < num_padded_tokens
246
+ query_start = tl.load(qsl_ptr + lo)
247
+ tl.store(seq_idx_ptr + offs, tl.where(actual, lo, 0), mask=padded)
248
+ tl.store(
249
+ query_start_ptr + offs,
250
+ tl.where(actual, query_start, 0),
251
+ mask=padded,
252
+ )
253
+
254
+
255
+def sconv_seq_metadata(
256
+ query_start_loc: torch.Tensor,
257
+ num_reqs: int,
258
+ num_actual_tokens: int,
259
+ seq_idx_out: torch.Tensor,
260
+ query_start_out: torch.Tensor,
261
+ num_padded_tokens: int | None = None,
262
+) -> None:
263
+ """Fill static per-token seq_idx / query_start buffers in one launch.
264
+
265
+ Replaces the arange + searchsorted + clamp + gather + 2x copy chain of the
266
+ sconv metadata build with a single kernel writing both persistent buffers.
267
+ Padded rows are filled with zero and must have ``slot_mapping == -1``.
268
+ """
269
+ if num_padded_tokens is None:
270
+ num_padded_tokens = num_actual_tokens
271
+ if num_padded_tokens < num_actual_tokens:
272
+ raise ValueError("num_padded_tokens must cover all actual tokens")
273
+ if num_padded_tokens > seq_idx_out.shape[0]:
274
+ raise ValueError("seq_idx_out is too small for the padded token count")
275
+ if num_padded_tokens > query_start_out.shape[0]:
276
+ raise ValueError("query_start_out is too small for the padded token count")
277
+
278
+ BLOCK = 256
279
+ n_iters = (num_reqs - 1).bit_length()
280
+ grid = (triton.cdiv(num_padded_tokens, BLOCK),)
281
+ _seq_metadata_kernel[grid](
282
+ query_start_loc,
283
+ seq_idx_out,
284
+ query_start_out,
285
+ num_reqs,
286
+ num_actual_tokens,
287
+ num_padded_tokens,
288
+ n_iters,
289
+ BLOCK=BLOCK,
290
+ )
vllm/models/inkling/amd/ops/silu_and_mul.py
new
+197
@@ -0,0 +1,197 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""SwiGLU kernels for the Inkling MLP layers.
4
+
5
+``silu_and_mul_triton``: SiLU-and-mul over the checkpoint's interleaved
6
+fused gate/up layout (dense MLP). ``sink_silu_mul_epilogue``: the sink-expert
7
+variant with the per-expert dequant scale and per-token gamma fused in.
8
+"""
9
+
10
+import torch
11
+
12
+from vllm.triton_utils import tl, triton
13
+
14
+
15
+@triton.jit(do_not_specialize=["M"])
16
+def _silu_and_mul_triton_kernel(
17
+ gateup_out_ptr,
18
+ down_inp_ptr,
19
+ M,
20
+ N: tl.constexpr,
21
+ GRID_SIZE: tl.constexpr,
22
+ NUM_STAGES: tl.constexpr,
23
+ BLOCK_SIZE_M: tl.constexpr,
24
+ BLOCK_SIZE_N: tl.constexpr,
25
+ EVEN_N: tl.constexpr,
26
+ INT64_INDEX: tl.constexpr,
27
+):
28
+ start_pid = tl.program_id(0)
29
+ if INT64_INDEX:
30
+ start_pid = start_pid.to(tl.int64)
31
+ M = M.to(tl.int64)
32
+
33
+ NUM_BLOCKS_N: tl.constexpr = tl.cdiv(N, BLOCK_SIZE_N)
34
+ num_blocks_mn = tl.cdiv(M, BLOCK_SIZE_M) * NUM_BLOCKS_N
35
+
36
+ for pid in tl.range(start_pid, num_blocks_mn, GRID_SIZE, num_stages=NUM_STAGES):
37
+ pid_m = pid // NUM_BLOCKS_N
38
+ pid_n = pid % NUM_BLOCKS_N
39
+
40
+ offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
41
+ offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
42
+ mask_m = offs_m < M
43
+ mask_n = offs_n < N
44
+
45
+ # Interleaved fused gate/up: [g0, u0, g1, u1, ...].
46
+ mask_offs_2n = pid_n * BLOCK_SIZE_N + tl.arange(0, 2 * BLOCK_SIZE_N) // 2
47
+ tl.static_assert(BLOCK_SIZE_N % 8 == 0, f"{BLOCK_SIZE_N=}")
48
+ mask_2n = mask_offs_2n < N
49
+ mask_2n = tl.max_constancy(mask_2n, [16])
50
+
51
+ offs_2n = pid_n * 2 * BLOCK_SIZE_N + tl.arange(0, 2 * BLOCK_SIZE_N)
52
+ offs_m2n = offs_m[:, None] * N * 2 + offs_2n[None, :]
53
+
54
+ if EVEN_N or pid_n * BLOCK_SIZE_N + BLOCK_SIZE_N <= N:
55
+ gateup_out = tl.load(
56
+ gateup_out_ptr + offs_m2n, mask=mask_m[:, None], other=0.0
57
+ )
58
+ else:
59
+ mask_m2n = mask_m[:, None] & mask_2n[None, :]
60
+ gateup_out = tl.load(gateup_out_ptr + offs_m2n, mask=mask_m2n, other=0.0)
61
+
62
+ gate_out, up_out = tl.split(
63
+ tl.reshape(gateup_out, (BLOCK_SIZE_M, BLOCK_SIZE_N, 2))
64
+ )
65
+ gate_out = gate_out.to(tl.float32)
66
+ up_out = up_out.to(tl.float32)
67
+
68
+ down_inp = gate_out * tl.sigmoid(gate_out) * up_out
69
+
70
+ mask_mn = mask_m[:, None] if EVEN_N else mask_m[:, None] & mask_n[None, :]
71
+ offs_mn = offs_m[:, None] * N + offs_n[None, :]
72
+ tl.store(down_inp_ptr + offs_mn, down_inp, mask=mask_mn)
73
+
74
+
75
+def silu_and_mul_triton(gateup_output: torch.Tensor) -> torch.Tensor:
76
+ """SiLU-and-mul for the interleaved fused gate/up layout.
77
+
78
+ Adapted from ``inkling_kernels.activation.silu_and_mul_fwd`` (without MXFP).
79
+ """
80
+ assert gateup_output.is_contiguous(), (
81
+ f"{gateup_output.shape=} {gateup_output.stride()=}"
82
+ )
83
+ assert gateup_output.ndim == 2, f"{gateup_output.shape=}"
84
+
85
+ M = gateup_output.shape[0]
86
+ hidden_size = gateup_output.shape[1]
87
+ assert hidden_size % 2 == 0, f"{hidden_size=}"
88
+ N = hidden_size // 2
89
+
90
+ down_input = torch.empty(
91
+ (M, N), device=gateup_output.device, dtype=gateup_output.dtype
92
+ )
93
+ if M == 0:
94
+ return down_input
95
+
96
+ BLOCK_SIZE_N = max(8, min(256, triton.next_power_of_2(N)))
97
+ if M <= 1:
98
+ BLOCK_SIZE_M = 4
99
+ elif M <= 256:
100
+ BLOCK_SIZE_M = 2
101
+ elif M < 4096:
102
+ BLOCK_SIZE_M = 4
103
+ else:
104
+ BLOCK_SIZE_M = 16
105
+ BLOCK_SIZE_N = max(8, min(128, triton.next_power_of_2(N)))
106
+ max_grid_size = triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N)
107
+ num_sms = torch.cuda.get_device_properties(
108
+ gateup_output.device
109
+ ).multi_processor_count
110
+ grid_size = min(num_sms * 4, max_grid_size)
111
+
112
+ _silu_and_mul_triton_kernel[(grid_size,)](
113
+ gateup_out_ptr=gateup_output,
114
+ down_inp_ptr=down_input,
115
+ M=M,
116
+ N=N,
117
+ GRID_SIZE=grid_size,
118
+ NUM_STAGES=1,
119
+ BLOCK_SIZE_M=BLOCK_SIZE_M,
120
+ BLOCK_SIZE_N=BLOCK_SIZE_N,
121
+ EVEN_N=N % BLOCK_SIZE_N == 0,
122
+ INT64_INDEX=gateup_output.nbytes >= 2**31,
123
+ num_warps=8,
124
+ )
125
+
126
+ return down_input
127
+
128
+
129
+@triton.jit(do_not_specialize=["T"])
130
+def _sink_epilogue_kernel(
131
+ raw_ptr, # [T, S * 2F] gemm1 output, interleaved g/u pairs per expert block
132
+ alpha_ptr, # [S] fp32 per-expert pre-SiLU dequant scale
133
+ gamma_ptr, # [T, S] fp32 per-token sink weights (may be strided)
134
+ ratio_ptr, # [S] fp32 per-expert post-SiLU scale (gemm2 alpha ratio)
135
+ out_ptr, # [T, S * F] output
136
+ T,
137
+ stride_raw_0,
138
+ stride_gamma_0,
139
+ F: tl.constexpr,
140
+ S: tl.constexpr,
141
+ BLOCK_F: tl.constexpr,
142
+):
143
+ pid_t = tl.program_id(0).to(tl.int64)
144
+ pid_sf = tl.program_id(1)
145
+ if pid_t >= T:
146
+ return
147
+ s = pid_sf // (F // BLOCK_F)
148
+ offs_f = (pid_sf % (F // BLOCK_F)) * BLOCK_F + tl.arange(0, BLOCK_F)
149
+
150
+ base = pid_t * stride_raw_0 + s * 2 * F
151
+ gate = tl.load(raw_ptr + base + 2 * offs_f).to(tl.float32)
152
+ up = tl.load(raw_ptr + base + 2 * offs_f + 1).to(tl.float32)
153
+ alpha = tl.load(alpha_ptr + s)
154
+ weight = tl.load(gamma_ptr + pid_t * stride_gamma_0 + s) * tl.load(ratio_ptr + s)
155
+
156
+ gate *= alpha
157
+ up *= alpha
158
+ h = gate * tl.sigmoid(gate) * up * weight
159
+ tl.store(out_ptr + pid_t * (S * F) + s * F + offs_f, h)
160
+
161
+
162
+def sink_silu_mul_epilogue(
163
+ raw: torch.Tensor, # [T, S * 2F] gemm1 output (interleaved gate/up rows)
164
+ alphas: torch.Tensor, # [S] fp32
165
+ gammas: torch.Tensor, # [T, S] fp32
166
+ ratios: torch.Tensor, # [S] fp32
167
+ n_experts: int,
168
+ out_dtype: torch.dtype,
169
+) -> torch.Tensor:
170
+ """Fused sink-expert epilogue: silu(g * a_e) * (u * a_e) * (gamma * r_e).
171
+
172
+ One kernel replaces the per-expert dequant column scale, the SwiGLU, and
173
+ the per-token gamma multiply between the two sink GEMMs.
174
+ """
175
+ tokens = raw.shape[0]
176
+ f = raw.shape[1] // (2 * n_experts)
177
+ out = torch.empty((tokens, n_experts * f), device=raw.device, dtype=out_dtype)
178
+ if tokens == 0:
179
+ return out
180
+ # raw may be a column-slice of a padded GEMM output (rows strided).
181
+ assert raw.stride(1) == 1 and gammas.stride(1) == 1
182
+ # Largest power-of-two divisor of f (f = 768 -> 256), capped at 512.
183
+ block_f = min(512, f & (-f))
184
+ _sink_epilogue_kernel[(tokens, n_experts * (f // block_f))](
185
+ raw,
186
+ alphas,
187
+ gammas,
188
+ ratios,
189
+ out,
190
+ tokens,
191
+ raw.stride(0),
192
+ gammas.stride(0),
193
+ F=f,
194
+ S=n_experts,
195
+ BLOCK_F=block_f,
196
+ )
197
+ return out
vllm/models/inkling/amd/sconv_swa_attn.py
new
+226
@@ -0,0 +1,226 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inkling short-conv state managed as a sliding-window KV cache.
4
+
5
+Each decoder layer owns one ``InklingConvState`` (an ``AttentionLayerBase``) that
6
+emits a single ``SlidingWindowSpec`` for the layer's 4 sconv streams (K, V,
7
+attn-output, mlp-output), packed head-major into one block:
8
+
9
+ H = num_kv_heads (per-rank), N = block_size = sconv_kernel_size,
10
+ D = head_dim(K) + head_dim(V) + hidden/H(attn) + hidden/H(mlp)
11
+
12
+``D`` is TP-invariant; per rank we store ``H/TP`` heads of width ``D``. The conv
13
+reads/writes this cache out-of-band via a custom backend; the (smaller) conv page
14
+is padded up to the uniform attention page by ``unify_kv_cache_spec_page_size``.
15
+"""
16
+
17
+from __future__ import annotations
18
+
19
+from dataclasses import dataclass
20
+from typing import ClassVar
21
+
22
+import torch
23
+from torch import nn
24
+
25
+from vllm.config import VllmConfig, get_current_vllm_config
26
+from vllm.distributed import get_tensor_model_parallel_world_size
27
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
28
+from vllm.v1.attention.backend import (
29
+ AttentionBackend,
30
+ AttentionCGSupport,
31
+ AttentionMetadata,
32
+ AttentionMetadataBuilder,
33
+ CommonAttentionMetadata,
34
+)
35
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec, SlidingWindowSpec
36
+
37
+from .ops.sconv import sconv_seq_metadata
38
+
39
+# Stream order within the per-head packed D (== contiguous sub-ranges).
40
+_K, _V, _ATTN, _MLP = 0, 1, 2, 3
41
+
42
+
43
+@dataclass
44
+class InklingSconvMetadata(AttentionMetadata):
45
+ block_table: torch.Tensor # [num_reqs, max_blocks] physical blocks per req
46
+ slot_mapping: torch.Tensor # [T] int64 flat slot of each token (-1 => skip)
47
+ seq_idx: torch.Tensor # [T] int32 token -> batch request
48
+ query_start: torch.Tensor # [T] int32 first x-row of each token's request
49
+
50
+
51
+class InklingSconvMetadataBuilder(AttentionMetadataBuilder[InklingSconvMetadata]):
52
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
53
+
54
+ def __init__(
55
+ self,
56
+ kv_cache_spec: AttentionSpec,
57
+ layer_names: list[str],
58
+ vllm_config: VllmConfig,
59
+ device: torch.device,
60
+ ) -> None:
61
+ super().__init__(kv_cache_spec, layer_names, vllm_config, device)
62
+ assert isinstance(kv_cache_spec, SlidingWindowSpec)
63
+ # Persistent per-token buffers for CUDA graph capture.
64
+ max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
65
+ self.seq_idx_buffer = torch.empty(
66
+ max_num_tokens, dtype=torch.int32, device=device
67
+ )
68
+ self.query_start_buffer = torch.empty(
69
+ max_num_tokens, dtype=torch.int32, device=device
70
+ )
71
+
72
+ def build(
73
+ self,
74
+ common_prefix_len: int,
75
+ common_attn_metadata: CommonAttentionMetadata,
76
+ fast_build: bool = False,
77
+ ) -> InklingSconvMetadata:
78
+ num_reqs = common_attn_metadata.num_reqs
79
+ num_actual_tokens = int(common_attn_metadata.query_start_loc_cpu[-1])
80
+ num_padded_tokens = common_attn_metadata.slot_mapping.shape[0]
81
+ assert num_padded_tokens >= num_actual_tokens
82
+
83
+ # Per-token seq_idx (owning request) and query_start (first x-row of
84
+ # that request; the fused kernel uses it to tell same-forward taps,
85
+ # read from x, from pre-forward taps, read from cache) in one launch.
86
+ sconv_seq_metadata(
87
+ common_attn_metadata.query_start_loc,
88
+ num_reqs,
89
+ num_actual_tokens,
90
+ self.seq_idx_buffer,
91
+ self.query_start_buffer,
92
+ num_padded_tokens,
93
+ )
94
+
95
+ return InklingSconvMetadata(
96
+ block_table=common_attn_metadata.block_table_tensor,
97
+ slot_mapping=common_attn_metadata.slot_mapping[:num_padded_tokens],
98
+ seq_idx=self.seq_idx_buffer[:num_padded_tokens],
99
+ query_start=self.query_start_buffer[:num_padded_tokens],
100
+ )
101
+
102
+
103
+class InklingSconvBackend(AttentionBackend):
104
+ """Custom dummy backend for the sconv sliding-window cache management."""
105
+
106
+ @staticmethod
107
+ def get_name() -> str:
108
+ return "INKLING_SCONV_SWA"
109
+
110
+ @classmethod
111
+ def indexes_kv_by_block_stride(cls) -> bool:
112
+ # num_blocks is the outermost dim (HND, see get_kv_cache_shape), so the
113
+ # padded conv page is read through a strided view.
114
+ return True
115
+
116
+ @staticmethod
117
+ def get_kv_cache_shape(
118
+ num_blocks: int,
119
+ block_size: int,
120
+ num_kv_heads: int,
121
+ head_size: int,
122
+ cache_dtype_str: str = "auto",
123
+ ) -> tuple[int, ...]:
124
+ # HND, num-blocks-first, head-major: [num_blocks, H, N, D].
125
+ return (num_blocks, num_kv_heads, block_size, head_size)
126
+
127
+ @staticmethod
128
+ def get_kv_cache_stride_order(
129
+ include_num_layers_dimension: bool = False,
130
+ ) -> tuple[int, ...]:
131
+ # Identity: physical layout == logical [num_blocks, H, N, D].
132
+ if include_num_layers_dimension:
133
+ return (0, 1, 2, 3, 4)
134
+ return (0, 1, 2, 3)
135
+
136
+ @staticmethod
137
+ def get_impl_cls():
138
+ raise NotImplementedError(
139
+ "InklingSconvBackend has no attention impl; the conv runs out-of-band."
140
+ )
141
+
142
+ @staticmethod
143
+ def get_builder_cls() -> type[InklingSconvMetadataBuilder]:
144
+ return InklingSconvMetadataBuilder
145
+
146
+
147
+class InklingConvState(nn.Module, AttentionLayerBase):
148
+ """Per-decoder-layer owner emitting one sliding-window conv-state spec."""
149
+
150
+ def __init__(
151
+ self,
152
+ *,
153
+ num_kv_heads: int,
154
+ head_dim: int,
155
+ hidden_size: int,
156
+ kernel_size: int,
157
+ prefix: str,
158
+ ) -> None:
159
+ super().__init__()
160
+ self.prefix = prefix
161
+ # Bound to the manager-allocated paged cache by bind_kv_cache; a
162
+ # placeholder until then. Read out-of-band by InklingShortConv.
163
+ self.kv_cache = torch.tensor([])
164
+ tp_size = get_tensor_model_parallel_world_size()
165
+ # Guardrails for the conv-state layout below; only these are exercised.
166
+ # tp_size <= num_kv_heads keeps >=1 whole KV head per rank (no
167
+ # replication/clamping), so the per-head width stays TP-invariant.
168
+ assert tp_size <= num_kv_heads, (
169
+ f"sconv SWA cache supports tp_size <= num_kv_heads ({num_kv_heads}), "
170
+ f"got {tp_size}"
171
+ )
172
+ # Per-rank head count; D is TP-invariant (K/V heads and the hidden
173
+ # chunk both scale 1/TP together). The attn-/mlp-output sconv streams
174
+ # are hidden-sharded: each rank owns its H/tp chunk (the sublayer
175
+ # outputs are reduce-scattered / all-gathered around the conv).
176
+ self.num_kv_heads = num_kv_heads // tp_size
177
+ hidden_per_head = hidden_size // num_kv_heads
178
+ # Packed per-head width: K + V + attn-output chunk + mlp-output chunk,
179
+ # padded to a power of two so every layer's conv page is the same size
180
+ # and an exact multiple of the attention page (the page unifier then
181
+ # scales attention block sizes instead of padding).
182
+ raw_head_size = 2 * head_dim + 2 * hidden_per_head
183
+ self.head_size = 1 << (raw_head_size - 1).bit_length()
184
+ self.sliding_window = kernel_size
185
+ self.block_size = kernel_size
186
+ # Per-head D-sub-range (offset, width) for each stream. Streams share
187
+ # the cache; each writes/reads its own width across all H heads.
188
+ self.stream_ranges: tuple[tuple[int, int], ...] = (
189
+ (0, head_dim), # _K
190
+ (head_dim, head_dim), # _V
191
+ (2 * head_dim, hidden_per_head), # _ATTN
192
+ (2 * head_dim + hidden_per_head, hidden_per_head), # _MLP
193
+ )
194
+ vllm_config = get_current_vllm_config()
195
+ self._dtype = vllm_config.model_config.dtype
196
+ assert self._dtype == torch.bfloat16, (
197
+ f"sconv SWA cache supports bfloat16 only, got {self._dtype}"
198
+ )
199
+ # Register in the forward context so the runner enumerates this owner as
200
+ # an attention-like layer (get_kv_cache_spec / get_attn_backend).
201
+ compilation_config = vllm_config.compilation_config
202
+ if prefix in compilation_config.static_forward_context:
203
+ raise ValueError(f"Duplicate layer name: {prefix}")
204
+ compilation_config.static_forward_context[prefix] = self
205
+
206
+ def forward(self): ...
207
+
208
+ @property
209
+ def cache_block_size(self) -> int:
210
+ """Return the block size used by cache metadata and physical indexing."""
211
+ if self.kv_cache.numel() > 0:
212
+ return self.kv_cache.shape[2]
213
+ return self.block_size
214
+
215
+ def get_attn_backend(self) -> type[AttentionBackend]:
216
+ return InklingSconvBackend
217
+
218
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
219
+ return SlidingWindowSpec(
220
+ block_size=self.block_size,
221
+ num_kv_heads=self.num_kv_heads,
222
+ head_size=self.head_size,
223
+ head_size_v=0, # all 4 streams packed into head_size
224
+ dtype=self._dtype,
225
+ sliding_window=self.sliding_window,
226
+ )
vllm/models/inkling/amd/short_conv.py
new
+98
@@ -0,0 +1,98 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inkling short convolution: depthwise causal conv1d (+ residual) over a paged
4
+sliding-window conv-state cache.
5
+
6
+Each decoder layer owns one ``InklingConvState`` (``sconv_swa_attn.py``) holding the
7
+manager-allocated paged cache for the layer's 4 sconv streams (K, V, attn-output,
8
+mlp-output), packed head-major into one block. Each ``InklingShortConv`` is a
9
+stateless weight + kernel launcher that, per forward (positions-addressed, the
10
+same path for prefill / decode / mixed), inserts the current tokens' inputs
11
+into their paged slot and convolves each token against the ``W`` taps ending
12
+at its absolute position, reading pre-forward window positions out of the
13
+paged cache via the block table.
14
+
15
+Per-forward metadata (``block_table`` / ``slot_mapping`` / ``seq_idx`` /
16
+``query_start``) is built once by ``InklingSconvMetadataBuilder`` and published under
17
+the owner's prefix in the forward context; the absolute ``positions`` are
18
+threaded in from the model. The insert + conv run in a single ``fused_sconv``
19
+launch (same path for prefill / decode / mixed / spec). All inputs are
20
+fixed-address persistent buffers and the grid is fixed, so the conv replays
21
+correctly under eager, PIECEWISE, and FULL cudagraphs.
22
+"""
23
+
24
+from __future__ import annotations
25
+
26
+import torch
27
+from torch import nn
28
+from torch.nn.parameter import Parameter
29
+
30
+from vllm.distributed import get_tensor_model_parallel_rank
31
+from vllm.forward_context import get_forward_context
32
+from vllm.model_executor.utils import set_weight_attrs
33
+
34
+from .ops import fused_sconv
35
+from .sconv_swa_attn import InklingConvState, InklingSconvMetadata
36
+
37
+
38
+class InklingShortConv(nn.Module):
39
+ def __init__(
40
+ self, dim: int, kernel_size: int, owner: InklingConvState, stream_idx: int
41
+ ) -> None:
42
+ super().__init__()
43
+ self.dim = dim
44
+ self.kernel_size = kernel_size
45
+ self.owner = owner
46
+ self.stream_idx = stream_idx
47
+ self.tp_rank = get_tensor_model_parallel_rank()
48
+
49
+ # Depthwise conv weight; checkpoint stores (dim, 1, W).
50
+ self.weight = Parameter(torch.empty(dim, 1, kernel_size), requires_grad=False)
51
+ set_weight_attrs(self.weight, {"weight_loader": self.weight_loader})
52
+
53
+ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor) -> None:
54
+ if loaded_weight.shape[0] != param.shape[0]:
55
+ shard = param.shape[0]
56
+ loaded_weight = loaded_weight.narrow(0, self.tp_rank * shard, shard)
57
+ param.data.copy_(loaded_weight)
58
+
59
+ def forward(self, x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
60
+ # x: (num_tokens, dim); positions: (num_tokens,) absolute positions.
61
+ attn_metadata = get_forward_context().attn_metadata
62
+ if not isinstance(attn_metadata, dict):
63
+ # Memory-profiling / no metadata: identity (residual).
64
+ return x
65
+ m = attn_metadata.get(self.owner.prefix)
66
+ if m is None:
67
+ return x
68
+ assert isinstance(m, InklingSconvMetadata)
69
+ cache = self.owner.kv_cache
70
+ if cache.numel() == 0:
71
+ # Cache not yet bound (profiling before KV alloc): identity.
72
+ return x
73
+
74
+ off_s, ws = self.owner.stream_ranges[self.stream_idx]
75
+ # The hybrid KV-cache planner can enlarge the logical conv block so
76
+ # that its physical page size matches the attention caches (for
77
+ # example, W=4 becomes 32 when --block-size=128). Metadata slot
78
+ # mappings are built with that enlarged size, so index the bound cache
79
+ # with its runtime token dimension rather than the kernel window size.
80
+ block_size = self.owner.cache_block_size
81
+ x = x.contiguous()
82
+ weight = self.weight.squeeze(1) # (dim, W)
83
+
84
+ return fused_sconv(
85
+ x,
86
+ weight,
87
+ cache,
88
+ positions,
89
+ m.block_table,
90
+ m.seq_idx,
91
+ m.slot_mapping,
92
+ m.query_start,
93
+ off_s,
94
+ ws,
95
+ block_size,
96
+ activation=None,
97
+ use_residual=True,
98
+ )
vllm/triton_utils/__init__.py
+9
@@ -12,10 +12,16 @@ if TYPE_CHECKING or HAS_TRITON:
12
import triton
13
import triton.language as tl
14
import triton.language.extra.libdevice as tldevice
15
+ from triton.experimental import gluon
16
+ from triton.experimental.gluon import language as gl
17
+ from triton.language.core import _aggregate as aggregate # noqa: E501
18
else:
19
triton = TritonPlaceholder()
20
tl = TritonLanguagePlaceholder()
21
tldevice = TritonLanguagePlaceholder()
22
+ gluon = TritonLanguagePlaceholder()
23
+ gl = TritonLanguagePlaceholder()
24
+ aggregate = TritonLanguagePlaceholder()
25
26
from vllm.triton_utils.tensor_descriptor import use_tensor_descriptor
27
@@ -29,5 +35,8 @@ __all__ = [
35
"tldevice",
36
"LOG2E",
37
"LOGE2",
38
+ "gluon",
39
+ "gl",
40
+ "aggregate",
41
"use_tensor_descriptor",
42
]